Distilasi Konsistensi Laten

Latent Consistency Models (LCMs) mampu menghasilkan gambar berkualitas tinggi hanya dalam beberapa langkah, yang merupakan lompatan besar karena banyak pipeline membutuhkan setidaknya 25+ langkah. LCMs dihasilkan dengan menerapkan metode latent consistency distillation pada model Stable Diffusion mana pun. Metode ini bekerja dengan menerapkan one-stage guided distillation pada ruang laten, serta mengintegrasikan metode skipping-step untuk secara konsisten melewati timestep guna mempercepat proses distilasi (lihat bagian 4.1, 4.2, dan 4.3 dari makalah untuk detail lebih lanjut).

Jika Anda melatih pada GPU dengan vRAM terbatas, coba aktifkan gradient_checkpointing, gradient_accumulation_steps, dan mixed_precision untuk mengurangi penggunaan memori dan mempercepat pelatihan. Anda dapat mengurangi penggunaan memori lebih lanjut dengan mengaktifkan perhatian hemat memori menggunakan xFormers dan optimizer 8-bit dari bitsandbytes.

Panduan ini akan membahas skrip train_lcm_distill_sd_wds.py untuk membantu Anda lebih mengenalnya, 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/consistency_distillation
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 (coba aktifkan torch.compile untuk mempercepat pelatihan secara signifikan):

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 fp16, tambahkan parameter --mixed_precision ke perintah pelatihan:

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

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

• --pretrained_teacher_model: jalur ke model latent diffusion terlatih sebelumnya yang akan digunakan sebagai model guru
• --pretrained_vae_model_name_or_path: path ke VAE pretrained; VAE SDXL diketahui mengalami ketidakstabilan numerik, sehingga parameter ini memungkinkan Anda untuk menentukan VAE alternatif (seperti VAE ini) yang dibuat oleh madebyollin dan berfungsi dalam fp16)
• --w_min dan --w_max: nilai minimum dan maksimum guidance scale untuk sampling guidance scale
• --num_ddim_timesteps: jumlah timestep untuk sampling DDIM
• --loss_type: jenis loss (L2 atau Huber) yang dihitung untuk latent consistency distillation; Huber loss umumnya lebih disukai karena lebih robust terhadap outlier
• --huber_c: parameter loss Huber

Skrip pelatihan

Skrip pelatihan dimulai dengan membuat kelas dataset - Text2ImageDataset - untuk prapemrosesan gambar dan membuat dataset pelatihan.

def transform(example):
    image = example["image"]
    image = TF.resize(image, resolution, interpolation=transforms.InterpolationMode.BILINEAR)

    c_top, c_left, _, _ = transforms.RandomCrop.get_params(image, output_size=(resolution, resolution))
    image = TF.crop(image, c_top, c_left, resolution, resolution)
    image = TF.to_tensor(image)
    image = TF.normalize(image, [0.5], [0.5])

    example["image"] = image
    return example

Untuk meningkatkan kinerja saat membaca dan menulis dataset besar yang disimpan di cloud, skrip ini menggunakan format WebDataset untuk membuat pipeline pra-pemrosesan guna menerapkan transformasi serta membuat dataset dan dataloader untuk pelatihan. Gambar diproses dan diberikan ke loop pelatihan tanpa perlu mengunduh seluruh dataset terlebih dahulu.

processing_pipeline = [
    wds.decode("pil", handler=wds.ignore_and_continue),
    wds.rename(image="jpg;png;jpeg;webp", text="text;txt;caption", handler=wds.warn_and_continue),
    wds.map(filter_keys({"image", "text"})),
    wds.map(transform),
    wds.to_tuple("image", "text"),
]

Dalam fungsi main(), semua komponen yang diperlukan seperti noise scheduler, tokenizers, text encoders, dan VAE dimuat. Teacher UNet juga dimuat di sini, kemudian Anda dapat membuat student UNet dari teacher UNet. Student UNet diperbarui oleh optimizer selama pelatihan.

teacher_unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_teacher_model, subfolder="unet", revision=args.teacher_revision
)

unet = UNet2DConditionModel(**teacher_unet.config)
unet.load_state_dict(teacher_unet.state_dict(), strict=False)
unet.train()

Sekarang Anda dapat membuat optimizer untuk memperbarui parameter UNet:

