Custom Diffusion

Custom Diffusion adalah teknik pelatihan untuk mempersonalisasi model generasi gambar. Seperti Textual Inversion, DreamBooth, dan LoRA, Custom Diffusion hanya membutuhkan beberapa (~4-5) gambar contoh. Teknik ini bekerja dengan hanya melatih bobot pada lapisan cross-attention, dan menggunakan kata khusus untuk merepresentasikan konsep yang baru dipelajari. Custom Diffusion unik karena juga dapat mempelajari beberapa konsep secara bersamaan.

Jika Anda melatih pada GPU dengan vRAM terbatas, Anda sebaiknya mencoba mengaktifkan xFormers dengan --enable_xformers_memory_efficient_attention untuk pelatihan yang lebih cepat dengan kebutuhan vRAM yang lebih rendah (16GB). Untuk menghemat lebih banyak memori, tambahkan --set_grads_to_none pada argumen pelatihan untuk mengatur gradien menjadi None alih-alih nol (opsi ini dapat menyebabkan beberapa masalah, jadi jika Anda mengalaminya, coba hapus parameter ini).

Panduan ini akan membahas skrip train_custom_diffusion.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 .

Beralih ke folder contoh dengan skrip pelatihan dan instal dependensi yang diperlukan:

cd examples/custom_diffusion
pip install -r requirements.txt
pip install clip-retrieval

🤗 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, 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.

Parameter skrip

Skrip pelatihan berisi semua parameter untuk membantu Anda menyesuaikan eksekusi pelatihan Anda. Parameter ini dapat ditemukan di fungsi parse_args(). Fungsi ini dilengkapi dengan nilai default, tetapi Anda juga dapat mengatur nilai Anda sendiri dalam perintah pelatihan jika Anda mau.

Sebagai contoh, untuk mengubah resolusi gambar input:

accelerate launch train_custom_diffusion.py \
  --resolution=256

Banyak parameter dasar dijelaskan dalam panduan pelatihan DreamBooth, sehingga panduan ini berfokus pada parameter yang unik untuk Custom Diffusion:

• --freeze_model: membekukan parameter kunci dan nilai pada lapisan cross-attention; nilai defaultnya adalah crossattn_kv, tetapi Anda dapat mengaturnya menjadi crossattn untuk melatih semua parameter pada lapisan cross-attention
• --concepts_list: untuk mempelajari beberapa konsep, berikan path ke file JSON yang berisi konsep-konsep tersebut
• --modifier_token: kata khusus yang digunakan untuk merepresentasikan konsep yang dipelajari
• --initializer_token: kata khusus yang digunakan untuk menginisialisasi embedding dari modifier_token

Prior preservation loss

Prior preservation loss adalah metode yang menggunakan sampel yang dihasilkan oleh model itu sendiri untuk membantunya belajar cara menghasilkan gambar yang lebih beragam. Karena gambar sampel yang dihasilkan ini termasuk dalam kelas yang sama dengan gambar yang Anda berikan, mereka membantu model mempertahankan apa yang telah dipelajarinya tentang kelas tersebut dan bagaimana ia dapat menggunakan pengetahuan yang sudah dimilikinya tentang kelas tersebut untuk membuat komposisi baru.

Banyak parameter untuk prior preservation loss dijelaskan dalam panduan pelatihan DreamBooth.

Regularisasi

Custom Diffusion mencakup pelatihan gambar target dengan set kecil gambar asli untuk mencegah overfitting. Seperti yang bisa Anda bayangkan, ini mudah dilakukan ketika Anda hanya melatih pada beberapa gambar! Unduh 200 gambar asli dengan clip_retrieval. class_prompt harus berupa kategori yang sama dengan gambar target. Gambar-gambar ini disimpan di class_data_dir.

python retrieve.py --class_prompt cat --class_data_dir real_reg/samples_cat --num_class_images 200

Untuk mengaktifkan regularisasi, tambahkan parameter berikut:

• --with_prior_preservation: apakah akan menggunakan prior preservation loss
• --prior_loss_weight: mengontrol pengaruh prior preservation loss pada model
• --real_prior: apakah akan menggunakan kumpulan kecil gambar asli untuk mencegah overfitting

