Stable Diffusion XL

Script này mang tính thử nghiệm, và rất dễ bị overfit cũng như gặp các vấn đề như catastrophic forgetting. Hãy thử khám phá các hyperparameter khác nhau để đạt được kết quả tốt nhất trên dataset của bạn.

Stable Diffusion XL (SDXL) là phiên bản lớn hơn và mạnh hơn của mô hình Stable Diffusion, có khả năng tạo ra các hình ảnh có độ phân giải cao hơn.

UNet của SDXL lớn gấp 3 lần và mô hình bổ sung một text encoder thứ hai vào kiến trúc. Tùy thuộc vào phần cứng bạn có, điều này có thể đòi hỏi rất nhiều tài nguyên tính toán và có thể không chạy được trên GPU tiêu dùng như Tesla T4. Để giúp đưa mô hình lớn hơn này vào bộ nhớ và tăng tốc quá trình huấn luyện, hãy thử bật gradient_checkpointing, mixed_precision và gradient_accumulation_steps. Bạn có thể giảm mức sử dụng bộ nhớ xuống thấp hơn nữa bằng cách bật attention tiết kiệm bộ nhớ với xFormers và sử dụng optimizer 8-bit của bitsandbytes.

Hướng dẫn này sẽ khám phá script huấn luyện train_text_to_image_sdxl.py để giúp bạn làm quen hơn với nó, cũng như cách bạn có thể điều chỉnh nó cho trường hợp sử dụng của riêng mình.

Trước khi chạy script, hãy đảm bảo bạn cài đặt thư viện từ mã nguồn:

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

Sau đó, điều hướng đến thư mục ví dụ chứa script huấn luyện và cài đặt các phụ thuộc cần thiết cho script bạn đang sử dụng:

cd examples/text_to_image
pip install -r requirements_sdxl.txt

🤗 Accelerate là một thư viện giúp bạn huấn luyện trên nhiều GPU/TPU hoặc với độ chính xác hỗn hợp. Nó sẽ tự động cấu hình thiết lập huấn luyện của bạn dựa trên phần cứng và môi trường của bạn. Hãy xem 🤗 Accelerate Quick tour để tìm hiểu thêm.

Khởi tạo một môi trường 🤗 Accelerate:

accelerate config

Để thiết lập môi trường 🤗 Accelerate mặc định mà không cần chọn bất kỳ cấu hình nào:

accelerate config default

Hoặc nếu môi trường của bạn không hỗ trợ shell tương tác, chẳng hạn như notebook, bạn có thể sử dụng:

from accelerate.utils import write_basic_config

write_basic_config()

Cuối cùng, nếu bạn muốn huấn luyện một mô hình trên bộ dữ liệu của riêng mình, hãy xem hướng dẫn Tạo bộ dữ liệu để huấn luyện để tìm hiểu cách tạo một bộ dữ liệu tương thích với script huấn luyện.

Tham số của script

Các phần tiếp theo làm nổi bật những phần quan trọng trong script huấn luyện để hiểu cách chỉnh sửa nó, nhưng không đi sâu vào mọi khía cạnh của script. Nếu bạn muốn tìm hiểu thêm, hãy tự do đọc qua script và cho chúng tôi biết nếu bạn có bất kỳ câu hỏi hoặc lo ngại nào.

Tệp script huấn luyện cung cấp nhiều tham số để giúp bạn tùy chỉnh phiên huấn luyện của mình. Tất cả các tham số và mô tả của chúng được tìm thấy trong hàm parse_args(). Hàm này cung cấp các giá trị mặc định cho từng tham số, chẳng hạn như kích thước lô huấn luyện và tốc độ học, nhưng bạn cũng có thể tự đặt giá trị của riêng mình trong lệnh huấn luyện nếu muốn.

Ví dụ, để tăng tốc độ huấn luyện với độ chính xác hỗn hợp sử dụng định dạng bf16, hãy thêm tham số --mixed_precision vào lệnh huấn luyện:

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

Hầu hết các tham số đều giống với các tham số trong hướng dẫn huấn luyện Text-to-image, vì vậy bạn sẽ tập trung vào các tham số liên quan đến việc huấn luyện SDXL trong hướng dẫn này.

--pretrained_vae_model_name_or_path: đường dẫn đến một VAE được huấn luyện trước; VAE của SDXL được biết là gặp phải sự bất ổn định về mặt số học, vì vậy tham số này cho phép bạn chỉ định một VAE tốt hơn
--proportion_empty_prompts: tỷ lệ các prompt hình ảnh sẽ được thay thế bằng chuỗi rỗng
--timestep_bias_strategy: vị trí (sớm hơn hay muộn hơn) trong timestep để áp dụng bias, có thể khuyến khích mô hình học chi tiết tần số thấp hoặc tần số cao
--timestep_bias_multiplier: trọng số của độ chệch cần áp dụng cho timestep
--timestep_bias_begin: bước thời gian bắt đầu áp dụng độ lệch
--timestep_bias_end: bước thời gian để kết thúc việc áp dụng độ lệch
--timestep_bias_portion: tỷ lệ các bước thời gian (timesteps) cần áp dụng độ lệch

