ControlNet

Model ControlNet adalah adapter yang dilatih di atas model pretrained lainnya. Hal ini memungkinkan tingkat kendali yang lebih besar terhadap generasi gambar dengan mengondisikan model menggunakan gambar input tambahan. Gambar input tersebut dapat berupa canny edge, depth map, human pose, dan banyak lagi.

Jika Anda melatih pada GPU dengan vRAM terbatas, Anda sebaiknya mencoba mengaktifkan parameter gradient_checkpointing, gradient_accumulation_steps, dan mixed_precision dalam perintah pelatihan. Anda juga dapat mengurangi jejak memori dengan menggunakan perhatian yang hemat memori dengan xFormers.

Panduan ini akan membahas skrip pelatihan train_controlnet.py untuk membantu Anda memahaminya, 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/controlnet
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, 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 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_controlnet.py \
  --mixed_precision="fp16"

Banyak parameter dasar dan penting dijelaskan dalam panduan pelatihan Text-to-image, sehingga panduan ini hanya berfokus pada parameter yang relevan untuk ControlNet:

• --max_train_samples: jumlah sampel pelatihan; ini dapat dikurangi untuk pelatihan yang lebih cepat, tetapi jika Anda ingin streaming dataset yang sangat besar, Anda perlu menyertakan parameter ini dan parameter --streaming dalam perintah pelatihan Anda
• --gradient_accumulation_steps: jumlah langkah pembaruan yang akan diakumulasi sebelum backward pass; ini memungkinkan Anda untuk melatih dengan ukuran batch yang lebih besar dari yang biasanya dapat ditangani oleh memori GPU Anda

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

Skrip pelatihan

Sama seperti parameter skrip, panduan umum mengenai skrip pelatihan tersedia di panduan pelatihan Text-to-image. Sebagai gantinya, panduan ini membahas bagian-bagian yang relevan dari skrip ControlNet.

Skrip pelatihan memiliki fungsi make_train_dataset untuk prapemrosesan dataset dengan transformasi gambar dan tokenisasi caption. Anda akan melihat bahwa selain tokenisasi caption dan transformasi gambar yang biasa, skrip ini juga mencakup transformasi untuk gambar kondisional.

Jika Anda melakukan streaming dataset di TPU, kinerja mungkin terhambat oleh pustaka 🤗 Datasets yang tidak dioptimalkan untuk gambar. Untuk memastikan throughput maksimum, Anda disarankan untuk menjelajahi format dataset lain seperti WebDataset, TorchData, dan TensorFlow Datasets.

conditioning_image_transforms = transforms.Compose(
    [
        transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
        transforms.CenterCrop(args.resolution),
        transforms.ToTensor(),
    ]
)

Di dalam fungsi main(), Anda akan menemukan kode untuk memuat tokenizer, text encoder, scheduler, dan model. Di sinilah model ControlNet dimuat, baik dari bobot yang sudah ada maupun diinisialisasi secara acak dari UNet:

if args.controlnet_model_name_or_path:
    logger.info("Loading existing controlnet weights")
    controlnet = ControlNetModel.from_pretrained(args.controlnet_model_name_or_path)
else:
    logger.info("Initializing controlnet weights from unet")
    controlnet = ControlNetModel.from_unet(unet)

optimizer diatur untuk memperbarui parameter ControlNet:

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

Terakhir, dalam loop pelatihan, embedding teks kondisional dan gambar diteruskan ke blok down dan mid dari model ControlNet:

encoder_hidden_states = text_encoder(batch["input_ids"])[0]
controlnet_image = batch["conditioning_pixel_values"].to(dtype=weight_dtype)

down_block_res_samples, mid_block_res_sample = controlnet(
    noisy_latents,
    timesteps,
    encoder_hidden_states=encoder_hidden_states,
    controlnet_cond=controlnet_image,
    return_dict=False,
)

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! 🚀

Panduan ini menggunakan dataset fusing/fill50k, tetapi ingat, Anda dapat membuat dan menggunakan dataset Anda sendiri jika diinginkan (lihat panduan Create a dataset for training).

Setel variabel lingkungan MODEL_NAME ke id model di Hub atau path ke model lokal, dan OUTPUT_DIR ke lokasi tempat Anda ingin menyimpan model.

Unduh gambar-gambar berikut untuk mengondisikan pelatihan Anda dengan:

wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_1.png
wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_2.png

Satu hal lagi sebelum Anda menjalankan skrip! Tergantung pada GPU yang Anda miliki, Anda mungkin perlu mengaktifkan optimasi tertentu untuk melatih ControlNet. Konfigurasi bawaan dalam skrip ini membutuhkan sekitar 38GB vRAM. Jika Anda melatih di lebih dari satu GPU, tambahkan parameter --multi_gpu ke perintah accelerate launch.

pip install bitsandbytes

accelerate launch train_controlnet.py \
  --gradient_checkpointing \
  --use_8bit_adam \

accelerate launch train_controlnet.py \
  --use_8bit_adam \
  --gradient_checkpointing \
  --enable_xformers_memory_efficient_attention \
  --set_grads_to_none \

accelerate config

compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 4
  offload_optimizer_device: cpu
  offload_param_device: cpu
  zero3_init_flag: false
  zero_stage: 2
distributed_type: DEEPSPEED

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

accelerate launch train_controlnet.py \
 --pretrained_model_name_or_path=$MODEL_DIR \
 --output_dir=$OUTPUT_DIR \
 --dataset_name=fusing/fill50k \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./conditioning_image_1.png" "./conditioning_image_2.png" \
 --validation_prompt "red circle with blue background" "cyan circle with brown floral background" \
 --train_batch_size=1 \
 --gradient_accumulation_steps=4 \
 --push_to_hub

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

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import torch

controlnet = ControlNetModel.from_pretrained("path/to/controlnet", dtype=torch.float16)
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
    "path/to/base/model", controlnet=controlnet, dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

control_image = load_image("./conditioning_image_1.png")
prompt = "pale golden rod circle with old lace background"

generator = torch.manual_seed(0)
image = pipeline(prompt, num_inference_steps=20, generator=generator, image=control_image).images[0]
image.save("./output.png")

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_controlnet_sdxl.py untuk melatih adapter ControlNet untuk model SDXL.

Skrip pelatihan SDXL dibahas lebih detail dalam panduan pelatihan SDXL.

Langkah selanjutnya

Selamat atas pelatihan ControlNet Anda sendiri! Untuk mempelajari lebih lanjut tentang cara menggunakan model baru Anda, panduan berikut mungkin dapat membantu:

• Pelajari cara menggunakan ControlNet untuk inferensi pada berbagai tugas.