DreamBooth

DreamBooth는 대상이나 스타일에 대한 이미지 몇 장만으로 전체 diffusion 모델을 업데이트하는 학습 기법입니다. 프롬프트의 특수한 단어를 예시 이미지와 연관시킴으로써 작동합니다.

추론을 위해 학습된 체크포인트를 로드하려면 추론을 위한 DreamBooth 모델 로드를 참조하세요.

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

이 가이드에서는 train_dreambooth.py 스크립트를 살펴보고, 이를 더 잘 이해하며 사용자의 용도에 맞게 수정하는 방법을 안내합니다.

…

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

학습 스크립트가 포함된 예제 폴더로 이동하고, 사용 중인 스크립트에 필요한 의존성을 설치하세요:

cd examples/dreambooth
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()

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

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

스크립트 매개변수

DreamBooth is very sensitive to training hyperparameters, and it is easy to overfit. Read the Training Stable Diffusion with Dreambooth using 🧨 Diffusers blog post for recommended settings for different subjects to help you choose the appropriate hyperparameters.

학습 스크립트는 학습 실행을 사용자 정의할 수 있는 다양한 매개변수를 제공합니다. 모든 매개변수 및 설명은 parse_args() 함수에서 확인할 수 있습니다. 매개변수는 기본적으로 잘 작동하는 기본값으로 설정되어 있지만, 원하시는 경우 학습 명령에서 직접 값을 설정할 수도 있습니다.

예를 들어, bf16 형식으로 학습하려면:

accelerate launch train_dreambooth.py \
    --mixed_precision="bf16"

알고 지정해야 할 몇 가지 기본적이고 중요한 매개변수는 다음과 같습니다:

--pretrained_model_name_or_path: Hub의 모델 이름 또는 사전 학습된 모델의 로컬 경로
--instance_data_dir: 학습 데이터셋(예시 이미지)이 포함된 폴더의 경로
--instance_prompt: 예시 이미지에 포함된 특수 단어가 들어간 텍스트 프롬프트
--train_text_encoder: 텍스트 인코더도 함께 학습할지 여부
--output_dir: 학습된 모델을 저장할 위치
--push_to_hub: 학습된 모델을 Hub에 푸시할지 여부
--checkpointing_steps: 모델 학습 중 체크포인트를 저장하는 주기입니다. 학습이 중단된 경우, 학습 명령에 --resume_from_checkpoint를 추가하여 해당 체크포인트에서 학습을 계속할 수 있습니다.

Min-SNR 가중

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

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

accelerate launch train_dreambooth.py \
  --snr_gamma=5.0

사전 보존 손실

Prior preservation loss는 모델이 스스로 생성한 샘플을 활용하여 더 다양한 이미지를 생성하는 방법을 학습하도록 돕는 기법입니다. 이러한 생성된 샘플 이미지는 사용자가 제공한 이미지와 동일한 클래스에 속하므로, 모델이 해당 클래스에 대해 학습한 내용을 유지하고 이미 알고 있는 클래스 지식을 활용하여 새로운 구성을 만들어내는 데 도움이 됩니다.

--with_prior_preservation: prior preservation loss를 사용할지 여부
--prior_loss_weight: 모델에 대한 사전 보존 손실의 영향을 제어합니다
--class_data_dir: 생성된 클래스 샘플 이미지가 포함된 폴더의 경로
--class_prompt: 생성된 샘플 이미지의 클래스를 설명하는 텍스트 프롬프트

accelerate launch train_dreambooth.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="path/to/class/images" \
  --class_prompt="text prompt describing class"

텍스트 인코더 학습

생성된 출력의 품질을 개선하려면 UNet 외에도 텍스트 인코더를 학습할 수 있습니다. 이를 위해서는 추가 메모리가 필요하며, 최소 24GB의 vRAM을 갖춘 GPU가 필요합니다. 필요한 하드웨어를 갖추고 있다면, 텍스트 인코더를 학습하는 것이 더 나은 결과를 제공합니다. 특히 얼굴 이미지를 생성할 때 그렇습니다. 이 옵션을 활성화하려면:

accelerate launch train_dreambooth.py \
  --train_text_encoder

학습 스크립트

DreamBooth에는 자체 데이터셋 클래스가 포함되어 있습니다:

DreamBoothDataset: 학습을 위해 이미지와 클래스 이미지를 전처리하고 프롬프트를 토큰화합니다
PromptDataset: 클래스 이미지를 생성하기 위한 프롬프트 임베딩을 생성합니다

prior preservation loss를 활성화한 경우, 클래스 이미지는 여기에 생성됩니다:

sample_dataset = PromptDataset(args.class_prompt, num_new_images)
sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size)

sample_dataloader = accelerator.prepare(sample_dataloader)
pipeline.to(accelerator.device)

for example in tqdm(
    sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process
):
    images = pipeline(example["prompt"]).images

