Hugging Face's logo

Diffusers documentation

Kandinsky 2.2

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Kandinsky 2.2

此脚本仍处于实验阶段，很容易过拟合，并出现灾难性遗忘等问题。请尝试探索不同的超参数，以在您的数据集上获得最佳效果。

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

根据您的 GPU 型号，您可能需要启用 gradient_checkpointing（⚠️ 不支持旧版模型！）、mixed_precision 和 gradient_accumulation_steps，以帮助将模型装入内存并加快训练速度。您还可以通过启用 xFormers 的内存高效注意力机制来进一步降低内存占用（v0.0.16 版本在某些 GPU 上训练时会失败，因此您可能需要改为安装开发版本）。

This guide explores the train_text_to_image_prior.py and the train_text_to_image_decoder.py scripts to help you become more familiar with it, and how you can adapt it for your own use-case.

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

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。

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

accelerate launch train_text_to_image_prior.py \
  --snr_gamma=5.0

训练脚本

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

先验模型
解码器模型

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

您会立即注意到的主要区别之一是，训练脚本除了加载调度器和分词器外，还会加载一个 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", dtype=weight_dtype
    ).eval()
    text_encoder = CLIPTextModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="text_encoder", 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

如果想进一步了解训练循环的工作原理，请查看理解流水线、模型和调度器教程，其中详细拆解了去噪过程的基本模式。

运行脚本

完成所有更改或对默认配置满意后，即可启动训练脚本！🚀

您将在 Naruto BLIP captions 数据集上进行训练，以生成您自己的 Naruto 角色，但您也可以按照 Create a dataset for training 指南创建并训练自己的数据集。将环境变量 DATASET_NAME 设置为 Hub 上数据集的名称，或者如果您在自己的文件上进行训练，则将环境变量 TRAIN_DIR 设置为指向您数据集的路径。

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

若要使用 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, 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, 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 模型。这两种训练技术甚至可以组合使用！
Update on GitHub

←Stable Diffusion XL ControlNet→