Stable Diffusion XL

이 스크립트는 실험적이며 과적합이 쉽게 발생하고 catastrophic forgetting과 같은 문제에 직면할 수 있습니다. 데이터셋에서 최상의 결과를 얻기 위해 다양한 하이퍼파라미터를 탐색해 보세요.

Stable Diffusion XL (SDXL)은 Stable Diffusion 모델의 더 크고 강력한 버전으로, 더 높은 해상도의 이미지를 생성할 수 있습니다.

SDXL의 UNet은 3배 더 크고 모델 아키텍처에 두 번째 텍스트 인코더가 추가됩니다. 사용 가능한 하드웨어에 따라 이는 매우 높은 계산 자원을 요구할 수 있으며, Tesla T4와 같은 소비자용 GPU에서는 실행되지 않을 수 있습니다. 이 더 큰 모델을 메모리에 적합하게 하고 학습 속도를 높이기 위해 gradient_checkpointing, mixed_precision, gradient_accumulation_steps를 활성화하는 것을 시도해 보세요. xFormers를 사용하여 메모리 효율적인 어텐션을 활성화하고 bitsandbytes’ 8비트 옵티마이저를 사용하면 메모리 사용량을 더욱 줄일 수 있습니다.

This guide will explore the train_text_to_image_sdxl.py training script to help you become more familiar with it, and how you can adapt it for your own use-case.

…

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

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

cd examples/text_to_image
pip install -r requirements_sdxl.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() 함수에서 찾을 수 있습니다. 이 함수는 학습 배치 크기나 학습률과 같이 각 매개변수에 대한 기본 값을 제공하지만, 원하시는 경우 학습 명령에서 직접 값을 설정할 수도 있습니다.

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

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

대부분의 파라미터는 Text-to-image 학습 가이드의 파라미터와 동일하므로, 이 가이드에서는 SDXL 학습에 관련되는 파라미터에 집중합니다.

• --pretrained_vae_model_name_or_path: 사전 학습된 VAE의 경로; SDXL VAE는 수치적 불안정성이 있는 것으로 알려져 있으므로, 이 매개변수를 통해 더 나은 VAE를 지정할 수 있습니다
• --proportion_empty_prompts: 빈 문자열로 대체할 이미지 프롬프트의 비율
• --timestep_bias_strategy: timestep의 어느 구간(초반 또는 후반)에 편향을 적용할지를 지정하며, 이를 통해 모델이 저주파 또는 고주파 세부 사항을 학습하도록 유도할 수 있습니다
• --timestep_bias_multiplier: timestep에 적용할 편향의 가중치
• --timestep_bias_begin: 편향을 적용하기 시작하는 타임스텝
• --timestep_bias_end: bias를 적용할 마지막 timestep
• --timestep_bias_portion: 편향을 적용할 타임스텝의 비율

Min-SNR 가중

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

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

accelerate launch train_text_to_image_sdxl.py \
  --snr_gamma=5.0

학습 스크립트

학습 스크립트는 Text-to-image 학습 가이드와 유사하지만, SDXL 학습을 지원하도록 수정되었습니다. 이 가이드에서는 SDXL 학습 스크립트에 고유한 코드에 초점을 맞춥니다.

프롬프트 임베딩을 계산하기 위해 프롬프트를 토큰화하는 함수와, VAE로 이미지 임베딩을 계산하는 함수를 만드는 것으로 시작합니다. 다음으로, 적용할 타임스텝 수와 타임스텝 편향 전략에 따라 타임스텝 가중치를 생성하는 함수를 만듭니다.

main() 함수 내에서 토크나이저를 로드하는 것 외에도, SDXL 아키텍처가 각각 두 개를 사용하므로 스크립트는 두 번째 토크나이저와 텍스트 인코더를 로드합니다:

tokenizer_one = AutoTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision, use_fast=False
)
tokenizer_two = AutoTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer_2", revision=args.revision, use_fast=False
)

text_encoder_cls_one = import_model_class_from_model_name_or_path(
    args.pretrained_model_name_or_path, args.revision
)
text_encoder_cls_two = import_model_class_from_model_name_or_path(
    args.pretrained_model_name_or_path, args.revision, subfolder="text_encoder_2"
)

프롬프트 및 이미지 임베딩은 먼저 계산되어 메모리에 유지되며, 이는 작은 데이터셋의 경우 일반적으로 문제가 되지 않지만, 더 큰 데이터셋의 경우 메모리 문제를 일으킬 수 있습니다. 이 경우, 사전 계산된 임베딩을 디스크에 별도로 저장한 후 학습 과정에서 메모리에 로드해야 합니다(이 주제에 대한 추가 논의는 이 PR을 참조하세요).

