Kandinsky 2.2

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

Kandinsky 2.2는 더 사실적인 이미지를 생성할 수 있는 다국어 텍스트-이미지 모델입니다. 이 모델은 텍스트 프롬프트에서 이미지 임베딩을 생성하기 위한 이미지 사전 모델과, 사전 모델의 임베딩을 기반으로 이미지를 생성하는 디코더 모델을 포함합니다. 따라서 Diffusers에서 Kandinsky 2.2를 위해 두 개의 별도 스크립트를 찾을 수 있으며, 하나는 사전 모델 학습용이고 다른 하나는 디코더 모델 학습용입니다. 두 모델을 각각 독립적으로 학습할 수 있지만, 최적의 결과를 얻으려면 사전 모델과 디코더 모델 모두를 학습해야 합니다.

GPU에 따라 모델이 메모리에 맞게 훈련 속도를 높이기 위해 gradient_checkpointing(⚠️ 이전 모델에서는 지원되지 않음!), mixed_precision, gradient_accumulation_steps를 활성화해야 할 수 있습니다. xFormers로 메모리 효율적인 어텐션을 활성화하면 메모리 사용량을 더 줄일 수 있습니다(버전 v0.0.16은 일부 GPU에서 훈련에 실패할 수 있으므로 대신 개발 버전을 설치해야 할 수 있습니다).

This guide explores the train_text_to_image_prior.py and the train_text_to_image_decoder.py scripts 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/kandinsky2_2/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_prior.py \
  --mixed_precision="fp16"

대부분의 파라미터는 Text-to-image 학습 가이드의 파라미터와 동일하므로, Kandinsky 학습 스크립트 설명으로 바로 넘어가 보겠습니다!

Min-SNR 가중

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

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

accelerate launch train_text_to_image_prior.py \
  --snr_gamma=5.0

학습 스크립트

학습 스크립트는 텍스트-투-이미지 학습 가이드와 유사하지만, prior 및 decoder 모델의 학습을 지원하도록 수정되었습니다. 이 가이드는 Kandinsky 2.2 학습 스크립트에 고유한 코드에 초점을 맞춥니다.

이전 모델

디코더 모델

main() 함수에는 데이터셋을 준비하고 모델을 학습하는 코드가 포함되어 있습니다.

가장 먼저 눈에 띄는 주요 차이점 중 하나는 학습 스크립트가 이미지 전처리를 위해 CLIPImageProcessor를, 이미지 인코딩을 위해 CLIPVisionModelWithProjection 모델을 로드한다는 것입니다. 이는 스케줄러와 토크나이저를 추가로 로드하는 것과 함께 이루어집니다.

noise_scheduler = DDPMScheduler(beta_schedule="squaredcos_cap_v2", prediction_type="sample")
image_processor = CLIPImageProcessor.from_pretrained(
    args.pretrained_prior_model_name_or_path, subfolder="image_processor"
)
tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="tokenizer")

with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="image_encoder", dtype=weight_dtype
    ).eval()
    text_encoder = CLIPTextModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="text_encoder", dtype=weight_dtype
    ).eval()

Kandinsky는 이미지 임베딩을 생성하기 위해 PriorTransformer를 사용하므로, prior 모드 매개변수를 학습하도록 옵티마이저를 설정해야 합니다.

prior = PriorTransformer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="prior")
prior.train()
optimizer = optimizer_cls(
    prior.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

다음으로, 입력 캡션은 토큰화되고 이미지는 CLIPImageProcessor에 의해 전처리됩니다:

def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[image_column]]
    examples["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
    examples["text_input_ids"], examples["text_mask"] = tokenize_captions(examples)
    return examples

마지막으로, training loop은 입력 이미지를 latents로 변환하고, 이미지 임베딩에 노이즈를 추가한 뒤 예측을 수행합니다:

model_pred = prior(
    noisy_latents,
    timestep=timesteps,
    proj_embedding=prompt_embeds,
    encoder_hidden_states=text_encoder_hidden_states,
    attention_mask=text_mask,
).predicted_image_embedding

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

스크립트 실행

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

Naruto BLIP captions 데이터셋을 학습하여 나루토 캐릭터를 생성할 수 있으며, Create a dataset for training 가이드를 따라 자체 데이터셋을 생성하고 학습할 수도 있습니다. Hub에 있는 데이터셋의 이름을 환경 변수 DATASET_NAME에 설정하거나, 자체 파일로 학습하는 경우 환경 변수 TRAIN_DIR에 데이터셋 경로를 설정하세요.

여러 개의 GPU에서 학습하는 경우 accelerate launch 명령에 --multi_gpu 매개변수를 추가하세요.

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

이전 모델

디코더 모델

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

이전 모델

디코더 모델

다음 단계

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

다양한 작업(텍스트-이미지 생성, 이미지-이미지 변환, 인페인팅, 보간)에 Kandinsky를 사용하는 방법과 ControlNet과 결합하는 방법을 배우려면 가이드를 읽어보세요.
DreamBooth 및 LoRA 학습 가이드를 참고하여 몇 장의 예시 이미지만으로 개인화된 Kandinsky 모델을 학습하는 방법을 배울 수 있습니다. 이 두 가지 학습 기법은 함께 사용할 수도 있습니다!

Update on GitHub