Kandinsky 2.2

Skrip ini bersifat eksperimental, dan sangat mudah mengalami overfitting serta masalah seperti catastrophic forgetting. Coba jelajahi berbagai hyperparameter untuk mendapatkan hasil terbaik pada dataset Anda.

Kandinsky 2.2 adalah model teks-ke-gambar multibahasa yang mampu menghasilkan gambar yang lebih fotorealistis. Model ini mencakup model prior gambar untuk membuat embedding gambar dari prompt teks, dan model decoder yang menghasilkan gambar berdasarkan embedding dari model prior. Itulah mengapa Anda akan menemukan dua skrip terpisah di Diffusers untuk Kandinsky 2.2, satu untuk melatih model prior dan satu untuk melatih model decoder. Anda dapat melatih kedua model secara terpisah, tetapi untuk mendapatkan hasil terbaik, Anda harus melatih model prior dan decoder.

Tergantung pada GPU Anda, Anda mungkin perlu mengaktifkan gradient_checkpointing (⚠️ tidak didukung untuk model sebelumnya!), mixed_precision, dan gradient_accumulation_steps untuk membantu memuat model ke dalam memori dan mempercepat pelatihan. Anda dapat mengurangi penggunaan memori lebih lanjut dengan mengaktifkan perhatian yang hemat memori menggunakan xFormers (versi v0.0.16 gagal untuk pelatihan pada beberapa GPU, jadi Anda mungkin perlu menginstal versi pengembangan sebagai gantinya).

Panduan ini membahas skrip train_text_to_image_prior.py dan train_text_to_image_decoder.py untuk membantu Anda lebih memahami alurnya, serta bagaimana Anda dapat menyesuaikannya untuk kebutuhan Anda sendiri.

Sebelum menjalankan skrip, pastikan Anda menginstal pustaka dari sumber:

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

Lalu, navigasikan ke folder contoh yang berisi skrip pelatihan dan instal dependensi yang diperlukan untuk skrip yang Anda gunakan:

cd examples/kandinsky2_2/text_to_image
pip install -r requirements.txt

🤗 Accelerate adalah pustaka untuk membantu Anda melatih pada beberapa GPU/TPU atau dengan presisi campuran. Pustaka ini akan mengonfigurasi pengaturan pelatihan Anda secara otomatis berdasarkan perangkat keras dan lingkungan Anda. Lihat Tour singkat 🤗 Accelerate untuk mempelajari lebih lanjut.

Inisialisasi lingkungan 🤗 Accelerate:

accelerate config

Untuk menyiapkan lingkungan 🤗 Accelerate default tanpa memilih konfigurasi apa pun:

accelerate config default

Atau jika lingkungan Anda tidak mendukung shell interaktif, seperti notebook, Anda dapat menggunakan:

from accelerate.utils import write_basic_config

write_basic_config()

Terakhir, jika Anda ingin melatih model pada dataset Anda sendiri, lihat panduan Membuat dataset untuk pelatihan untuk mempelajari cara membuat dataset yang berfungsi dengan skrip pelatihan.

Bagian-bagian berikut menyoroti bagian dari skrip pelatihan yang penting untuk dipahami agar Anda dapat memodifikasinya, tetapi tidak mencakup setiap aspek skrip secara detail. Jika Anda tertarik untuk mempelajari lebih lanjut, silakan membaca skripnya dan beri tahu kami jika Anda memiliki pertanyaan atau kekhawatiran.

Parameter skrip

Skrip pelatihan menyediakan banyak parameter untuk membantu Anda menyesuaikan eksekusi pelatihan Anda. Semua parameter dan deskripsinya dapat ditemukan di fungsi parse_args(). Skrip pelatihan menyediakan nilai default untuk setiap parameter, seperti ukuran batch pelatihan dan learning rate, tetapi Anda juga dapat mengatur nilai Anda sendiri dalam perintah pelatihan jika Anda mau.

Sebagai contoh, untuk mempercepat pelatihan dengan presisi campuran menggunakan format fp16, tambahkan parameter --mixed_precision ke perintah pelatihan:

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

Sebagian besar parameter identik dengan parameter dalam panduan pelatihan Text-to-image, jadi mari langsung saja ke walkthrough skrip pelatihan Kandinsky!

Penimbangan Min-SNR

Strategi penimbangan Min-SNR dapat membantu pelatihan dengan menyeimbangkan ulang loss untuk mencapai konvergensi yang lebih cepat. Skrip pelatihan mendukung prediksi epsilon (noise) atau v_prediction, tetapi Min-SNR kompatibel dengan kedua jenis prediksi tersebut. Strategi penimbangan ini hanya didukung oleh PyTorch.

