ControlNet

ControlNet 모델은 다른 사전 학습된 모델 위에 학습된 어댑터입니다. 추가 입력 이미지로 모델을 조건화하여 이미지 생성에 대한 제어도를 높일 수 있습니다. 입력 이미지는 canny edge, depth map, human pose 등 다양한 형태일 수 있습니다.

vRAM이 제한된 GPU에서 학습하는 경우, 학습 명령에서 gradient_checkpointing, gradient_accumulation_steps, mixed_precision 매개변수를 활성화하는 것을 시도해 보세요. 또한 xFormers를 사용하여 메모리 효율적인 어텐션을 적용하면 메모리 사용량을 줄일 수 있습니다.

이 가이드에서는 train_controlnet.py 학습 스크립트를 살펴보고, 이를 이해한 후 자신의 사용 사례에 맞게 수정하는 방법을 안내합니다.

…

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

이후 학습 스크립트가 포함된 예제 폴더로 이동하여 사용하는 스크립트에 필요한 의존성을 설치합니다:

cd examples/controlnet
pip install -r requirements.txt

🤗 Accelerate는 여러 GPU/TPU 또는 혼합 정밀도로 학습하는 것을 돕는 라이브러리입니다. 하드웨어와 환경에 따라 학습 설정을 자동으로 구성합니다. 자세한 내용은 🤗 Accelerate Quick tour를 참고하세요.

🤗 Accelerate 환경을 초기화합니다:

accelerate config

어떤 구성도 선택하지 않고 기본 🤗 Accelerate 환경을 설정하려면:

accelerate config default

또는 노트북과 같이 인터랙티브 셸을 지원하지 않는 환경에서는 다음을 사용할 수 있습니다:

from accelerate.utils import write_basic_config

write_basic_config()

마지막으로, 자체 데이터셋으로 모델을 학습하려면 학습용 데이터셋 만들기 가이드를 참고하여 학습 스크립트와 호환되는 데이터셋을 만드는 방법을 배워보세요.

다음 섹션에서는 스크립트를 수정하는 방법을 이해하는 데 중요한 학습 스크립트의 일부 부분을 강조합니다. 그러나 스크립트의 모든 측면을 자세히 다루지는 않습니다. 더 자세히 알고 싶으시다면 스크립트를 읽어보시고, 질문이나 우려 사항이 있으시면 알려주시기 바랍니다.

스크립트 매개변수

학습 스크립트는 학습 실행을 사용자 정의하는 데 도움이 되는 많은 매개변수를 제공합니다. 모든 매개변수와 그 설명은 parse_args() 함수에서 찾을 수 있습니다. 이 함수는 학습 배치 크기나 학습률과 같이 각 매개변수에 대한 기본 값을 제공하지만, 원하시는 경우 학습 명령에서 직접 값을 설정할 수도 있습니다.

예를 들어, fp16 형식을 사용하여 혼합 정밀도로 학습 속도를 높이려면 학습 명령에 --mixed_precision 매개변수를 추가합니다:

accelerate launch train_controlnet.py \
  --mixed_precision="fp16"

기본적이고 중요한 대부분의 매개변수는 텍스트-이미지 학습 가이드에 설명되어 있으므로, 이 가이드에서는 ControlNet에 관련되는 매개변수에만 초점을 맞춥니다:

--max_train_samples: 학습 샘플 수입니다. 이 값을 낮추면 학습 속도를 높일 수 있지만, 매우 큰 데이터셋을 스트리밍하려면 학습 명령에 이 매개변수와 --streaming 매개변수를 함께 포함해야 합니다.
--gradient_accumulation_steps: backward pass를 실행하기 전에 누적할 업데이트 단계 수. 이를 통해 GPU 메모리가 일반적으로 처리할 수 있는 것보다 더 큰 배치 크기로 학습할 수 있습니다

Min-SNR 가중

Min-SNR 가중 전략은 손실을 재균형화하여 수렴 속도를 높이는 데 도움이 됩니다. 학습 스크립트는 epsilon(노이즈) 또는 v_prediction 예측을 지원하지만, Min-SNR은 두 가지 예측 유형 모두와 호환됩니다. 이 가중 전략은 PyTorch에서만 지원됩니다.

--snr_gamma 매개변수를 추가하고 권장 값인 5.0으로 설정합니다:

accelerate launch train_controlnet.py \
  --snr_gamma=5.0

학습 스크립트

스크립트 매개변수와 마찬가지로, Text-to-image 학습 가이드에서 학습 스크립트의 일반적인 개요를 제공합니다. 대신 이 가이드에서는 ControlNet 스크립트의 관련 부분을 살펴봅니다.

학습 스크립트에는 이미지 변환 및 캡션 토큰화를 통해 데이터셋을 전처리하는 make_train_dataset 함수가 있습니다. 일반적인 캡션 토큰화와 이미지 변환 외에도, 스크립트에는 조건부 이미지에 대한 변환도 포함되어 있음을 확인할 수 있습니다.

