Hugging Face's logo

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 Quick tour 以了解更多信息。

初始化一个 🤗 Accelerate 环境

accelerate config

要设置一个默认的 🤗 Accelerate 环境而无需选择任何配置

accelerate config default

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

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：在时间步长中的哪个位置（较早 vs. 较晚）应用偏差，这可以鼓励模型学习低频或高频细节
  • --timestep_bias_multiplier：应用于时间步长的偏差权重
  • --timestep_bias_begin：开始应用偏差的时间步长
  • --timestep_bias_end：结束应用偏差的时间步长
  • --timestep_bias_portion：将偏差应用于的时间步长比例

Min-SNR 权重

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 训练脚本独有的代码。

它首先创建函数来 标记化 prompts 以计算 prompt embeddings，并使用 VAE 计算图像 embeddings。接下来，您将获得一个函数，该函数根据时间步长数和要应用的时间步长偏差策略来 生成时间步长权重。

在 main() 函数中，除了加载 tokenizer 之外，脚本还加载了第二个 tokenizer 和文本编码器，因为 SDXL 架构使用了两个 tokenizer 和文本编码器

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

Prompt 和图像 embeddings 首先被计算并保存在内存中，对于较小的数据集来说，这通常不是问题，但对于较大的数据集，这可能会导致内存问题。如果是这种情况，您应该将预计算的 embeddings 单独保存到磁盘，并在训练过程中将它们加载到内存中（有关此主题的更多讨论，请参见此 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,
)

在计算 embeddings 后，文本编码器、VAE 和 tokenizer 将被删除以释放一些内存

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)

如果您想了解有关训练循环工作原理的更多信息，请查看 理解 pipelines、模型和 schedulers 教程，该教程分解了去噪过程的基本模式。

启动脚本

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

让我们在 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 指南，了解如何将其用于各种不同的任务（文本到图像、图像到图像、图像修复）、如何使用其 refiner 模型以及不同类型的微条件控制。
  • 查看 DreamBooth 和 LoRA 训练指南，了解如何仅使用少量示例图像来训练个性化的 SDXL 模型。这两种训练技术甚至可以结合使用！
< > 更新 在 GitHub 上

←文本到图像 Kandinsky 2.2→