ControlNet 학습하기, huggingface 코드

ControlNet이란?

ControlNet은 Stable Diffusion의 text prompt 이외 추가적인 condition을 통해 내가 원하는 이미지를 생성할 수 있게 해주는 network입니다.

 

자세한 내용은 다음의 링크를 참고해 주세요!

논문, 짧은 정리, 자세한 정리

 

---

ControlNet 학습하기

ControlNet을 학습하기 위해 제가 사용해 볼 코드는 다음과 같습니다.

https://github.com/huggingface/diffusers/blob/main/examples/controlnet/train_controlnet.py

diffusers/examples/controlnet/train_controlnet.py at main · huggingface/diffusers

 

유명한 huggingface의 diffusers library에서 제공하는 train_controlnet.py라는 코드를 바탕으로 huggingface에 올라와 있는 데이터에 학습하는 방법을 알아볼게요.

 

우선 위의 링크를 따라 train_controlnet.py를 복사해 주세요!

위의 버튼을 클릭하면 복사가 됩니다. 그리고 똑같이 train_conrolnet.py라는 파일에 코드를 붙여주세요.

 

만약 huggingface에 올라와있는 데이터에 대해 학습을 하고자 하면 아래의 bash 파일을 실행해주기만 해도 됩니다.

!accelerate launch train_controlnet.py \
 --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-1-base" \
 --output_dir="model_out" \
 --dataset_name=multimodalart/facesyntheticsspigacaptioned \
 --conditioning_image_column=spiga_seg \
 --image_column=image \
 --caption_column=image_caption \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./face_landmarks1.jpeg" "./face_landmarks2.jpeg" "./face_landmarks3.jpeg" \
 --validation_prompt "High-quality close-up dslr photo of man wearing a hat with trees in the background" "Girl smiling, professional dslr photograph, dark background, studio lights, high quality" "Portrait of a clown face, oil on canvas, bittersweet expression" \
 --train_batch_size=4 \
 --num_train_epochs=3 \
 --tracker_project_name="controlnet" \
 --enable_xformers_memory_efficient_attention \
 --checkpointing_steps=5000 \
 --validation_steps=5000 \
 --report_to wandb \
 --push_to_hub

 

간단한 설명

• pretrained_model_name_or_path="stabilityai/stable-diffusion-2-1-base" : Stable diffusion을 stabilityai/stable-diffusion-2-1-base의 checkpoitns를 이용하겠다.
•  --output_dir="model_out" : 학습한 controlnet의 저장 폴더 설정
•  --dataset_name=multimodalart/facesyntheticsspigacaptioned : huggingface에 등록된 얼굴 합성 데이터를 이용하겠다.

facesyntheticsspigacaptioned에

multimodalart/facesyntheticsspigacaptioned · Datasets at Hugging Face

학습 데이터가 100,000개가 있으며 image, image_seg, landmarks, spiga, spiga_seg, image_caption으로 구성되어 있습니다.

• --conditioning_image_column=spiga_seg : huggingface에 등록된 얼굴 합성 데이터 중 spiga_seg라는 column을 conditioning 이미지로 활용하겠다.
• --image_column=image : GT image를 image라는 column으로 설정.
• --caption_column=image_caption : text prompt를 image_caption이라는 column으로 설정.
• --resolution=512 : 이미지 크기를 512로 설정.
• --learning_rate=1e-5 : 학습률을 1e-5로 설정.
• --validation_image "./face_landmarks1.jpeg" "./face_landmarks2.jpeg" "./face_landmarks3.jpeg" : validation을 위한 controlnet의 입력인 condition 이미지 설정.
• --validation_prompt "High-quality close-up dslr photo of man wearing a hat with trees in the background" "Girl smiling, professional dslr photograph, dark background, studio lights, high quality" "Portrait of a clown face, oil on canvas, bittersweet expression" : validation을 위한 text prompt들. --validation_image와 --validation_prompt의 개수가 같아야 합니다.
• --train_batch_size=4 : 학습 batch_size 4. GPU VRAM 대략 16GB~17GB 정도.
• --num_train_epochs=3 : 학습 3 epoch 돌리겠다.
• --tracker_project_name="controlnet" : wandb에서 controlnet이라는 프로젝트 명으로 학습 과정이 기록됨.
• --enable_xformers_memory_efficient_attention : 메모리와 연산 효율을 위한.
• --checkpointing_steps=5000 : 5000 step 마다 checkpoint 저장.
• --validation_steps=5000 : 5000 step 마다 validation 실행. Wandb에 결과 저장.
• --report_to wandb : wandb에 기록.
• --push_to_hub : 학습된 모델을 huggingface 허브에 올림

---

학습 결과는 다음 링크에서 확인해 볼 수 있어요.

학습결과

Uncanny Faces Training

 

결과 예시

step 5000

step 75000

 

학습이 진행될수록 conditioning 이미지에 맞게 얼굴이 생성되는 것을 확인할 수 있어요.

---

다음에는 custom data에 controlnet을 학습하는 방법에 대해 알아볼게요.