Stable Diffusion XL

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.

Stable Diffusion XL (SDXL) adalah iterasi yang lebih besar dan lebih kuat dari model Stable Diffusion, yang mampu menghasilkan gambar dengan resolusi lebih tinggi.

UNet dari SDXL berukuran 3x lebih besar dan model ini menambahkan encoder teks kedua ke dalam arsitekturnya. Tergantung pada perangkat keras yang tersedia, hal ini dapat sangat intensif secara komputasi dan mungkin tidak dapat berjalan pada GPU konsumen seperti Tesla T4. Untuk membantu memuat model yang lebih besar ini ke dalam memori dan mempercepat pelatihan, coba aktifkan gradient_checkpointing, mixed_precision, dan gradient_accumulation_steps. Anda dapat mengurangi penggunaan memori lebih lanjut dengan mengaktifkan perhatian yang hemat memori menggunakan xFormers dan menggunakan optimizer 8-bit dari bitsandbytes.

Panduan ini akan membahas skrip pelatihan train_text_to_image_sdxl.py untuk membantu Anda lebih mengenalnya, serta bagaimana Anda dapat menyesuaikannya untuk kasus penggunaan 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/text_to_image
pip install -r requirements_sdxl.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.

Parameter skrip

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

Skrip pelatihan menyediakan banyak parameter untuk membantu Anda menyesuaikan eksekusi pelatihan Anda. Semua parameter dan deskripsinya dapat ditemukan di fungsi parse_args(). Fungsi ini 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 bf16, tambahkan parameter --mixed_precision ke perintah pelatihan:

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

Sebagian besar parameter identik dengan parameter pada panduan pelatihan Text-to-image, sehingga Anda akan berfokus pada parameter yang relevan untuk pelatihan SDXL dalam panduan ini.

• --pretrained_vae_model_name_or_path: path ke VAE pretrained; VAE SDXL diketahui mengalami ketidakstabilan numerik, sehingga parameter ini memungkinkan Anda untuk menentukan VAE yang lebih baik
• --proportion_empty_prompts: proporsi prompt gambar yang akan diganti dengan string kosong
• --timestep_bias_strategy: di mana (awal vs. akhir) dalam timestep untuk menerapkan bias, yang dapat mendorong model untuk mempelajari detail frekuensi rendah atau tinggi
• --timestep_bias_multiplier: bobot bias yang diterapkan pada timestep
• --timestep_bias_begin: timestep untuk mulai menerapkan bias
• --timestep_bias_end: timestep untuk mengakhiri penerapan bias
• --timestep_bias_portion: proporsi timestep untuk menerapkan bias

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_sdxl.py \
  --snr_gamma=5.0

Skrip pelatihan

Skrip pelatihan juga mirip dengan panduan pelatihan Text-to-image, tetapi telah dimodifikasi untuk mendukung pelatihan SDXL. Panduan ini akan berfokus pada kode yang unik untuk skrip pelatihan SDXL.

Proses dimulai dengan membuat fungsi untuk meng-tokenisasi prompt guna menghitung embedding prompt, serta untuk menghitung embedding gambar dengan VAE. Selanjutnya, Anda akan membuat fungsi untuk menghasilkan bobot timestep berdasarkan jumlah timestep dan strategi bias timestep yang akan diterapkan.

Di dalam fungsi main(), selain memuat tokenizer, skrip ini juga memuat tokenizer kedua dan text encoder karena arsitektur SDXL menggunakan dua dari masing-masing:

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

Embedding prompt dan gambar dihitung terlebih dahulu dan disimpan di memori, yang biasanya bukan masalah untuk dataset yang lebih kecil, tetapi untuk dataset yang lebih besar dapat menyebabkan masalah memori. Jika ini terjadi, Anda sebaiknya menyimpan embedding yang telah dihitung sebelumnya ke disk secara terpisah dan memuatnya ke memori selama proses pelatihan (lihat PR ini untuk diskusi lebih lanjut tentang topik ini).

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

Setelah menghitung embedding, text encoder, VAE, dan tokenizer dihapus untuk membebaskan sebagian memori:

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

Terakhir, training loop menangani sisanya. Jika Anda memilih untuk menerapkan strategi bias timestep, Anda akan melihat bobot timestep dihitung dan ditambahkan sebagai noise:

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)

Jika Anda ingin mempelajari lebih lanjut tentang cara kerja loop pelatihan, lihat tutorial Memahami pipeline, model, dan scheduler yang menguraikan pola dasar dari proses denoising.

Jalankan skrip

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

Mari latih pada dataset Naruto BLIP captions untuk menghasilkan karakter Naruto Anda sendiri. Atur variabel lingkungan MODEL_NAME dan DATASET_NAME ke model dan dataset (baik dari Hub maupun jalur lokal). Anda juga harus menentukan VAE selain SDXL VAE (baik dari Hub maupun jalur lokal) dengan VAE_NAME untuk menghindari ketidakstabilan numerik.

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

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

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

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

Langkah selanjutnya

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

• Baca panduan Stable Diffusion XL untuk mempelajari cara menggunakannya untuk berbagai tugas yang berbeda (teks-ke-gambar, gambar-ke-gambar, inpainting), cara menggunakan model refiner-nya, dan berbagai jenis micro-conditioning.
• Lihat panduan pelatihan DreamBooth dan LoRA untuk mempelajari cara melatih model SDXL yang dipersonalisasi hanya dengan beberapa gambar contoh. Kedua teknik pelatihan ini bahkan dapat digabungkan!