optimizer = optimizer_class(
    unet.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Buat dataset:

dataset = Text2ImageDataset(
    train_shards_path_or_url=args.train_shards_path_or_url,
    num_train_examples=args.max_train_samples,
    per_gpu_batch_size=args.train_batch_size,
    global_batch_size=args.train_batch_size * accelerator.num_processes,
    num_workers=args.dataloader_num_workers,
    resolution=args.resolution,
    shuffle_buffer_size=1000,
    pin_memory=True,
    persistent_workers=True,
)
train_dataloader = dataset.train_dataloader

Selanjutnya, Anda siap untuk menyiapkan training loop dan mengimplementasikan metode latent consistency distillation (lihat Algoritma 1 di paper untuk detail lebih lanjut). Bagian skrip ini menangani penambahan noise ke latents, sampling dan pembuatan guidance scale embedding, serta prediksi gambar asli dari noise.

pred_x_0 = predicted_origin(
    noise_pred,
    start_timesteps,
    noisy_model_input,
    noise_scheduler.config.prediction_type,
    alpha_schedule,
    sigma_schedule,
)

model_pred = c_skip_start * noisy_model_input + c_out_start * pred_x_0

Selanjutnya, mengambil prediksi model guru dan prediksi LCM, menghitung loss, lalu melakukan backpropagation ke LCM.

if args.loss_type == "l2":
    loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
elif args.loss_type == "huber":
    loss = torch.mean(
        torch.sqrt((model_pred.float() - target.float()) ** 2 + args.huber_c**2) - args.huber_c
    )

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

Sekarang Anda siap untuk menjalankan skrip pelatihan dan mulai melakukan distilasi!

Untuk panduan ini, Anda akan menggunakan --train_shards_path_or_url untuk menentukan path ke dataset Conceptual Captions 12M yang disimpan di Hub di sini. Setel variabel lingkungan MODEL_DIR ke nama model teacher dan OUTPUT_DIR ke lokasi tempat Anda ingin menyimpan model.

export MODEL_DIR="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="path/to/saved/model"

accelerate launch train_lcm_distill_sd_wds.py \
    --pretrained_teacher_model=$MODEL_DIR \
    --output_dir=$OUTPUT_DIR \
    --mixed_precision=fp16 \
    --resolution=512 \
    --learning_rate=1e-6 --loss_type="huber" --ema_decay=0.95 --adam_weight_decay=0.0 \
    --max_train_steps=1000 \
    --max_train_samples=4000000 \
    --dataloader_num_workers=8 \
    --train_shards_path_or_url="pipe:curl -L -s https://huggingface.co/datasets/laion/conceptual-captions-12m-webdataset/resolve/main/data/{00000..01099}.tar?download=true" \
    --validation_steps=200 \
    --checkpointing_steps=200 --checkpoints_total_limit=10 \
    --train_batch_size=12 \
    --gradient_checkpointing --enable_xformers_memory_efficient_attention \
    --gradient_accumulation_steps=1 \
    --use_8bit_adam \
    --resume_from_checkpoint=latest \
    --report_to=wandb \
    --seed=453645634 \
    --push_to_hub

Setelah pelatihan selesai, Anda dapat menggunakan LCM baru Anda untuk inferensi.

from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
import torch

unet = UNet2DConditionModel.from_pretrained("your-username/your-model", dtype=torch.float16, variant="fp16")
pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, dtype=torch.float16, variant="fp16")

pipeline.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipeline.to("cuda")  # or "mps", "xpu", "cpu"

prompt = "sushi rolls in the form of panda heads, sushi platter"

image = pipeline(prompt, num_inference_steps=4, guidance_scale=1.0).images[0]

LoRA

LoRA adalah teknik pelatihan untuk mengurangi secara signifikan jumlah parameter yang dapat dilatih. Akibatnya, pelatihan menjadi lebih cepat dan lebih mudah untuk menyimpan bobot yang dihasilkan karena ukurannya jauh lebih kecil (~100MB). Gunakan skrip train_lcm_distill_lora_sd_wds.py atau train_lcm_distill_lora_sdxl.wds.py untuk melatih dengan LoRA.

Skrip pelatihan LoRA dibahas lebih detail dalam panduan pelatihan LoRA.

Stable Diffusion XL

Stable Diffusion XL (SDXL) adalah model teks-ke-gambar yang kuat yang menghasilkan gambar resolusi tinggi, dan menambahkan text-encoder kedua ke arsitekturnya. Gunakan skrip train_lcm_distill_sdxl_wds.py untuk melatih model SDXL dengan LoRA.

Skrip pelatihan SDXL dibahas lebih detail dalam panduan pelatihan SDXL.

Langkah selanjutnya

Selamat karena Anda berhasil mendistilasi model LCM! Untuk mempelajari lebih lanjut tentang LCM, berikut ini mungkin akan membantu:

• Pelajari cara menggunakan LCMs untuk inferensi untuk teks-ke-gambar, gambar-ke-gambar, dan dengan checkpoint LoRA.
• Baca artikel blog SDXL dalam 4 langkah dengan Latent Consistency LoRAs untuk mempelajari lebih lanjut tentang SDXL LCM-LoRA untuk inferensi super cepat, perbandingan kualitas, benchmark, dan lainnya.