文本到图像

文本到图像脚本目前仍处于实验阶段，很容易过拟合，并出现灾难性遗忘等问题。请尝试探索不同的超参数，以在你的数据集上获得最佳效果。

Stable Diffusion 等文本到图像模型以文本提示为条件，用于生成图像。

训练模型对硬件的要求可能很高，但如果您启用了 gradient_checkpointing 和 mixed_precision，则可以在单张 24GB 显存的 GPU 上完成训练。如果您使用更大的 batch size 或希望加快训练速度，建议使用显存大于 30GB 的 GPU。您可以通过启用 xFormers 提供的内存高效注意力机制来降低内存占用。

本指南将深入讲解 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 是一个用于在多个 GPU/TPU 上训练或使用混合精度训练的库。它会根据你的硬件和环境自动配置训练设置。请查看 🤗 Accelerate 的快速入门以了解更多信息。

初始化一个 🤗 Accelerate 环境：

accelerate config

要设置一个默认的 🤗 Accelerate 环境，而不选择任何配置：

accelerate config default

或者，如果你的环境不支持交互式 shell（例如 notebook），你可以使用：

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。

添加 --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。如果你想进一步了解训练循环的工作原理，请查看理解流水线、模型和调度器教程，该教程详细拆解了去噪过程的基本模式。

运行脚本

完成所有更改或对默认配置满意后，即可启动训练脚本！🚀

让我们使用 Naruto BLIP captions 数据集进行训练，以生成你自己的 Naruto 角色。将环境变量 MODEL_NAME 和 dataset_name 设置为模型和数据集（来自 Hub 或本地路径）。如果你使用多于一块 GPU 进行训练，请在 accelerate launch 命令中添加 --multi_gpu 参数。

要在本地数据集上进行训练，请将 TRAIN_DIR 和 OUTPUT_DIR 环境变量分别设置为数据集的路径和模型保存的路径。

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", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"

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

后续步骤

恭喜您训练了自己的文本到图像模型！如需了解如何使用您的新模型，以下指南可能对您有所帮助：

如果您使用 LoRA 训练了模型，请学习如何加载 LoRA 权重以进行推理。
有关 guidance scale 等特定参数或 prompt weighting 等技术如何帮助你控制推理，请参阅 Text-to-image 任务指南。

Update on GitHub

Diffusers

文本到图像

脚本参数

Min-SNR 加权

训练脚本

运行脚本

后续步骤