Tambahkan parameter --snr_gamma dan atur nilainya ke nilai yang direkomendasikan, yaitu 5.0:

accelerate launch train_text_to_image_prior.py \
  --snr_gamma=5.0

Skrip pelatihan

Skrip pelatihan juga mirip dengan panduan pelatihan Text-to-image, tetapi telah dimodifikasi untuk mendukung pelatihan model prior dan decoder. Panduan ini berfokus pada kode yang unik untuk skrip pelatihan Kandinsky 2.2.

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()

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,
)

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

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

with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
    vae = VQModel.from_pretrained(
        args.pretrained_decoder_model_name_or_path, subfolder="movq", dtype=weight_dtype
    ).eval()
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="image_encoder", dtype=weight_dtype
    ).eval()
unet = UNet2DConditionModel.from_pretrained(args.pretrained_decoder_model_name_or_path, subfolder="unet")

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["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
    return examples

model_pred = unet(noisy_latents, timesteps, None, added_cond_kwargs=added_cond_kwargs).sample[:, :4]

Jalankan skrip

Setelah Anda melakukan semua perubahan atau merasa puas dengan konfigurasi default, Anda siap untuk menjalankan skrip pelatihan! 🚀

Anda akan melatih pada dataset Naruto BLIP captions untuk menghasilkan karakter Naruto Anda sendiri, tetapi Anda juga dapat membuat dan melatih pada dataset Anda sendiri dengan mengikuti panduan Create a dataset for training. Atur variabel lingkungan DATASET_NAME ke nama dataset di Hub atau jika Anda melatih pada file Anda sendiri, atur variabel lingkungan TRAIN_DIR ke path dataset Anda.

Jika Anda melatih pada lebih dari satu GPU, tambahkan parameter --multi_gpu ke perintah accelerate launch.

Untuk memantau kemajuan pelatihan dengan Weights & Biases, tambahkan parameter --report_to=wandb ke perintah pelatihan. Anda juga perlu menambahkan --validation_prompt ke perintah pelatihan untuk melacak hasil. Hal ini sangat berguna untuk debugging model dan melihat hasil antara.

export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_prior.py \
  --dataset_name=$DATASET_NAME \
  --resolution=768 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --checkpoints_total_limit=3 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --validation_prompts="A robot naruto, 4k photo" \
  --report_to="wandb" \
  --push_to_hub \
  --output_dir="kandi2-prior-naruto-model"

export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_decoder.py \
  --dataset_name=$DATASET_NAME \
  --resolution=768 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --checkpoints_total_limit=3 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --validation_prompts="A robot naruto, 4k photo" \
  --report_to="wandb" \
  --push_to_hub \
  --output_dir="kandi2-decoder-naruto-model"

Setelah pelatihan selesai, Anda dapat menggunakan model yang baru dilatih untuk inferensi!

from diffusers import AutoPipelineForText2Image, DiffusionPipeline
import torch

prior_pipeline = DiffusionPipeline.from_pretrained(output_dir, dtype=torch.float16)
prior_components = {"prior_" + k: v for k,v in prior_pipeline.components.items()}
pipeline = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", **prior_components, dtype=torch.float16)

pipe.enable_model_cpu_offload()
prompt="A robot naruto, 4k photo"
image = pipeline(prompt=prompt, negative_prompt=negative_prompt).images[0]

from diffusers import AutoPipelineForText2Image
import torch

pipeline = AutoPipelineForText2Image.from_pretrained("path/to/saved/model", dtype=torch.float16)
pipeline.enable_model_cpu_offload()

prompt="A robot naruto, 4k photo"
image = pipeline(prompt=prompt).images[0]

from diffusers import AutoPipelineForText2Image, UNet2DConditionModel

unet = UNet2DConditionModel.from_pretrained("path/to/saved/model" + "/checkpoint-<N>/unet")

pipeline = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", unet=unet, dtype=torch.float16)
pipeline.enable_model_cpu_offload()

image = pipeline(prompt="A robot naruto, 4k photo").images[0]

Langkah selanjutnya

Selamat karena Anda telah melatih model Kandinsky 2.2! Untuk mempelajari lebih lanjut tentang cara menggunakan model baru Anda, panduan berikut mungkin dapat membantu:

• Baca panduan Kandinsky untuk mempelajari cara menggunakannya untuk berbagai tugas yang berbeda (teks-ke-gambar, gambar-ke-gambar, inpainting, interpolasi), dan bagaimana ia dapat dikombinasikan dengan ControlNet.
• Lihat panduan pelatihan DreamBooth dan LoRA untuk mempelajari cara melatih model Kandinsky yang dipersonalisasi dengan hanya beberapa gambar contoh. Kedua teknik pelatihan ini bahkan dapat digabungkan!