Hugging Face Diffusers로 ControlNet 훈련하기

Hugging Face는 개발자들이 자체 ControlNet 모델을 훈련할 수 있도록 diffusers 라이브러리 내에 상세 가이드와 훈련 스크립트를 공개했습니다. ControlNet은 포즈 추정, 깊이 맵, 스케치와 같은 추가 조건을 생성 과정에 추가함으로써 확산 모델을 세밀하게 제어할 수 있는 신경망 구조입니다.

Stable Diffusion용 맞춤 ControlNet을 훈련하는 과정은 크게 세 단계로 이루어집니다: 조건 계획, 데이터셋 구축, 그리고 훈련 실행.

3단계 ControlNet 훈련 워크플로우

1. 조건 계획하기

첫 번째 단계는 작업에 필요한 구체적인 조건을 정의하는 것입니다. 여기에는 원하는 제어 메커니즘을 결정하고 기존 모델이 일반 이미지를 해당 조건으로 변환할 수 있는지 확인하는 과정이 포함됩니다. 예를 들어, Hugging Face 팀은 얼굴 랜드마크 모델을 만들어 Stable Diffusion이 특정 얼굴 표정이나 포즈를 따르도록 목표했습니다.

2. 데이터셋 구축하기

ControlNet 데이터셋은 세 가지 특정 열이 필요합니다:

  • Ground Truth Image: 대상 이미지 (예: 얼굴).
  • Conditioning Image: 조건을 나타내는 이미지 (예: 시각화된 얼굴 랜드마크 마스크).
  • Prompt: 이미지를 설명하는 텍스트 캡션.

"Uncanny Faces" 예시에서 팀은 100K 합성 얼굴을 포함한 Microsoft FaceSynthetics 데이터셋을 사용했습니다. 기존 모델 중 얼굴을 직접 해당 데이터셋의 랜드마크 형식으로 변환할 수 있는 것이 없었기 때문에, 팀은 최첨단 모델인 SPIGA를 사용해 iBUG 형식의 68개 얼굴 랜드마크를 추출하고, 이를 일러스트 마스크로 변환한 뒤, BLIP 캡셔닝을 이용해 각 이미지에 대한 설명을 생성했습니다.

3. 모델 훈련하기

diffusers 예제에 제공된 train_controlnet.py 스크립트를 사용해 훈련을 진행했습니다. 팀은 단일 A100 GPU를 사용해 배치 크기 4로 3 epoch 동안 훈련했습니다.

훈련 관찰 및 과적합

팀은 3 epoch 훈련이 과적합을 초래한다는 것을 발견했습니다. 모델이 스타일을 무시하고 실제 얼굴과 다른 개념(예: 프롬프트에 "고양이"나 "슈렉"을 생성하지 못함)을 잊어버렸습니다. 약 1 epoch(약 25K 스텝) 정도에서 수렴이 이루어졌으며, 이 시점에서 모델은 과적합 없이 포즈를 성공적으로 따랐습니다. FaceSynthetics 데이터셋이 합성 이미지로 구성되어 있었기 때문에, 최종 모델은 사진처럼 사실적인 얼굴이 아니라 "섬뜩한" 3D 느낌의 얼굴을 생성했습니다.

기술 구현 및 하드웨어 최적화

훈련 설정

train_controlnet.py 스크립트는 출력 제어를 위해 여러 핵심 매개변수를 활용합니다:

  • pretrained_model_name_or_path: 기본 Stable Diffusion 모델 (얼굴 렌더링을 개선하기 위해 v2-1-base 사용).
  • learning_rate: 예시에서는 1e-5로 설정했으며, 1e-4와 2e-6 사이의 값을 권장합니다.
  • resolution: 조건 이미지와 실제 이미지 모두 512x512로 설정합니다.
  • validation_steps: 모델이 검증 프롬프트와 이미지를 실행하여 진행 상황을 추적하는 빈도를 결정합니다.

저사양 GPU를 위한 VRAM 최적화

GPU VRAM 필요한 최적화 / 매개변수
16GB train_batch_size=1, gradient_accumulation_steps=4, gradient_checkpointing, 그리고 use_8bit_adam (bitsandbytes 사용).
12GB 16GB 최적화에 모두 추가로 set_grads_to_none.
8GB diffusers GitHub 훈련 가이드에 자세히 설명된 특정 구성.

배치 크기 1에 4개의 gradient accumulation step을 사용하면, 사용자는 A100 훈련에서 사용된 배치 크기 4와 동일한 효과를 시뮬레이션하면서 메모리 사용량을 크게 줄일 수 있습니다.

Sources