Skip to content

Stable Diffusion XL

Stable Diffusion XL (SDXL) 是 Stable Diffusion 模型的一个更大、更强大的版本,能够生成更高分辨率的图像。

SDXL 的 UNet 比之前的版本大 3 倍,并且模型在架构中添加了第二个文本编码器。根据你可用的硬件,这可能会非常计算密集,可能无法在像 Tesla T4 这样的消费级 GPU 上运行。为了帮助将这个更大的模型装入内存并加速训练,可以尝试启用 gradient_checkpointingmixed_precisiongradient_accumulation_steps。你还可以通过启用 xFormers 的内存高效注意力机制和使用 bitsandbytes' 的 8 位优化器来进一步减少内存使用。

本指南将探讨 train_text_to_image_sdxl.py 训练脚本,帮助你更熟悉它,并了解如何根据自己的需求进行调整。

在运行脚本之前,请确保从源代码安装库:

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

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

bash
cd examples/text_to_image
pip install -r requirements_sdxl.txt

初始化一个 🤗 Accelerate 环境:

bash
accelerate config

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

bash
accelerate config default

或者,如果你的环境不支持交互式 shell,比如笔记本,你可以使用:

py
from accelerate.utils import write_basic_config

write_basic_config()

最后,如果你想在自己的数据集上训练模型,请参阅创建用于训练的数据集指南,了解如何创建与训练脚本兼容的数据集。

脚本参数

训练脚本提供了许多参数,帮助你自定义训练过程。所有参数及其描述都在 parse_args() 函数中。此函数为每个参数提供了默认值,例如训练批次大小和学习率,但你也可以在训练命令中设置自己的值。

例如,要使用 bf16 格式通过混合精度加速训练,可以在训练命令中添加 --mixed_precision 参数:

bash
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: 在时间步的早期或晚期应用偏差,这可以鼓励模型学习低频或高频细节
  • --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:

bash
accelerate launch train_text_to_image_sdxl.py \
  --snr_gamma=5.0

训练脚本

训练脚本与 文本到图像 训练指南类似,但已修改以支持 SDXL 训练。本指南将重点介绍 SDXL 训练脚本中独有的代码。

脚本首先创建函数来 对提示进行分词,以计算提示嵌入,并使用 VAE 计算图像嵌入。接下来,你将创建一个函数来 生成时间步权重,这取决于时间步的数量和要应用的时间步偏置策略。

main() 函数中,除了加载一个分词器外,脚本还加载了第二个分词器和文本编码器,因为 SDXL 架构使用了两个分词器和两个文本编码器。

py
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 和图像嵌入 首先被计算并保存在内存中,对于较小的数据集这通常不是问题,但对于较大的数据集,这可能会导致内存问题。如果出现这种情况,你应该将预计算的嵌入单独保存到磁盘上,并在训练过程中将它们加载到内存中(参见这个 PR 以获取更多关于这个主题的讨论)。

py
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,
)

在计算完嵌入后,文本编码器、VAE 和分词器会被删除以释放一些内存:

py
del text_encoders, tokenizers, vae
gc.collect()
torch.cuda.empty_cache()

最后,训练循环 会处理剩余的部分。如果你选择应用时间步偏置策略,你会看到时间步权重会被计算并作为噪声添加:

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

如果你想了解更多关于训练循环的工作原理,可以查看 理解管道、模型和调度器 教程,该教程详细介绍了去噪过程的基本模式。

启动脚本

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

让我们使用 Naruto BLIP 标题 数据集来生成你自己的 Naruto 角色。将环境变量 MODEL_NAMEDATASET_NAME 设置为模型和数据集(可以是来自 Hub 的或本地路径)。你还应该使用 VAE_NAME 指定一个不同于 SDXL VAE 的 VAE(可以是来自 Hub 的或本地路径),以避免数值不稳定性。

bash
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 模型进行推理!

下一步

恭喜你训练了一个 SDXL 模型!要了解更多如何使用你的新模型,以下指南可能会对你有所帮助:

  • 阅读 Stable Diffusion XL 指南,了解如何将其用于各种不同的任务(文本到图像、图像到图像、修复),如何使用其精炼模型,以及不同类型的微调条件。
  • 查看 DreamBoothLoRA 训练指南,了解如何使用少量示例图像训练个性化的 SDXL 模型。这两种训练技术甚至可以结合使用!