文字到影像

像 Stable Diffusion 這樣的文字到影像模型經過條件訓練，可以根據文字提示生成影像。

訓練模型可能會對您的硬體造成負擔，但如果啟用 `gradient_checkpointing` 和 `mixed_precision`，則可以在單個 24GB GPU 上訓練模型。如果您使用更大的批處理大小進行訓練或希望更快地訓練，最好使用記憶體超過 30GB 的 GPU。您可以透過使用 xFormers 啟用記憶體高效的注意力來減少記憶體佔用。JAX/Flax 訓練也支援在 TPU 和 GPU 上進行高效訓練，但它不支援梯度檢查點、梯度累積或 xFormers。建議使用至少 30GB 記憶體的 GPU 或 TPU v3 來使用 Flax 進行訓練。

本指南將探討 train_text_to_image.py 訓練指令碼，以幫助您熟悉它，以及如何將其用於您自己的用例。

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

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

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

cd examples/text_to_image
pip install -r requirements.txt

初始化 🤗 Accelerate 環境

accelerate config

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

accelerate config default

或者如果您的環境不支援互動式 shell（例如筆記本），您可以使用

from accelerate.utils import write_basic_config

write_basic_config()

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

指令碼引數

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

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

accelerate launch train_text_to_image.py \
  --mixed_precision="fp16"

一些基本和重要的引數包括

• --pretrained_model_name_or_path：Hub 上的模型名稱或預訓練模型的本地路徑
• --dataset_name：Hub 上的資料集名稱或要訓練的資料集的本地路徑
• --image_column：資料集中用於訓練的影像列的名稱
• --caption_column：資料集中用於訓練的文字列的名稱
• --output_dir：訓練好的模型的儲存位置
• --push_to_hub：是否將訓練好的模型推送到 Hub
• --checkpointing_steps：模型訓練時儲存檢查點的頻率；這對於訓練因某種原因中斷的情況很有用，您可以透過在訓練命令中新增 --resume_from_checkpoint 來從該檢查點繼續訓練

Min-SNR 加權

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

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

accelerate launch train_text_to_image.py \
  --snr_gamma=5.0

您可以在此 Weights and Biases 報告中比較不同 `snr_gamma` 值的損失曲面。對於較小的資料集，Min-SNR 的效果可能不如在較大資料集上那麼明顯。

訓練指令碼

資料集預處理程式碼和訓練迴圈位於 main() 函式中。如果您需要調整訓練指令碼，這裡是您需要進行更改的地方。

train_text_to_image 指令碼首先載入排程器和分詞器。如果需要，您可以在此處選擇使用不同的排程器

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
)

然後指令碼載入 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())

接下來，需要對資料集的文字和影像列進行預處理。`tokenize_captions` 函式負責對輸入進行分詞，`train_transforms` 函式指定了應用於影像的變換型別。這兩個函式都打包在 `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

最後，訓練迴圈處理其他所有事情。它將影像編碼到潛在空間，向潛在空間新增噪聲，計算要條件化的文字嵌入，更新模型引數，並將模型儲存並推送到 Hub。如果您想了解有關訓練迴圈如何工作的更多資訊，請檢視瞭解管道、模型和排程器教程，其中分解了去噪過程的基本模式。

啟動指令碼

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

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

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

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("path/to/saved_model", torch_dtype=torch.float16, use_safetensors=True).to("cuda")

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

下一步

恭喜您訓練出自己的文字到影像模型！要了解如何使用新模型，以下指南可能會有所幫助

• 如果您使用 LoRA 訓練了模型，請了解如何載入 LoRA 權重進行推理。
• 在文字到影像任務指南中，瞭解更多關於指導比例等某些引數或提示加權等技術如何幫助您控制推理。