Diffusers 文件

Stable Diffusion XL

Hugging Face's logo
加入 Hugging Face 社群

並獲得增強的文件體驗

開始使用

Stable Diffusion XL

此指令碼是實驗性的,很容易過擬合併遇到災難性遺忘等問題。嘗試探索不同的超引數以在您的資料集上獲得最佳結果。

Stable Diffusion XL (SDXL) 是 Stable Diffusion 模型的一個更大更強大的迭代,能夠生成更高解析度的影像。

SDXL 的 UNet 大 3 倍,並且模型向架構中添加了第二個文字編碼器。根據您可用的硬體,這可能非常計算密集,並且可能無法在像 Tesla T4 這樣的消費級 GPU 上執行。為了幫助將這個更大的模型裝入記憶體並加快訓練速度,請嘗試啟用 `gradient_checkpointing`、`mixed_precision` 和 `gradient_accumulation_steps`。您可以透過使用 xFormers 啟用記憶體高效注意力並使用 bitsandbytes 的 8 位最佳化器來進一步減少記憶體使用。

本指南將探討 train_text_to_image_sdxl.py 訓練指令碼,以幫助您更熟悉它,以及如何根據您自己的用例對其進行調整。

在執行指令碼之前,請確保從原始碼安裝庫

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

然後導航到包含訓練指令碼的示例資料夾並安裝您正在使用的指令碼所需的依賴項

cd examples/text_to_image
pip install -r requirements_sdxl.txt

🤗 Accelerate 是一個幫助您在多個 GPU/TPU 上或使用混合精度進行訓練的庫。它將根據您的硬體和環境自動配置您的訓練設定。請檢視 🤗 Accelerate 快速入門 以瞭解更多資訊。

初始化 🤗 Accelerate 環境

accelerate config

要設定預設的 🤗 Accelerate 環境而不選擇任何配置

accelerate config default

或者如果您的環境不支援互動式 shell(例如筆記本),您可以使用

from accelerate.utils import write_basic_config

write_basic_config()

最後,如果您想在自己的資料集上訓練模型,請檢視 建立訓練資料集 指南,瞭解如何建立與訓練指令碼相容的資料集。

指令碼引數

以下部分重點介紹了訓練指令碼中對於理解如何修改它很重要的部分,但它沒有詳細介紹指令碼的每個方面。如果您有興趣瞭解更多資訊,請隨時閱讀 指令碼,如果您有任何問題或疑慮,請告訴我們。

訓練指令碼提供了許多引數來幫助您自定義訓練執行。所有引數及其描述都可以在 parse_args() 函式中找到。此函式為每個引數(例如訓練批次大小和學習率)提供了預設值,但您也可以在訓練命令中設定自己的值(如果需要)。

例如,要使用 bf16 格式的混合精度加速訓練,請將 `--mixed_precision` 引數新增到訓練命令中

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

大多數引數與 文字到影像 訓練指南中的引數相同,因此本指南將重點介紹與訓練 SDXL 相關的引數。

  • --pretrained_vae_model_name_or_path:預訓練 VAE 的路徑;SDXL VAE 已知存在數值不穩定性問題,因此此引數允許您指定更好的 VAE
  • --proportion_empty_prompts:將影像提示替換為空字串的比例
  • --timestep_bias_strategy:在時間步長的哪個位置(更早或更晚)應用偏差,這可以促使模型學習低頻或高頻細節
  • --timestep_bias_multiplier:應用於時間步長的偏差權重
  • --timestep_bias_begin:開始應用偏差的時間步長
  • --timestep_bias_end:結束應用偏差的時間步長
  • --timestep_bias_portion:應用偏差的時間步長比例

最小信噪比加權

最小信噪比 (Min-SNR) 加權策略可以透過重新平衡損失來實現更快收斂來幫助訓練。訓練指令碼支援預測 `epsilon`(噪聲)或 `v_prediction`,但 Min-SNR 相容這兩種預測型別。此加權策略僅受 PyTorch 支援,在 Flax 訓練指令碼中不可用。

新增 --snr_gamma 引數並將其設定為推薦值 5.0

accelerate launch train_text_to_image_sdxl.py \
  --snr_gamma=5.0

訓練指令碼

該訓練指令碼也類似於文字到影像訓練指南,但已針對 SDXL 訓練進行了修改。本指南將重點介紹 SDXL 訓練指令碼獨有的程式碼。

它首先建立函式以分詞提示以計算提示嵌入,並使用VAE計算影像嵌入。接下來,您將建立一個函式,根據時間步數和要應用的時間步偏差策略來生成時間步權重

main()函式中,除了載入一個分詞器,指令碼還載入了第二個分詞器和文字編碼器,因為 SDXL 架構使用了兩個這樣的元件。

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

提示和影像嵌入首先計算並儲存在記憶體中,這對於較小的資料集來說通常不是問題,但對於較大的資料集可能會導致記憶體問題。如果出現這種情況,您應該將預計算的嵌入單獨儲存到磁碟,並在訓練過程中將它們載入到記憶體中(有關此主題的更多討論,請參閱此 PR)。

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

計算嵌入後,文字編碼器、VAE 和分詞器將被刪除以釋放一些記憶體

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

最後,訓練迴圈負責其餘部分。如果您選擇應用時間步長偏差策略,您會看到時間步長權重被計算並作為噪聲新增

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)

如果您想了解更多關於訓練迴圈如何工作的資訊,請檢視 理解管道、模型和排程器 教程,它分解了去噪過程的基本模式。

啟動指令碼

完成所有更改或對預設配置滿意後,您就可以啟動訓練指令碼了!🚀

讓我們在 Naruto BLIP captions 資料集上進行訓練,以生成您自己的火影忍者角色。將環境變數 `MODEL_NAME` 和 `DATASET_NAME` 設定為模型和資料集(來自 Hub 或本地路徑)。您還應該使用 `VAE_NAME` 指定一個不同於 SDXL VAE 的 VAE(來自 Hub 或本地路徑),以避免數值不穩定性。

要使用 Weights & Biases 監控訓練進度,請將 `--report_to=wandb` 引數新增到訓練命令中。您還需要將 `--validation_prompt` 和 `--validation_epochs` 新增到訓練命令中以跟蹤結果。這對於除錯模型和檢視中間結果非常有用。

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

訓練完成後,您可以使用新訓練的 SDXL 模型進行推理!

PyTorch
PyTorch XLA
from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("path/to/your/model", torch_dtype=torch.float16).to("cuda")

prompt = "A naruto with green eyes and red legs."
image = pipeline(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
image.save("naruto.png")

後續步驟

恭喜您訓練了一個 SDXL 模型!要了解如何使用您的新模型,以下指南可能會有所幫助

  • 閱讀 Stable Diffusion XL 指南,瞭解如何將其用於各種不同任務(文字到影像、影像到影像、影像修復),如何使用其精煉器模型以及不同型別的微條件。
  • 檢視 DreamBoothLoRA 訓練指南,瞭解如何僅使用少量示例影像來訓練個性化的 SDXL 模型。這兩種訓練技術甚至可以結合使用!
< > 在 GitHub 上更新

© . This site is unofficial and not affiliated with Hugging Face, Inc.