accelerate launch train_custom_diffusion.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="./real_reg/samples_cat" \
  --class_prompt="cat" \
  --real_prior=True \

Skrip pelatihan

Banyak kode dalam skrip pelatihan Custom Diffusion mirip dengan skrip DreamBooth. Panduan ini justru berfokus pada kode yang relevan dengan Custom Diffusion.

Skrip pelatihan Custom Diffusion memiliki dua kelas dataset:

• CustomDiffusionDataset: memproses prapenanganan gambar, gambar kelas, dan prompt untuk pelatihan
• PromptDataset: menyiapkan prompt untuk menghasilkan gambar kelas

Selanjutnya, modifier_token ditambahkan ke tokenizer, dikonversi menjadi id token, dan ukuran embedding token disesuaikan untuk mengakomodasi modifier_token baru. Kemudian, embedding modifier_token diinisialisasi dengan embedding dari initializer_token. Semua parameter dalam text encoder dibekukan, kecuali embedding token, karena inilah yang model coba pelajari untuk dikaitkan dengan konsep-konsep tersebut.

params_to_freeze = itertools.chain(
    text_encoder.text_model.encoder.parameters(),
    text_encoder.text_model.final_layer_norm.parameters(),
    text_encoder.text_model.embeddings.position_embedding.parameters(),
)
freeze_params(params_to_freeze)

Sekarang Anda perlu menambahkan Custom Diffusion weights ke lapisan attention. Ini adalah langkah yang sangat penting untuk mendapatkan bentuk dan ukuran bobot attention yang benar, serta untuk mengatur jumlah processor attention yang tepat di setiap blok UNet.

st = unet.state_dict()
for name, _ in unet.attn_processors.items():
    cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
    if name.startswith("mid_block"):
        hidden_size = unet.config.block_out_channels[-1]
    elif name.startswith("up_blocks"):
        block_id = int(name[len("up_blocks.")])
        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
    elif name.startswith("down_blocks"):
        block_id = int(name[len("down_blocks.")])
        hidden_size = unet.config.block_out_channels[block_id]
    layer_name = name.split(".processor")[0]
    weights = {
        "to_k_custom_diffusion.weight": st[layer_name + ".to_k.weight"],
        "to_v_custom_diffusion.weight": st[layer_name + ".to_v.weight"],
    }
    if train_q_out:
        weights["to_q_custom_diffusion.weight"] = st[layer_name + ".to_q.weight"]
        weights["to_out_custom_diffusion.0.weight"] = st[layer_name + ".to_out.0.weight"]
        weights["to_out_custom_diffusion.0.bias"] = st[layer_name + ".to_out.0.bias"]
    if cross_attention_dim is not None:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=train_kv,
            train_q_out=train_q_out,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        ).to(unet.device)
        custom_diffusion_attn_procs[name].load_state_dict(weights)
    else:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=False,
            train_q_out=False,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        )
del st
unet.set_attn_processor(custom_diffusion_attn_procs)
custom_diffusion_layers = AttnProcsLayers(unet.attn_processors)

optimizer diinisialisasi untuk memperbarui parameter lapisan cross-attention:

