Teks-ke-gambar

Skrip teks-ke-gambar bersifat eksperimental, dan sangat mudah mengalami overfit serta menghadapi masalah seperti catastrophic forgetting. Coba jelajahi berbagai hiperparameter untuk mendapatkan hasil terbaik pada dataset Anda.

Model text-to-image seperti Stable Diffusion dikondisikan untuk menghasilkan gambar berdasarkan prompt teks.

Melatih model dapat membebani perangkat keras Anda, tetapi jika Anda mengaktifkan gradient_checkpointing dan mixed_precision, Anda tetap dapat melatih model pada satu GPU 24GB. Jika Anda melatih dengan ukuran batch yang lebih besar atau ingin melatih lebih cepat, sebaiknya gunakan GPU dengan memori lebih dari 30GB. Anda dapat mengurangi jejak memori dengan mengaktifkan perhatian hemat memori menggunakan xFormers.

Panduan ini akan membahas skrip pelatihan train_text_to_image.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/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.

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_text_to_image.py \
  --mixed_precision="fp16"

Beberapa parameter dasar dan penting meliputi:

• --pretrained_model_name_or_path: nama model di Hub atau jalur lokal ke model pretrained
• --dataset_name: nama dataset di Hub atau jalur lokal ke dataset yang akan digunakan untuk pelatihan
• --image_column: nama kolom gambar dalam dataset untuk pelatihan
• --caption_column: nama kolom teks dalam dataset untuk dilatih
• --output_dir: tempat menyimpan model yang telah dilatih
• --push_to_hub: apakah model yang dilatih akan didorong ke Hub
• --checkpointing_steps: frekuensi penyimpanan checkpoint saat model dilatih; ini berguna jika pelatihan terputus karena alasan tertentu, Anda dapat melanjutkan pelatihan dari checkpoint tersebut dengan menambahkan --resume_from_checkpoint ke perintah pelatihan 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_text_to_image.py \
  --snr_gamma=5.0

Anda dapat membandingkan permukaan loss untuk nilai snr_gamma yang berbeda dalam laporan Weights and Biases ini. Untuk dataset yang lebih kecil, efek Min-SNR mungkin tidak seobvious dibandingkan dengan dataset yang lebih besar.

Skrip pelatihan

Kode pra-pemrosesan dataset dan loop pelatihan dapat ditemukan di fungsi main(). Jika Anda perlu menyesuaikan skrip pelatihan, di sinilah Anda perlu melakukan perubahan.

Skrip train_text_to_image dimulai dengan memuat scheduler dan tokenizer. Anda dapat memilih untuk menggunakan scheduler yang berbeda di sini jika diinginkan:

noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
tokenizer = CLIPTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
)

Lalu skrip memuat model UNet:

load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
model.register_to_config(**load_model.config)

model.load_state_dict(load_model.state_dict())

Selanjutnya, kolom teks dan gambar dari dataset perlu dipreproses. Fungsi tokenize_captions menangani tokenisasi input, dan fungsi train_transforms menentukan jenis transformasi yang diterapkan pada gambar. Kedua fungsi ini dikemas ke dalam preprocess_train:

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["input_ids"] = tokenize_captions(examples)
    return examples

Terakhir, training loop menangani semua hal lainnya. Loop ini meng-encode gambar ke ruang laten, menambahkan noise ke latents, menghitung text embeddings untuk dikondisikan, memperbarui parameter model, serta menyimpan dan mendorong model ke Hub. Jika Anda ingin mempelajari lebih lanjut tentang cara kerja training loop, lihat tutorial Understanding pipelines, models and schedulers 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). Jika Anda melatih pada lebih dari satu GPU, tambahkan parameter --multi_gpu ke perintah accelerate launch.

Untuk melatih pada dataset lokal, atur variabel lingkungan TRAIN_DIR dan OUTPUT_DIR ke path dataset dan lokasi penyimpanan model.

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export dataset_name="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$dataset_name \
  --use_ema \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --enable_xformers_memory_efficient_attention \
  --lr_scheduler="constant" --lr_warmup_steps=0 \
  --output_dir="sd-naruto-model" \
  --push_to_hub

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

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("path/to/saved_model", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"

image = pipeline(prompt="yoda").images[0]
image.save("yoda-naruto.png")

Langkah selanjutnya

Selamat karena Anda telah melatih model teks-ke-gambar Anda sendiri! Untuk mempelajari lebih lanjut tentang cara menggunakan model baru Anda, panduan berikut mungkin dapat membantu:

• Pelajari cara memuat bobot LoRA untuk inferensi jika Anda melatih model Anda dengan LoRA.
• Pelajari lebih lanjut tentang bagaimana parameter tertentu seperti guidance scale atau teknik seperti prompt weighting dapat membantu Anda mengontrol inferensi dalam panduan tugas Text-to-image.