TPU에서 데이터셋을 스트리밍하는 경우, 이미지 최적화가 되지 않은 🤗 Datasets 라이브러리로 인해 성능이 병목 현상을 겪을 수 있습니다. 최대 처리량을 보장하기 위해 WebDataset, TorchData, TensorFlow Datasets와 같은 다른 데이터셋 형식을 탐색해 보시기 바랍니다.

conditioning_image_transforms = transforms.Compose(
    [
        transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
        transforms.CenterCrop(args.resolution),
        transforms.ToTensor(),
    ]
)

main() 함수 안에는 tokenizer, text encoder, scheduler 및 모델을 로드하는 코드가 있습니다. 또한 ControlNet 모델을 기존 가중치에서 로드하거나 UNet에서 무작위로 초기화하는 로직도 여기에 있습니다:

if args.controlnet_model_name_or_path:
    logger.info("Loading existing controlnet weights")
    controlnet = ControlNetModel.from_pretrained(args.controlnet_model_name_or_path)
else:
    logger.info("Initializing controlnet weights from unet")
    controlnet = ControlNetModel.from_unet(unet)

optimizer는 ControlNet 매개변수를 업데이트하도록 설정됩니다:

params_to_optimize = controlnet.parameters()
optimizer = optimizer_class(
    params_to_optimize,
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

마지막으로, training loop에서 conditioning text embeddings와 이미지가 ControlNet 모델의 down 및 mid-blocks에 전달됩니다:

encoder_hidden_states = text_encoder(batch["input_ids"])[0]
controlnet_image = batch["conditioning_pixel_values"].to(dtype=weight_dtype)

down_block_res_samples, mid_block_res_sample = controlnet(
    noisy_latents,
    timesteps,
    encoder_hidden_states=encoder_hidden_states,
    controlnet_cond=controlnet_image,
    return_dict=False,
)

학습 루프가 어떻게 작동하는지 더 자세히 알고 싶다면, 디노이징 과정의 기본 패턴을 단계별로 설명하는 Understanding pipelines, models and schedulers 튜토리얼을 확인해 보세요.

스크립트 실행

이제 학습 스크립트를 실행할 준비가 되었습니다! 🚀

이 가이드는 fusing/fill50k 데이터셋을 사용하지만, 원하시는 경우 직접 데이터셋을 생성하여 사용할 수 있습니다(학습용 데이터셋 만들기 가이드 참고).

환경 변수 MODEL_NAME을 Hub의 모델 ID 또는 로컬 모델 경로로, OUTPUT_DIR을 모델을 저장할 위치로 설정하세요.

다음 이미지들을 다운로드하여 학습을 조건화하세요:

wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_1.png
wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_2.png

스크립트를 실행하기 전에 한 가지 더! 사용 중인 GPU에 따라 ControlNet을 학습하기 위해 특정 최적화를 활성화해야 할 수 있습니다. 이 스크립트의 기본 구성은 약 38GB의 vRAM이 필요합니다. 여러 GPU에서 학습하는 경우 accelerate launch 명령에 --multi_gpu 매개변수를 추가하세요.

16GB

12GB

8GB

export MODEL_DIR="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="path/to/save/model"

accelerate launch train_controlnet.py \
 --pretrained_model_name_or_path=$MODEL_DIR \
 --output_dir=$OUTPUT_DIR \
 --dataset_name=fusing/fill50k \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./conditioning_image_1.png" "./conditioning_image_2.png" \
 --validation_prompt "red circle with blue background" "cyan circle with brown floral background" \
 --train_batch_size=1 \
 --gradient_accumulation_steps=4 \
 --push_to_hub

학습이 완료되면, 새로 학습한 모델을 추론에 사용할 수 있습니다!

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import torch

controlnet = ControlNetModel.from_pretrained("path/to/controlnet", dtype=torch.float16)
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
    "path/to/base/model", controlnet=controlnet, dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

control_image = load_image("./conditioning_image_1.png")
prompt = "pale golden rod circle with old lace background"

generator = torch.manual_seed(0)
image = pipeline(prompt, num_inference_steps=20, generator=generator, image=control_image).images[0]
image.save("./output.png")

Stable Diffusion XL

Stable Diffusion XL(SDXL)은 고해상도 이미지를 생성하는 강력한 텍스트-투-이미지 모델이며, 아키텍처에 두 번째 텍스트 인코더를 추가했습니다. train_controlnet_sdxl.py 스크립트를 사용하여 SDXL 모델의 ControlNet 어댑터를 학습하세요.

SDXL 학습 스크립트는 SDXL 학습 가이드에서 더 자세히 설명되어 있습니다.

다음 단계

직접 ControlNet을 학습하신 것을 축하드립니다! 새로 학습한 모델을 사용하는 방법에 대해 더 알고 싶으시다면, 다음 가이드가 도움이 될 수 있습니다:

다양한 작업에서 추론을 위해 ControlNet을 사용하는 방법을 알아보세요.

Update on GitHub

Diffusers

ControlNet

스크립트 매개변수

Min-SNR 가중

학습 스크립트

스크립트 실행

Stable Diffusion XL

다음 단계