다음은 학습용 데이터셋 설정과 학습 루프 자체를 처리하는 main() 함수입니다. 스크립트는 tokenizer, 스케줄러 및 모델을 로드합니다:

# Load the tokenizer
if args.tokenizer_name:
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, revision=args.revision, use_fast=False)
elif args.pretrained_model_name_or_path:
    tokenizer = AutoTokenizer.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="tokenizer",
        revision=args.revision,
        use_fast=False,
    )

# Load scheduler and models
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
text_encoder = text_encoder_cls.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
)

if model_has_vae(args):
    vae = AutoencoderKL.from_pretrained(
        args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision
    )
else:
    vae = None

unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
)

이제 DreamBoothDataset에서 학습 데이터셋과 DataLoader를 생성할 차례입니다:

train_dataset = DreamBoothDataset(
    instance_data_root=args.instance_data_dir,
    instance_prompt=args.instance_prompt,
    class_data_root=args.class_data_dir if args.with_prior_preservation else None,
    class_prompt=args.class_prompt,
    class_num=args.num_class_images,
    tokenizer=tokenizer,
    size=args.resolution,
    center_crop=args.center_crop,
    encoder_hidden_states=pre_computed_encoder_hidden_states,
    class_prompt_encoder_hidden_states=pre_computed_class_prompt_encoder_hidden_states,
    tokenizer_max_length=args.tokenizer_max_length,
)

train_dataloader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=args.train_batch_size,
    shuffle=True,
    collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation),
    num_workers=args.dataloader_num_workers,
)

마지막으로, training loop은 이미지를 latent space로 변환하고, 입력에 노이즈를 추가하고, 노이즈 잔차를 예측하며, 손실을 계산하는 등 나머지 단계를 처리합니다.

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

스크립트 실행

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

이 가이드에서는 개 이미지의 일부를 다운로드하여 디렉터리에 저장합니다. 원하시는 경우, 직접 데이터셋을 생성하여 사용할 수 있습니다(학습용 데이터셋 만들기 가이드 참고).

from huggingface_hub import snapshot_download

local_dir = "./dog"
snapshot_download(
    "diffusers/dog-example",
    local_dir=local_dir,
    repo_type="dataset",
    ignore_patterns=".gitattributes",
)

환경 변수 MODEL_NAME을 Hub의 모델 ID 또는 로컬 모델 경로로, INSTANCE_DIR을 방금 다운로드한 개 이미지 경로를, OUTPUT_DIR을 모델을 저장할 위치로 설정하세요. 학습을 연결할 특수한 단어로 sks를 사용합니다.

학습 과정을 따라가 보고 싶다면, 학습이 진행되는 동안 생성된 이미지를 주기적으로 저장할 수 있습니다. 학습 명령에 다음 매개변수를 추가하세요:

--validation_prompt="a photo of a sks dog"
--num_validation_images=4
--validation_steps=100

스크립트를 실행하기 전에 한 가지 더! 사용 중인 GPU에 따라 DreamBooth 학습을 위해 특정 최적화를 활성화해야 할 수 있습니다.

16GB

12GB

8GB

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export INSTANCE_DIR="./dog"
export OUTPUT_DIR="path_to_saved_model"

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=400 \
  --push_to_hub

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

학습이 완료되기 전에 추론으로 모델을 사용해 보고 싶으신가요? 🤭 최신 버전의 🤗 Accelerate가 설치되어 있는지 확인하세요.

from diffusers import DiffusionPipeline, UNet2DConditionModel
from transformers import CLIPTextModel
import torch

unet = UNet2DConditionModel.from_pretrained("path/to/model/checkpoint-100/unet")

# if you have trained with `--args.train_text_encoder` make sure to also load the text encoder
text_encoder = CLIPTextModel.from_pretrained("path/to/model/checkpoint-100/checkpoint-100/text_encoder")

pipeline = DiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, text_encoder=text_encoder, dtype=torch.float16,
).to("cuda")

