文本到图像

诸如 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 脚本首先 加载 scheduler 和分词器。如果您愿意，可以在此处选择使用其他 scheduler

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。如果您想了解有关训练循环如何工作的更多信息，请查看 理解 pipelines、models 和 schedulers 教程，该教程分解了去噪过程的基本模式。

启动脚本

一旦您完成了所有更改，或者对默认配置感到满意，就可以启动训练脚本了！🚀

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 权重 以进行推理。
• 在 文本到图像 任务指南中，详细了解诸如 guidance scale 之类的某些参数或诸如 prompt weighting 之类的技术如何帮助您控制推理。