Skip to content

Kandinsky 2.2

Kandinsky 2.2 是一个多语言的文本到图像模型,能够生成更逼真的图像。该模型包括一个图像先验模型,用于从文本提示中创建图像嵌入,以及一个解码器模型,根据先验模型的嵌入生成图像。这就是为什么在 Diffusers 中你会找到两个单独的脚本用于 Kandinsky 2.2,一个用于训练先验模型,另一个用于训练解码器模型。你可以分别训练这两个模型,但为了获得最佳效果,你应该同时训练先验模型和解码器模型。

根据你的 GPU,你可能需要启用 gradient_checkpointing(⚠️ 先验模型不支持!)、mixed_precisiongradient_accumulation_steps 以帮助将模型装入内存并加速训练。你还可以通过启用 xFormers(版本 v0.0.16 在某些 GPU 上训练时会失败,因此你可能需要安装开发版本)的内存高效注意力机制来进一步减少内存使用。

本指南将探讨 train_text_to_image_prior.pytrain_text_to_image_decoder.py 脚本,帮助你更熟悉这些脚本,并了解如何根据自己的用例进行调整。

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

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

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

bash
cd examples/kandinsky2_2/text_to_image
pip install -r requirements.txt

初始化一个 🤗 Accelerate 环境:

bash
accelerate config

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

bash
accelerate config default

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

py
from accelerate.utils import write_basic_config

write_basic_config()

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

脚本参数

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

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

bash
accelerate launch train_text_to_image_prior.py \
  --mixed_precision="fp16"

大多数参数与 文本到图像 训练指南中的参数相同,因此我们直接进入 Kandinsky 训练脚本的介绍!

Min-SNR 权重

Min-SNR 权重策略可以通过重新平衡损失来加速收敛,从而帮助训练。训练脚本支持预测 epsilon(噪声)或 v_prediction,但 Min-SNR 与这两种预测类型都兼容。此权重策略仅由 PyTorch 支持,Flax 训练脚本中不可用。

添加 --snr_gamma 参数并将其设置为推荐值 5.0:

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

训练脚本

训练脚本也类似于 文本到图像 训练指南,但已修改以支持训练先验模型和解码器模型。本指南重点关注 Kandinsky 2.2 训练脚本中独有的代码。

启动脚本

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

你将在 Naruto BLIP captions 数据集上进行训练,以生成你自己的《火影忍者》角色,但你也可以通过遵循 创建训练数据集 指南来创建和训练你自己的数据集。将环境变量 DATASET_NAME 设置为 Hub 上的数据集名称,或者如果你使用自己的文件进行训练,将环境变量 TRAIN_DIR 设置为你的数据集路径。

如果你使用多于一个 GPU 进行训练,请在 accelerate launch 命令中添加 --multi_gpu 参数。

一旦训练完成,你就可以使用新训练的模型进行推理!

下一步

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

  • 阅读 Kandinsky 指南,了解如何将其用于各种不同的任务(文本到图像、图像到图像、修复、插值),以及如何将其与 ControlNet 结合使用。
  • 查看 DreamBoothLoRA 训练指南,了解如何使用少量示例图像训练个性化的 Kandinsky 模型。这两种训练技术甚至可以结合使用!