Hugging Face 的标志

Diffusers 文档

Kandinsky 2.2

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

Kandinsky 2.2

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

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

根据您的 GPU，您可能需要启用 gradient_checkpointing（⚠️ 先验模型不支持！）、mixed_precision 和 gradient_accumulation_steps，以帮助将模型放入内存并加速训练。您还可以通过启用具有 xFormers 的内存高效注意力来进一步减少内存使用（v0.0.16 版本在某些 GPU 上训练会失败，因此您可能需要安装开发版本）。

本指南探讨了 train_text_to_image_prior.py 和 train_text_to_image_decoder.py 脚本，以帮助您更熟悉它们，以及如何将它们用于您自己的用例。

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

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

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

cd examples/kandinsky2_2/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_prior.py \
  --mixed_precision="fp16"

大多数参数与 文本到图像 训练指南中的参数相同，因此让我们直接进入 Kandinsky 训练脚本的演练！

Min-SNR 权重

Min-SNR 权重策略可以通过重新平衡损失以实现更快的收敛来帮助训练。训练脚本支持预测 epsilon（噪声）或 v_prediction，但 Min-SNR 与两种预测类型都兼容。此权重策略仅受 PyTorch 支持，在 Flax 训练脚本中不可用。

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

accelerate launch train_text_to_image_prior.py \
  --snr_gamma=5.0

训练脚本

训练脚本也类似于 文本到图像 训练指南，但已修改为支持训练先验模型和解码器模型。本指南重点介绍 Kandinsky 2.2 训练脚本特有的代码。

先验模型
解码器模型

main() 函数包含用于准备数据集和训练模型的代码。

您会立即注意到的主要区别之一是，除了 scheduler 和 tokenizer 之外，训练脚本还加载了一个 CLIPImageProcessor - 用于预处理图像，以及一个 CLIPVisionModelWithProjection 模型 - 用于编码图像

noise_scheduler = DDPMScheduler(beta_schedule="squaredcos_cap_v2", prediction_type="sample")
image_processor = CLIPImageProcessor.from_pretrained(
    args.pretrained_prior_model_name_or_path, subfolder="image_processor"
)
tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="tokenizer")

with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="image_encoder", torch_dtype=weight_dtype
    ).eval()
    text_encoder = CLIPTextModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="text_encoder", torch_dtype=weight_dtype
    ).eval()

Kandinsky 使用 PriorTransformer 生成图像嵌入，因此您需要设置优化器来学习先验模型的参数。

prior = PriorTransformer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="prior")
prior.train()
optimizer = optimizer_cls(
    prior.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

接下来，输入标题被标记化，图像由 预处理，通过 CLIPImageProcessor

def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[image_column]]
    examples["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
    examples["text_input_ids"], examples["text_mask"] = tokenize_captions(examples)
    return examples

最后，训练循环 将输入图像转换为潜在空间，向图像嵌入添加噪声，并进行预测

model_pred = prior(
    noisy_latents,
    timestep=timesteps,
    proj_embedding=prompt_embeds,
    encoder_hidden_states=text_encoder_hidden_states,
    attention_mask=text_mask,
).predicted_image_embedding

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

启动脚本

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

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

如果您在多个 GPU 上进行训练，请将 --multi_gpu 参数添加到 accelerate launch 命令中。

要使用 Weights & Biases 监控训练进度，请将 --report_to=wandb 参数添加到训练命令中。您还需要将 --validation_prompt 添加到训练命令中，以跟踪结果。这对于调试模型和查看中间结果非常有用。

先验模型
解码器模型
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_prior.py \
  --dataset_name=$DATASET_NAME \
  --resolution=768 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --checkpoints_total_limit=3 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --validation_prompts="A robot naruto, 4k photo" \
  --report_to="wandb" \
  --push_to_hub \
  --output_dir="kandi2-prior-naruto-model"

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

先验模型
解码器模型
from diffusers import AutoPipelineForText2Image, DiffusionPipeline
import torch

prior_pipeline = DiffusionPipeline.from_pretrained(output_dir, torch_dtype=torch.float16)
prior_components = {"prior_" + k: v for k,v in prior_pipeline.components.items()}
pipeline = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", **prior_components, torch_dtype=torch.float16)

pipe.enable_model_cpu_offload()
prompt="A robot naruto, 4k photo"
image = pipeline(prompt=prompt, negative_prompt=negative_prompt).images[0]

随意将 kandinsky-community/kandinsky-2-2-decoder 替换为您自己训练的解码器检查点！

下一步

恭喜您训练了 Kandinsky 2.2 模型！要了解有关如何使用新模型的更多信息，以下指南可能会有所帮助

  • 阅读 Kandinsky 指南，了解如何将其用于各种不同的任务（文本到图像、图像到图像、图像修复、插值），以及如何将其与 ControlNet 结合使用。
  • 查看 DreamBooth 和 LoRA 训练指南，了解如何仅使用少量示例图像来训练个性化的 Kandinsky 模型。这两种训练技术甚至可以结合使用！
< > 更新 在 GitHub 上

←Stable Diffusion XL Wuerstchen→