텍스트-이미지

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

Stable Diffusion과 같은 텍스트-이미지 모델은 텍스트 프롬프트를 조건으로 받아 이미지를 생성하도록 학습됩니다.

모델을 학습하는 것은 하드웨어에 상당한 부하를 줄 수 있지만, gradient_checkpointing와 mixed_precision를 활성화하면 단일 24GB GPU에서도 모델을 학습할 수 있습니다. 더 큰 배치 크기로 학습하거나 더 빠르게 학습하려면 30GB 이상의 메모리를 가진 GPU를 사용하는 것이 좋습니다. xFormers를 사용하여 메모리 효율적인 어텐션을 활성화하면 메모리 사용량을 줄일 수 있습니다.

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

…

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

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

cd examples/text_to_image
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_text_to_image.py \
  --mixed_precision="fp16"

일부 기본적이고 중요한 매개변수는 다음과 같습니다:

--pretrained_model_name_or_path: Hub의 모델 이름 또는 사전 학습된 모델의 로컬 경로
--dataset_name: Hub에 있는 데이터셋의 이름 또는 학습에 사용할 데이터셋의 로컬 경로
--image_column: 학습에 사용할 데이터셋의 이미지 열 이름
--caption_column: 학습에 사용할 데이터셋의 텍스트 열 이름
--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_text_to_image.py \
  --snr_gamma=5.0

이 Weights and Biases 리포트에서 서로 다른 snr_gamma 값에 대한 손실 표면을 비교할 수 있습니다. 작은 데이터셋의 경우, 큰 데이터셋과 비교했을 때 Min-SNR의 효과가 덜 뚜렷할 수 있습니다.

학습 스크립트

데이터셋 전처리 코드와 학습 루프는 main() 함수에 있습니다. 학습 스크립트를 수정해야 한다면 이 부분을 변경해야 합니다.

train_text_to_image 스크립트는 스케줄러와 토크나이저를 로드하는 것으로 시작합니다. 원하시는 경우 여기서 다른 스케줄러를 사용할 수 있습니다:

noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
tokenizer = CLIPTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
)

그 후 스크립트는 UNet 모델을 로드합니다:

load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
model.register_to_config(**load_model.config)

model.load_state_dict(load_model.state_dict())

다음으로, 데이터셋의 텍스트 및 이미지 컬럼을 전처리해야 합니다. tokenize_captions 함수는 입력을 토큰화하는 역할을 하고, train_transforms 함수는 이미지에 적용할 변환의 유형을 지정합니다. 이 두 함수는 모두 preprocess_train에 묶여 있습니다:

def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[image_column]]
    examples["pixel_values"] = [train_transforms(image) for image in images]
    examples["input_ids"] = tokenize_captions(examples)
    return examples

마지막으로, training loop은 나머지 모든 작업을 처리합니다. 이미지를 latent space로 인코딩하고, latents에 노이즈를 추가하며, 조건으로 사용할 텍스트 임베딩을 계산하고, 모델 파라미터를 업데이트한 후, 모델을 저장하고 Hub에 push합니다. training loop의 작동 방식에 대해 더 자세히 알고 싶다면, denoising 과정의 기본 패턴을 분석하는 Understanding pipelines, models and schedulers 튜토리얼을 확인해 보세요.

스크립트 실행

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

Naruto BLIP captions 데이터셋을 학습하여 나만의 Naruto 캐릭터를 생성해 보겠습니다. 환경 변수 MODEL_NAME와 dataset_name를 각각 모델과 데이터셋(Hub 또는 로컬 경로)으로 설정하세요. GPU가 하나 이상인 경우 accelerate launch 명령에 --multi_gpu 매개변수를 추가하세요.

로컬 데이터셋으로 학습하려면 TRAIN_DIR 및 OUTPUT_DIR 환경 변수를 데이터셋 경로와 모델을 저장할 위치로 설정하세요.

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export dataset_name="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$dataset_name \
  --use_ema \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --enable_xformers_memory_efficient_attention \
  --lr_scheduler="constant" --lr_warmup_steps=0 \
  --output_dir="sd-naruto-model" \
  --push_to_hub

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

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("path/to/saved_model", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"

image = pipeline(prompt="yoda").images[0]
image.save("yoda-naruto.png")

다음 단계

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

LoRA로 모델을 학습했다면 추론을 위해 LoRA 가중치를 로드하는 방법을 알아보세요.
guidance scale와 같은 특정 파라미터나 prompt weighting과 같은 기법이 추론을 제어하는 데 어떻게 도움이 되는지 자세히 알아보려면 Text-to-image 작업 가이드를 참고하세요.

Update on GitHub

Diffusers

텍스트-이미지

스크립트 매개변수

Min-SNR 가중

학습 스크립트

스크립트 실행

다음 단계