optimizer = optimizer_class(
    itertools.chain(text_encoder.get_input_embeddings().parameters(), custom_diffusion_layers.parameters())
    if args.modifier_token is not None
    else custom_diffusion_layers.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Dalam loop pelatihan, penting untuk hanya memperbarui embedding untuk konsep yang sedang Anda pelajari. Ini berarti mengatur gradien dari semua embedding token lainnya menjadi nol:

if args.modifier_token is not None:
    if accelerator.num_processes > 1:
        grads_text_encoder = text_encoder.module.get_input_embeddings().weight.grad
    else:
        grads_text_encoder = text_encoder.get_input_embeddings().weight.grad
    index_grads_to_zero = torch.arange(len(tokenizer)) != modifier_token_id[0]
    for i in range(len(modifier_token_id[1:])):
        index_grads_to_zero = index_grads_to_zero & (
            torch.arange(len(tokenizer)) != modifier_token_id[i]
        )
    grads_text_encoder.data[index_grads_to_zero, :] = grads_text_encoder.data[
        index_grads_to_zero, :
    ].fill_(0)

Jalankan skrip

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

Dalam panduan ini, Anda akan mengunduh dan menggunakan contoh gambar kucing ini. Anda juga dapat membuat dan menggunakan dataset Anda sendiri jika diinginkan (lihat panduan Membuat dataset untuk pelatihan).

Setel variabel lingkungan MODEL_NAME ke id model di Hub atau path ke model lokal, INSTANCE_DIR ke path tempat Anda baru saja mengunduh gambar kucing, dan OUTPUT_DIR ke tempat Anda ingin menyimpan model. Anda akan menggunakan <new1> sebagai kata khusus untuk mengaitkan embedding yang baru dipelajari. Skrip ini membuat dan menyimpan checkpoint model serta file pytorch_custom_diffusion_weights.bin ke repositori Anda.

Untuk memantau kemajuan pelatihan dengan Weights and Biases, tambahkan parameter --report_to=wandb ke perintah pelatihan dan tentukan prompt validasi dengan --validation_prompt. Hal ini berguna untuk debugging dan menyimpan hasil antara.

Jika Anda melatih pada wajah manusia, tim Custom Diffusion telah menemukan bahwa parameter berikut bekerja dengan baik:

• --learning_rate=5e-6
• --max_train_steps dapat berada di antara 1000 dan 2000
• --freeze_model=crossattn
• gunakan setidaknya 15-20 gambar untuk pelatihan

export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export OUTPUT_DIR="path-to-save-model"
export INSTANCE_DIR="./data/cat"

accelerate launch train_custom_diffusion.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --class_data_dir=./real_reg/samples_cat/ \
  --with_prior_preservation \
  --real_prior \
  --prior_loss_weight=1.0 \
  --class_prompt="cat" \
  --num_class_images=200 \
  --instance_prompt="photo of a <new1> cat"  \
  --resolution=512  \
  --train_batch_size=2  \
  --learning_rate=1e-5  \
  --lr_warmup_steps=0 \
  --max_train_steps=250 \
  --scale_lr \
  --hflip  \
  --modifier_token "<new1>" \
  --validation_prompt="<new1> cat sitting in a bucket" \
  --report_to="wandb" \
  --push_to_hub

pip install clip-retrieval
python retrieve.py --class_prompt {} --class_data_dir {} --num_class_images 200

export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export OUTPUT_DIR="path-to-save-model"

accelerate launch train_custom_diffusion.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --output_dir=$OUTPUT_DIR \
  --concepts_list=./concept_list.json \
  --with_prior_preservation \
  --real_prior \
  --prior_loss_weight=1.0 \
  --resolution=512  \
  --train_batch_size=2  \
  --learning_rate=1e-5  \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --num_class_images=200 \
  --scale_lr \
  --hflip  \
  --modifier_token "<new1>+<new2>" \
  --push_to_hub

Setelah pelatihan selesai, Anda dapat menggunakan model Custom Diffusion baru Anda untuk inferensi.

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet.load_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin")
pipeline.load_textual_inversion("path-to-save-model", weight_name="<new1>.bin")

image = pipeline(
    "<new1> cat sitting in a bucket",
    num_inference_steps=100,
    guidance_scale=6.0,
    eta=1.0,
).images[0]
image.save("cat.png")

import torch
from huggingface_hub.repocard import RepoCard
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
model_id = "sayakpaul/custom-diffusion-cat-wooden-pot"
pipeline.unet.load_attn_procs(model_id, weight_name="pytorch_custom_diffusion_weights.bin")
pipeline.load_textual_inversion(model_id, weight_name="<new1>.bin")
pipeline.load_textual_inversion(model_id, weight_name="<new2>.bin")

image = pipeline(
    "the <new1> cat sculpture in the style of a <new2> wooden pot",
    num_inference_steps=100,
    guidance_scale=6.0,
    eta=1.0,
).images[0]
image.save("multi-subject.png")

Langkah selanjutnya

Selamat karena Anda telah melatih model dengan Custom Diffusion! 🎉 Untuk mempelajari lebih lanjut:

• Baca artikel blog Multi-Concept Customization of Text-to-Image Diffusion untuk mempelajari lebih lanjut tentang hasil eksperimen dari tim Custom Diffusion.