Diffusers 文档

文本到图像

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

文本到图像

文本到图像脚本是实验性的,很容易过拟合并遇到灾难性遗忘等问题。尝试探索不同的超参数,以在您的数据集上获得最佳结果。

像 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 .

然后导航到包含训练脚本的示例文件夹并安装您正在使用的脚本所需的依赖项

PyTorch
Flax
cd examples/text_to_image
pip install -r requirements.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() 函数中找到。此函数为每个参数提供了默认值,例如训练批处理大小和学习率,但您也可以在训练命令中设置自己的值。

例如,要使用 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。如果您想了解有关训练循环如何工作的更多信息,请查看了解管道、模型和调度器教程,其中分解了去噪过程的基本模式。

启动脚本

完成所有更改或对默认配置满意后,您就可以启动训练脚本了!🚀

PyTorch
Flax

让我们在 Naruto BLIP captions 数据集上训练,以生成您自己的火影忍者角色。将环境变量 `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

训练完成后,您可以使用新训练的模型进行推理

PyTorch
Flax
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 权重进行推理。
  • 文本到图像任务指南中,了解更多关于指导比例等某些参数或提示加权等技术如何帮助您控制推理。
< > 在 GitHub 上更新