Trọng số Min-SNR

Chiến lược trọng số Min-SNR có thể hỗ trợ quá trình huấn luyện bằng cách cân bằng lại loss để đạt được sự hội tụ nhanh hơn. Script huấn luyện hỗ trợ dự đoán epsilon (noise) hoặc v_prediction, nhưng Min-SNR tương thích với cả hai loại dự đoán. Chiến lược trọng số này chỉ được hỗ trợ bởi PyTorch.

Thêm tham số --snr_gamma và đặt nó bằng giá trị khuyến nghị là 5.0:

accelerate launch train_text_to_image_sdxl.py \
  --snr_gamma=5.0

Script huấn luyện

Tệp kịch bản huấn luyện cũng tương tự như hướng dẫn huấn luyện Text-to-image, nhưng đã được chỉnh sửa để hỗ trợ huấn luyện SDXL. Hướng dẫn này sẽ tập trung vào phần mã riêng có trong tệp kịch bản huấn luyện SDXL.

Quá trình bắt đầu bằng việc tạo các hàm để tokenize các prompt nhằm tính toán các embedding của prompt, và để tính toán các embedding của ảnh bằng VAE. Tiếp theo, bạn sẽ tạo một hàm để tạo các trọng số timesteps tùy thuộc vào số lượng timesteps và chiến lược bias của timestep cần áp dụng.

Bên trong hàm main(), ngoài việc nạp một tokenizer, script còn nạp một tokenizer và text encoder thứ hai vì kiến trúc SDXL sử dụng hai bộ cho mỗi loại:

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

prompt và image embeddings được tính toán trước và giữ trong bộ nhớ, điều này thường không phải là vấn đề với bộ dữ liệu nhỏ, nhưng với bộ dữ liệu lớn hơn, nó có thể dẫn đến các vấn đề về bộ nhớ. Nếu gặp trường hợp này, bạn nên lưu các embeddings đã tính toán trước vào đĩa riêng và nạp chúng vào bộ nhớ trong quá trình huấn luyện (xem PR này để biết thêm thảo luận về chủ đề này).

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

Sau khi tính toán các embedding, text encoder, VAE và tokenizer sẽ bị xóa để giải phóng một phần bộ nhớ:

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

Cuối cùng, vòng lặp huấn luyện sẽ xử lý phần còn lại. Nếu bạn chọn áp dụng chiến lược thiên lệch bước thời gian, bạn sẽ thấy trọng số bước thời gian được tính toán và thêm vào như nhiễu:

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)

Nếu bạn muốn tìm hiểu thêm về cách hoạt động của vòng lặp huấn luyện, hãy xem hướng dẫn Hiểu về pipeline, mô hình và trình lập lịch, nơi phân tích chi tiết mẫu cơ bản của quá trình khử nhiễu.

Chạy script

Khi bạn đã thực hiện tất cả các thay đổi hoặc hài lòng với cấu hình mặc định, bạn đã sẵn sàng để chạy script huấn luyện! 🚀

Hãy huấn luyện trên bộ dữ liệu Naruto BLIP captions để tạo ra các nhân vật Naruto của riêng bạn. Đặt các biến môi trường MODEL_NAME và DATASET_NAME thành tên mô hình và tên bộ dữ liệu (từ Hub hoặc đường dẫn cục bộ). Bạn cũng nên chỉ định một VAE khác với VAE của SDXL (từ Hub hoặc đường dẫn cục bộ) bằng VAE_NAME để tránh các bất ổn về mặt số.

Để theo dõi tiến độ huấn luyện với Weights & Biases, hãy thêm tham số --report_to=wandb vào lệnh huấn luyện. Bạn cũng cần thêm --validation_prompt và --validation_epochs vào lệnh huấn luyện để theo dõi kết quả. Điều này thực sự hữu ích để gỡ lỗi mô hình và xem các kết quả trung gian.

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

Sau khi hoàn tất huấn luyện, bạn có thể sử dụng mô hình SDXL vừa được huấn luyện để suy luận!

PyTorch

PyTorch XLA

Các bước tiếp theo

Chúc mừng bạn đã huấn luyện thành công một mô hình SDXL! Để tìm hiểu thêm về cách sử dụng mô hình mới của bạn, các hướng dẫn sau có thể hữu ích:

Đọc hướng dẫn Stable Diffusion XL để tìm hiểu cách sử dụng nó cho nhiều tác vụ khác nhau (text-to-image, image-to-image, inpainting), cách sử dụng mô hình refiner của nó, và các loại micro-conditioning khác nhau.
Xem hướng dẫn huấn luyện DreamBooth và LoRA để tìm hiểu cách huấn luyện một mô hình SDXL cá nhân hóa chỉ với một vài hình ảnh mẫu. Hai kỹ thuật huấn luyện này thậm chí có thể được kết hợp với nhau!

Update on GitHub

Diffusers

Stable Diffusion XL

Tham số của script

Trọng số Min-SNR

Script huấn luyện

Chạy script

Các bước tiếp theo