text_encoders = [text_encoder_one, text_encoder_two]
tokenizers = [tokenizer_one, tokenizer_two]
compute_embeddings_fn = functools.partial(
    encode_prompt,
    text_encoders=text_encoders,
    tokenizers=tokenizers,
    proportion_empty_prompts=args.proportion_empty_prompts,
    caption_column=args.caption_column,
)

train_dataset = train_dataset.map(compute_embeddings_fn, batched=True, new_fingerprint=new_fingerprint)
train_dataset = train_dataset.map(
    compute_vae_encodings_fn,
    batched=True,
    batch_size=args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps,
    new_fingerprint=new_fingerprint_for_vae,
)

임베딩을 계산한 후, 텍스트 인코더, VAE, 토크나이저를 삭제하여 메모리를 일부 확보합니다:

del text_encoders, tokenizers, vae
gc.collect()
torch.cuda.empty_cache()

마지막으로, training loop이 나머지 작업을 처리합니다. timestep bias 전략을 적용하기로 선택했다면, timestep 가중치가 계산되어 노이즈로 추가되는 것을 확인할 수 있습니다:

weights = generate_timestep_weights(args, noise_scheduler.config.num_train_timesteps).to(
        model_input.device
    )
    timesteps = torch.multinomial(weights, bsz, replacement=True).long()

noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps)

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

스크립트 실행

모든 변경 사항을 완료했거나 기본 구성으로 진행하는 데 문제가 없다면, 학습 스크립트를 실행할 준비가 된 것입니다! 🚀

Naruto BLIP captions 데이터셋을 학습하여 나루토 캐릭터를 생성해 보겠습니다. 모델과 데이터셋(Hub 또는 로컬 경로)을 MODEL_NAME 및 DATASET_NAME 환경 변수로 설정하세요. 또한 수치 불안정성을 피하기 위해 SDXL VAE가 아닌 다른 VAE(Hub 또는 로컬 경로)를 VAE_NAME으로 지정해야 합니다.

Weights & Biases로 학습 진행 상황을 모니터링하려면 학습 명령에 --report_to=wandb 매개변수를 추가하세요. 결과를 추적하려면 학습 명령에 --validation_prompt와 --validation_epochs도 추가해야 합니다. 이는 모델 디버깅 및 중간 결과 확인에 매우 유용합니다.

export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0"
export VAE_NAME="madebyollin/sdxl-vae-fp16-fix"
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch train_text_to_image_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --pretrained_vae_model_name_or_path=$VAE_NAME \
  --dataset_name=$DATASET_NAME \
  --enable_xformers_memory_efficient_attention \
  --resolution=512 \
  --center_crop \
  --random_flip \
  --proportion_empty_prompts=0.2 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=10000 \
  --use_8bit_adam \
  --learning_rate=1e-06 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --mixed_precision="fp16" \
  --report_to="wandb" \
  --validation_prompt="a cute Sundar Pichai creature" \
  --validation_epochs 5 \
  --checkpointing_steps=5000 \
  --output_dir="sdxl-naruto-model" \
  --push_to_hub

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

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("path/to/your/model", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "A naruto with green eyes and red legs."
image = pipeline(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
image.save("naruto.png")

from diffusers import DiffusionPipeline
import torch
import torch_xla.core.xla_model as xm

device = xm.xla_device()
pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0").to(device)

prompt = "A naruto with green eyes and red legs."
start = time()
image = pipeline(prompt, num_inference_steps=inference_steps).images[0]
print(f'Compilation time is {time()-start} sec')
image.save("naruto.png")

start = time()
image = pipeline(prompt, num_inference_steps=inference_steps).images[0]
print(f'Inference time is {time()-start} sec after compilation')

다음 단계

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

• Stable Diffusion XL 가이드를 읽어 다양한 작업(텍스트-이미지 생성, 이미지-이미지 변환, 인페인팅)에 사용하는 방법, 리파이너 모델 사용법, 그리고 다양한 유형의 마이크로 컨디셔닝에 대해 배워보세요.
• DreamBooth 및 LoRA 학습 가이드를 참고하여 몇 장의 예시 이미지만으로 개인화된 SDXL 모델을 학습하는 방법을 배울 수 있습니다. 이 두 가지 학습 기법은 함께 사용할 수도 있습니다!