# Custom Diffusion

[Custom Diffusion](https://huggingface.co/papers/2212.04488) 是一种用于个性化图像生成模型的训练技术。与 Textual Inversion、DreamBooth 和 LoRA 类似，Custom Diffusion 仅需少量（约 4-5 张）示例图像。该技术仅训练交叉注意力层中的权重，并使用一个特殊单词来表示新学习的概念。Custom Diffusion 的独特之处在于它可以同时学习多个概念。

如果您在 vRAM 有限的 GPU 上进行训练，建议通过 `--enable_xformers_memory_efficient_attention` 启用 xFormers，以在更低的 vRAM 需求（16GB）下实现更快的训练速度。若要进一步节省内存，可在训练参数中添加 `--set_grads_to_none`，将梯度设置为 `None` 而非零（此选项可能会引发一些问题，如果遇到任何问题，请尝试移除该参数）。

本指南将深入讲解 [train_custom_diffusion.py](https://github.com/huggingface/diffusers/blob/main/examples/custom_diffusion/train_custom_diffusion.py) 脚本，帮助你更熟悉该脚本，并了解如何将其适配到自身的使用场景中。

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

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

导航至包含训练脚本的示例文件夹，并安装所需的依赖项：

```bash
cd examples/custom_diffusion
pip install -r requirements.txt
pip install clip-retrieval
```

> [!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]
> 以下各节重点介绍了训练脚本中对于理解如何修改脚本至关重要的部分，但并未详细涵盖脚本的每个方面。如果您希望了解更多，请随时阅读[脚本](https://github.com/huggingface/diffusers/blob/main/examples/custom_diffusion/train_custom_diffusion.py)，并告诉我们您是否有任何疑问或顾虑。

## 脚本参数

训练脚本包含了所有用于帮助自定义训练过程的参数。这些参数位于 [`parse_args()`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L319) 函数中。该函数带有默认值，但您也可以根据需要在训练命令中设置自己的值。

例如，要更改输入图像的分辨率：

```bash
accelerate launch train_custom_diffusion.py \
  --resolution=256
```

许多基本参数已在 [DreamBooth](dreambooth#script-parameters) 训练指南中介绍，因此本指南将重点介绍 Custom Diffusion 特有的参数：

- `--freeze_model`：冻结交叉注意力层中的 key 和 value 参数；默认值为 `crossattn_kv`，但您可以将其设置为 `crossattn` 以训练交叉注意力层中的所有参数
- `--concepts_list`：要学习多个概念，请提供一个包含这些概念的 JSON 文件的路径
- `--modifier_token`：用于表示所学概念的特殊词
- `--initializer_token`：用于初始化 `modifier_token` 嵌入的特殊词

### 先验保持损失

先验保持损失是一种利用模型自身生成的样本来辅助其学习如何生成更多样化图像的方法。由于这些生成的样本图像与您提供的图像属于同一类别，它们有助于模型保留其对该类别已学到的知识，并利用其对该类别的既有认知来构建新的组合。

先验保持损失的许多参数在 [DreamBooth](dreambooth#prior-preservation-loss) 训练指南中有详细说明。

### 正则化

Custom Diffusion 包括使用少量真实图像对目标图像进行训练，以防止过拟合。正如你可以想象的那样，当只使用少量图像进行训练时，这很容易做到！使用 `clip_retrieval` 下载 200 张真实图像。`class_prompt` 应与目标图像属于同一类别。这些图像存储在 `class_data_dir` 中。

```bash
python retrieve.py --class_prompt cat --class_data_dir real_reg/samples_cat --num_class_images 200
```

要启用正则化，请添加以下参数：

- `--with_prior_preservation`：是否使用先验保持损失
- `--prior_loss_weight`：控制先验保持损失对模型的影响
- `--real_prior`：是否使用少量真实图像以防止过拟合

```bash
accelerate launch train_custom_diffusion.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="./real_reg/samples_cat" \
  --class_prompt="cat" \
  --real_prior=True \
```

## 训练脚本

> [!TIP]
> Custom Diffusion 训练脚本中的许多代码与 [DreamBooth](dreambooth#training-script) 脚本类似。本指南将重点介绍与 Custom Diffusion 相关的代码。

Custom Diffusion 训练脚本包含两个数据集类：

- [`CustomDiffusionDataset`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L165)：预处理用于训练的图像、类别图像和提示词
- [`PromptDataset`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L148)：为生成类别图像准备提示词

接下来，`modifier_token` 被[添加到分词器](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L811)中，转换为 token id，并调整 token 嵌入的大小以容纳新的 `modifier_token`。然后，`modifier_token` 的嵌入使用 `initializer_token` 的嵌入进行初始化。文本编码器中的所有参数都被冻结，但 token 嵌入除外，因为模型试图学习将 token 嵌入与概念关联起来。

```py
params_to_freeze = itertools.chain(
    text_encoder.text_model.encoder.parameters(),
    text_encoder.text_model.final_layer_norm.parameters(),
    text_encoder.text_model.embeddings.position_embedding.parameters(),
)
freeze_params(params_to_freeze)
```

接下来，你需要将 [Custom Diffusion 权重](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L911C3-L911C3) 添加到注意力层中。这一步对于确保注意力权重的形状和尺寸正确，以及为每个 UNet 块设置适当数量的注意力处理器至关重要。

```py
st = unet.state_dict()
for name, _ in unet.attn_processors.items():
    cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
    if name.startswith("mid_block"):
        hidden_size = unet.config.block_out_channels[-1]
    elif name.startswith("up_blocks"):
        block_id = int(name[len("up_blocks.")])
        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
    elif name.startswith("down_blocks"):
        block_id = int(name[len("down_blocks.")])
        hidden_size = unet.config.block_out_channels[block_id]
    layer_name = name.split(".processor")[0]
    weights = {
        "to_k_custom_diffusion.weight": st[layer_name + ".to_k.weight"],
        "to_v_custom_diffusion.weight": st[layer_name + ".to_v.weight"],
    }
    if train_q_out:
        weights["to_q_custom_diffusion.weight"] = st[layer_name + ".to_q.weight"]
        weights["to_out_custom_diffusion.0.weight"] = st[layer_name + ".to_out.0.weight"]
        weights["to_out_custom_diffusion.0.bias"] = st[layer_name + ".to_out.0.bias"]
    if cross_attention_dim is not None:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=train_kv,
            train_q_out=train_q_out,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        ).to(unet.device)
        custom_diffusion_attn_procs[name].load_state_dict(weights)
    else:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=False,
            train_q_out=False,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        )
del st
unet.set_attn_processor(custom_diffusion_attn_procs)
custom_diffusion_layers = AttnProcsLayers(unet.attn_processors)
```

[optimizer](https://github.com/huggingface/diffusers/blob/84cd9e8d01adb47f046b1ee449fc76a0c32dc4e2/examples/custom_diffusion/train_custom_diffusion.py#L982) 被初始化以更新交叉注意力层的参数：

```py
optimizer = optimizer_class(
    itertools.chain(text_encoder.get_input_embeddings().parameters(), custom_diffusion_layers.parameters())
    if args.modifier_token is not None
    else custom_diffusion_layers.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)
```

在[训练循环](https://github.com/huggingface/diffusers/blob/84cd9e8d01adb47f046b1ee449fc76a0c32dc4e2/examples/custom_diffusion/train_custom_diffusion.py#L1048)中，重要的是只更新你试图学习的那个概念的嵌入。这意味着将所有其他 token 嵌入的梯度设置为零：

```py
if args.modifier_token is not None:
    if accelerator.num_processes > 1:
        grads_text_encoder = text_encoder.module.get_input_embeddings().weight.grad
    else:
        grads_text_encoder = text_encoder.get_input_embeddings().weight.grad
    index_grads_to_zero = torch.arange(len(tokenizer)) != modifier_token_id[0]
    for i in range(len(modifier_token_id[1:])):
        index_grads_to_zero = index_grads_to_zero & (
            torch.arange(len(tokenizer)) != modifier_token_id[i]
        )
    grads_text_encoder.data[index_grads_to_zero, :] = grads_text_encoder.data[
        index_grads_to_zero, :
    ].fill_(0)
```

## 运行脚本

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

在本指南中，你将下载并使用这些示例 [猫图片](https://www.cs.cmu.edu/~custom-diffusion/assets/data.zip)。如果你愿意，也可以创建并使用自己的数据集（参见 [创建训练数据集](create_dataset) 指南）。

将环境变量 `MODEL_NAME` 设置为 Hub 上的模型 ID 或本地模型的路径，将 `INSTANCE_DIR` 设置为您刚刚下载猫图片的路径，将 `OUTPUT_DIR` 设置为您希望保存模型的位置。您将使用 `<new1>` 作为特殊词，以将新学习的嵌入与之关联。该脚本会创建并保存模型检查点以及 pytorch_custom_diffusion_weights.bin 文件到您的仓库中。

若要使用 Weights and Biases 监控训练进度，请在训练命令中添加 `--report_to=wandb` 参数，并使用 `--validation_prompt` 指定验证提示。这对于调试和保存中间结果非常有用。

> [!TIP]
> 如果您正在使用人脸图像进行训练，Custom Diffusion 团队发现以下参数效果良好：
>
> - `--learning_rate=5e-6` - `--max_train_steps` 可以设置在 1000 到 2000 之间 - `--freeze_model=crossattn` - 训练时至少使用 15-20 张图片

```bash
export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export OUTPUT_DIR="path-to-save-model"
export INSTANCE_DIR="./data/cat"

accelerate launch train_custom_diffusion.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --class_data_dir=./real_reg/samples_cat/ \
  --with_prior_preservation \
  --real_prior \
  --prior_loss_weight=1.0 \
  --class_prompt="cat" \
  --num_class_images=200 \
  --instance_prompt="photo of a <new1> cat"  \
  --resolution=512  \
  --train_batch_size=2  \
  --learning_rate=1e-5  \
  --lr_warmup_steps=0 \
  --max_train_steps=250 \
  --scale_lr \
  --hflip  \
  --modifier_token "<new1>" \
  --validation_prompt="<new1> cat sitting in a bucket" \
  --report_to="wandb" \
  --push_to_hub
```

如果您提供一个包含每个待学习概念相关详情的 [JSON](https://github.com/adobe-research/custom-diffusion/blob/main/assets/concept_list.json) 文件，Custom Diffusion 也可以学习多个概念。

运行 clip-retrieval 以收集一些真实图像，用于正则化：

```bash
pip install clip-retrieval
python retrieve.py --class_prompt {} --class_data_dir {} --num_class_images 200
```

然后，您可以运行该脚本：

```bash
export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export OUTPUT_DIR="path-to-save-model"

accelerate launch train_custom_diffusion.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --output_dir=$OUTPUT_DIR \
  --concepts_list=./concept_list.json \
  --with_prior_preservation \
  --real_prior \
  --prior_loss_weight=1.0 \
  --resolution=512  \
  --train_batch_size=2  \
  --learning_rate=1e-5  \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --num_class_images=200 \
  --scale_lr \
  --hflip  \
  --modifier_token "<new1>+<new2>" \
  --push_to_hub
```

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

```py
import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet.load_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin")
pipeline.load_textual_inversion("path-to-save-model", weight_name="<new1>.bin")

image = pipeline(
    "<new1> cat sitting in a bucket",
    num_inference_steps=100,
    guidance_scale=6.0,
    eta=1.0,
).images[0]
image.save("cat.png")
```

```py
import torch
from huggingface_hub.repocard import RepoCard
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
model_id = "sayakpaul/custom-diffusion-cat-wooden-pot"
pipeline.unet.load_attn_procs(model_id, weight_name="pytorch_custom_diffusion_weights.bin")
pipeline.load_textual_inversion(model_id, weight_name="<new1>.bin")
pipeline.load_textual_inversion(model_id, weight_name="<new2>.bin")

image = pipeline(
    "the <new1> cat sculpture in the style of a <new2> wooden pot",
    num_inference_steps=100,
    guidance_scale=6.0,
    eta=1.0,
).images[0]
image.save("multi-subject.png")
```

## 后续步骤

恭喜您使用 Custom Diffusion 训练了一个模型！🎉 如需了解更多：

- Read the [Multi-Concept Customization of Text-to-Image Diffusion](https://www.cs.cmu.edu/~custom-diffusion/) blog post to learn more details about the experimental results from the Custom Diffusion team.