# Kandinsky 2.2

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

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

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

This guide explores the [train_text_to_image_prior.py](https://github.com/huggingface/diffusers/blob/main/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py) and the [train_text_to_image_decoder.py](https://github.com/huggingface/diffusers/blob/main/examples/kandinsky2_2/text_to_image/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.

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

```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
```

> [!TIP]
> 🤗 Accelerate 是一个用于在多个 GPU/TPU 上训练或使用混合精度训练的库。它会根据你的硬件和环境自动配置训练设置。请查看 🤗 Accelerate 的 [快速入门](https://huggingface.co/docs/accelerate/quicktour) 以了解更多信息。

初始化一个 🤗 Accelerate 环境：

```bash
accelerate config
```

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

```bash
accelerate config default
```

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

```py
from accelerate.utils import write_basic_config

write_basic_config()
```

最后，如果您想在自己的数据集上训练模型，请查看[创建训练数据集](create_dataset)指南，了解如何创建可与训练脚本配合使用的数据集。

> [!TIP]
> 以下各部分重点介绍了训练脚本中对于理解如何修改脚本至关重要的内容，但并未详细涵盖脚本的每个方面。如果您希望深入了解，请随时阅读脚本，并告诉我们您是否有任何疑问或顾虑。

## 脚本参数

训练脚本提供了许多参数，帮助你自定义训练过程。所有参数及其说明都可以在 [`parse_args()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L190) 函数中找到。训练脚本为每个参数提供了默认值，例如训练批次大小和学习率，但你也可以在训练命令中设置自己的值。

例如，若要使用 fp16 格式通过混合精度加速训练，请在训练命令中添加 `--mixed_precision` 参数：

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

大多数参数与[文本到图像](text2image#script-parameters)训练指南中的参数相同，因此我们直接进入 Kandinsky 训练脚本的讲解！

### Min-SNR 加权

[Min-SNR](https://huggingface.co/papers/2303.09556) 加权策略可以通过重新平衡损失来加速收敛，从而辅助训练。训练脚本支持预测 `epsilon`（噪声）或 `v_prediction`，而 Min-SNR 与这两种预测类型均兼容。此加权策略仅支持 PyTorch。

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

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

## 训练脚本

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

[`main()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L441) 函数包含用于准备数据集和训练模型的代码。

您会立即注意到的主要区别之一是，训练脚本除了加载调度器和分词器外，还会加载一个 [CLIPImageProcessor](https://huggingface.co/docs/transformers/main/en/model_doc/clip#transformers.CLIPImageProcessor) 用于图像预处理，以及一个 [CLIPVisionModelWithProjection](https://huggingface.co/docs/transformers/main/en/model_doc/clip#transformers.CLIPVisionModelWithProjection) 模型用于对图像进行编码：

```py
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](/docs/diffusers/main/zh/api/models/prior_transformer#diffusers.PriorTransformer) 来生成图像嵌入，因此你需要配置优化器以学习先验模式的参数。

```py
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](https://huggingface.co/docs/transformers/main/en/model_doc/clip#transformers.CLIPImageProcessor) 进行 [预处理](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L632)：

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

最后，[训练循环](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L718)将输入图像转换为潜在表示，为图像嵌入添加噪声，并进行预测：

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

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

[`main()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L440) 函数包含用于准备数据集和训练模型的代码。

与之前的模型不同，解码器会初始化一个 [VQModel](/docs/diffusers/main/zh/api/models/vq#diffusers.VQModel) 将潜在变量解码为图像，并使用 [UNet2DConditionModel](/docs/diffusers/main/zh/api/models/unet2d-cond#diffusers.UNet2DConditionModel)：

```py
with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
    vae = VQModel.from_pretrained(
        args.pretrained_decoder_model_name_or_path, subfolder="movq", dtype=weight_dtype
    ).eval()
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="image_encoder", dtype=weight_dtype
    ).eval()
unet = UNet2DConditionModel.from_pretrained(args.pretrained_decoder_model_name_or_path, subfolder="unet")
```

接下来，脚本包含若干图像变换以及一个用于对图像应用这些变换并返回像素值的 [preprocessing](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L622) 函数：

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

最后，[训练循环](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L706)负责将图像转换为潜在表示、添加噪声并预测噪声残差。

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

```py
model_pred = unet(noisy_latents, timesteps, None, added_cond_kwargs=added_cond_kwargs).sample[:, :4]
```

## 运行脚本

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

您将在 [Naruto BLIP captions](https://huggingface.co/datasets/lambdalabs/naruto-blip-captions) 数据集上进行训练，以生成您自己的 Naruto 角色，但您也可以按照 [Create a dataset for training](create_dataset) 指南创建并训练自己的数据集。将环境变量 `DATASET_NAME` 设置为 Hub 上数据集的名称，或者如果您要使用自己的文件进行训练，则将环境变量 `TRAIN_DIR` 设置为指向您数据集的路径。

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

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

```bash
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"
```

```bash
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_decoder.py \
  --dataset_name=$DATASET_NAME \
  --resolution=768 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --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-decoder-naruto-model"
```

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

```py
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]
```

> [!TIP]
> 您可以随意将 `kandinsky-community/kandinsky-2-2-decoder` 替换为您自己训练的解码器检查点！

```py
from diffusers import AutoPipelineForText2Image
import torch

pipeline = AutoPipelineForText2Image.from_pretrained("path/to/saved/model", dtype=torch.float16)
pipeline.enable_model_cpu_offload()

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

对于解码器模型，你也可以从保存的检查点执行推理，这对于查看中间结果可能很有用。在这种情况下，将检查点加载到 UNet 中：

```py
from diffusers import AutoPipelineForText2Image, UNet2DConditionModel

unet = UNet2DConditionModel.from_pretrained("path/to/saved/model" + "/checkpoint-<N>/unet")

pipeline = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", unet=unet, dtype=torch.float16)
pipeline.enable_model_cpu_offload()

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

## 后续步骤

恭喜您成功训练了 Kandinsky 2.2 模型！如需了解如何使用您的新模型，以下指南可能对您有所帮助：

- 阅读 [Kandinsky](../api/pipelines/kandinsky) 指南，了解如何使用它完成多种不同的任务（文本生成图像、图像生成图像、图像修复、插值），以及如何将其与 ControlNet 结合使用。
- 请参阅 [DreamBooth](dreambooth) 和 [LoRA](lora) 训练指南，了解如何使用少量示例图像训练个性化的 Kandinsky 模型。这两种训练技术甚至可以组合使用！