image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("dog-bucket.png")

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("path_to_saved_model", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"
image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("dog-bucket.png")

LoRA

LoRA는 학습 가능한 매개변수의 수를 크게 줄이는 학습 기법입니다. 그 결과, 학습 속도가 빨라지고, 생성된 가중치가 훨씬 작아(~100MB) 저장하기가 쉬워집니다. LoRA로 학습하려면 train_dreambooth_lora.py 스크립트를 사용하세요.

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

Stable Diffusion XL

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

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

DeepFloyd IF

DeepFloyd IF는 세 단계로 구성된 캐스케이딩 픽셀 확산 모델입니다. 첫 번째 단계에서 기본 이미지를 생성하고, 두 번째와 세 번째 단계에서 기본 이미지를 점진적으로 업스케일링하여 고해상도 1024x1024 이미지를 생성합니다. train_dreambooth_lora.py 또는 train_dreambooth.py 스크립트를 사용하여 LoRA 또는 전체 모델로 DeepFloyd IF 모델을 학습할 수 있습니다.

DeepFloyd IF는 예측된 분산을 사용하지만, Diffusers 학습 스크립트는 예측된 오류를 사용하므로 학습된 DeepFloyd IF 모델은 고정된 분산 스케줄로 전환됩니다. 학습 스크립트는 완전히 학습된 모델의 스케줄러 구성을 자동으로 업데이트해 줍니다. 그러나 저장된 LoRA 가중치를 로드할 때는 파이프라인의 스케줄러 구성도 직접 업데이트해야 합니다.

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", use_safetensors=True)

pipe.load_lora_weights("<lora weights path>")

# Update scheduler config to fixed variance schedule
pipe.scheduler = pipe.scheduler.__class__.from_config(pipe.scheduler.config, variance_type="fixed_small")

stage 2 모델은 업스케일링을 위해 추가 검증 이미지가 필요합니다. 이를 위해 학습 이미지의 축소된 버전을 다운로드하여 사용할 수 있습니다.

from huggingface_hub import snapshot_download

local_dir = "./dog_downsized"
snapshot_download(
    "diffusers/dog-example-downsized",
    local_dir=local_dir,
    repo_type="dataset",
    ignore_patterns=".gitattributes",
)

아래 코드 샘플은 DreamBooth와 LoRA의 조합으로 DeepFloyd IF 모델을 학습하는 방법에 대한 간단한 개요를 제공합니다. 주목해야 할 몇 가지 중요한 매개변수는 다음과 같습니다:

--resolution=64, DeepFloyd IF는 픽셀 확산 모델이므로 압축되지 않은 픽셀을 다루기 위해 입력 이미지가 더 작아야 하기 때문에 훨씬 작은 해상도가 필요합니다
--pre_compute_text_embeddings, T5Model이 많은 메모리를 차지할 수 있으므로 메모리를 절약하기 위해 텍스트 임베딩을 미리 계산합니다
--tokenizer_max_length=77를 사용하면 T5를 텍스트 인코더로 사용할 때 더 긴 기본 텍스트 길이를 사용할 수 있지만, 기본 모델 인코딩 절차는 더 짧은 텍스트 길이를 사용합니다.
--text_encoder_use_attention_mask를 사용하여 attention mask를 text encoder에 전달합니다

Stage 1 LoRA DreamBooth

Stage 2 LoRA DreamBooth

Stage 1 DreamBooth

Stage 2 DreamBooth

LoRA와 DreamBooth를 사용하여 DeepFloyd IF의 1단계 학습을 수행하려면 약 28GB의 메모리가 필요합니다.

export MODEL_NAME="DeepFloyd/IF-I-XL-v1.0"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="dreambooth_dog_lora"

accelerate launch train_dreambooth_lora.py \
  --report_to wandb \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a sks dog" \
  --resolution=64 \
  --train_batch_size=4 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --scale_lr \
  --max_train_steps=1200 \
  --validation_prompt="a sks dog" \
  --validation_epochs=25 \
  --checkpointing_steps=100 \
  --pre_compute_text_embeddings \
  --tokenizer_max_length=77 \
  --text_encoder_use_attention_mask

학습 팁

DeepFloyd IF 모델의 학습은 어려울 수 있지만, 유용하다고 판단한 몇 가지 팁을 공유합니다:

모델의 낮은 해상도로 인해 더 세밀한 세부 사항을 표현하는 것이 본질적으로 어렵기 때문에, 1단계 모델 학습에는 LoRA만으로도 충분합니다.
일반적이거나 단순한 객체의 경우, 업스케일러를 파인튜닝할 필요가 없을 수 있습니다. 업스케일러에 전달되는 프롬프트가 인스턴스 프롬프트에서 새로운 토큰을 제거하도록 조정되었는지 확인하세요. 예를 들어, 1단계 프롬프트가 “a sks dog”라면 2단계 프롬프트는 “a dog”여야 합니다.
얼굴과 같은 더 세밀한 디테일의 경우, stage 2 업스케일러를 완전히 학습하는 것이 LoRA로 stage 2 모델을 학습하는 것보다 더 좋습니다. 또한 더 큰 배치 크기를 사용하면서 더 낮은 학습률을 사용하는 것도 도움이 됩니다.
2단계 모델을 학습할 때는 더 낮은 학습률을 사용해야 합니다.
DDPMScheduler는 학습 스크립트에서 사용된 DPMSolver보다 더 잘 작동합니다.

다음 단계

DreamBooth 모델 학습을 완료하신 것을 축하드립니다! 새로 학습한 모델을 사용하는 방법에 대해 자세히 알아보려면 다음 가이드가 도움이 될 수 있습니다:

추론을 위해 DreamBooth 체크포인트를 로드하는 방법을 알아봅니다.
LoRA로 학습했다면 추론을 위한 LoRA를 참조하세요.

Update on GitHub