# ControlNet

[ControlNet](https://hf.co/papers/2302.05543) 模型是在另一个预训练模型之上训练的适配器。它通过引入额外的输入图像作为条件，从而实现对图像生成过程更精细的控制。该输入图像可以是 Canny 边缘、深度图、人体姿态等。

如果用于训练的 GPU 显存有限，建议在训练命令中启用 `gradient_checkpointing`、`gradient_accumulation_steps` 和 `mixed_precision` 参数。此外，还可以使用 [xFormers](../optimization/attention_backends) 提供的内存高效注意力机制来降低内存占用。

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

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

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

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

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

## 脚本参数

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

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

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

许多基础且重要的参数已在[文本到图像](text2image#script-parameters)训练指南中介绍，因此本指南仅聚焦于与 ControlNet 相关的参数：

- `--max_train_samples`：训练样本的数量；您可以降低此值以加快训练速度，但如果您想要流式处理非常大的数据集，则需要在训练命令中包含此参数以及 `--streaming` 参数
- `--gradient_accumulation_steps`：执行反向传播之前累积的更新步数；这允许你使用比 GPU 显存通常能处理更大的批量大小进行训练

### Min-SNR 加权

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

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

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

## 训练脚本

与脚本参数类似，[文本到图像](text2image#training-script)训练指南中提供了训练脚本的总体介绍。本指南将重点介绍 ControlNet 脚本中的相关部分。

训练脚本中有一个 [`make_train_dataset`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/controlnet/train_controlnet.py#L582) 函数，用于通过图像变换和字幕分词对数据集进行预处理。您会看到，除了常规的字幕分词和图像变换外，该脚本还包含针对条件图像的变换。

> [!TIP]
> 如果您正在 TPU 上流式传输数据集，性能可能会受到 🤗 Datasets 库的瓶颈限制，因为该库并未针对图像进行优化。为确保最大吞吐量，建议您探索其他数据集格式，例如 [WebDataset](https://github.com/webdataset/webdataset)、[TorchData](https://github.com/pytorch/data) 和 [TensorFlow Datasets](https://www.tensorflow.org/datasets/tfless_tfds)。

```py
conditioning_image_transforms = transforms.Compose(
    [
        transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
        transforms.CenterCrop(args.resolution),
        transforms.ToTensor(),
    ]
)
```

在 [`main()`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/controlnet/train_controlnet.py#L713) 函数中，你可以找到加载 tokenizer、text encoder、scheduler 和模型的代码。ControlNet 模型也是在这里加载的，它可以从现有权重加载，也可以从 UNet 随机初始化：

```py
if args.controlnet_model_name_or_path:
    logger.info("Loading existing controlnet weights")
    controlnet = ControlNetModel.from_pretrained(args.controlnet_model_name_or_path)
else:
    logger.info("Initializing controlnet weights from unet")
    controlnet = ControlNetModel.from_unet(unet)
```

[optimizer](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/controlnet/train_controlnet.py#L871) 已配置为更新 ControlNet 参数：

```py
params_to_optimize = controlnet.parameters()
optimizer = optimizer_class(
    params_to_optimize,
    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/64603389da01082055a901f2883c4810d1144edb/examples/controlnet/train_controlnet.py#L943) 中，条件文本嵌入和图像会被传入 ControlNet 模型的 down 和 mid 块：

```py
encoder_hidden_states = text_encoder(batch["input_ids"])[0]
controlnet_image = batch["conditioning_pixel_values"].to(dtype=weight_dtype)

down_block_res_samples, mid_block_res_sample = controlnet(
    noisy_latents,
    timesteps,
    encoder_hidden_states=encoder_hidden_states,
    controlnet_cond=controlnet_image,
    return_dict=False,
)
```

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

## 运行脚本

现在，您可以启动训练脚本了！🚀

本指南使用 [fusing/fill50k](https://huggingface.co/datasets/fusing/fill50k) 数据集，但请记住，如果您愿意，也可以创建并使用自己的数据集（参见 [为训练创建数据集](create_dataset) 指南）。

将环境变量 `MODEL_NAME` 设置为 Hub 上的模型 ID 或本地模型的路径，并将 `OUTPUT_DIR` 设置为你希望保存模型的位置。

下载以下图像，用于条件化您的训练：

```bash
wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_1.png
wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_2.png
```

在运行脚本之前，还有一件事需要注意！根据你所使用的 GPU，你可能需要启用某些优化来训练 ControlNet。此脚本的默认配置大约需要 38GB 的 vRAM。如果你使用多块 GPU 进行训练，请在 `accelerate launch` 命令中添加 `--multi_gpu` 参数。

在 16GB 的 GPU 上，你可以使用 bitsandbytes 8 位优化器和梯度检查点来优化训练过程。安装 bitsandbytes：

```py
pip install bitsandbytes
```

然后，在训练命令中添加以下参数：

```bash
accelerate launch train_controlnet.py \
  --gradient_checkpointing \
  --use_8bit_adam \
```

在 12GB 的 GPU 上，你需要使用 bitsandbytes 8-bit 优化器、梯度检查点、xFormers，并将梯度设置为 `None` 而非零，以降低内存占用。

```bash
accelerate launch train_controlnet.py \
  --use_8bit_adam \
  --gradient_checkpointing \
  --enable_xformers_memory_efficient_attention \
  --set_grads_to_none \
```

在 8GB 的 GPU 上，你需要使用 [DeepSpeed](https://www.deepspeed.ai/) 将部分张量从 vRAM 卸载到 CPU 或 NVME，以便在更少的 GPU 显存下完成训练。

运行以下命令以配置您的 🤗 Accelerate 环境：

```bash
accelerate config
```

在配置过程中，请确认您希望使用 DeepSpeed stage 2。现在，通过结合 DeepSpeed stage 2、fp16 混合精度以及将模型参数和优化器状态卸载到 CPU，应该可以在 8GB 以下的 vRAM 上完成训练。其缺点是这需要更多的系统内存（约 25 GB）。有关更多配置选项，请参阅 [DeepSpeed 文档](https://huggingface.co/docs/accelerate/usage_guides/deepspeed)。您的配置文件应类似于以下内容：

```bash
compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 4
  offload_optimizer_device: cpu
  offload_param_device: cpu
  zero3_init_flag: false
  zero_stage: 2
distributed_type: DEEPSPEED
```

您还应将默认的 Adam 优化器更改为 DeepSpeed 的优化版本 Adam [`deepspeed.ops.adam.DeepSpeedCPUAdam`](https://deepspeed.readthedocs.io/en/latest/optimizers.html#adam-cpu)，以获得显著的速度提升。启用 `DeepSpeedCPUAdam` 要求系统的 CUDA 工具链版本与 PyTorch 安装时使用的版本一致。

bitsandbytes 的 8 位优化器目前似乎与 DeepSpeed 不兼容。

就是这样！你无需在训练命令中添加任何额外参数。

```bash
export MODEL_DIR="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="path/to/save/model"

accelerate launch train_controlnet.py \
 --pretrained_model_name_or_path=$MODEL_DIR \
 --output_dir=$OUTPUT_DIR \
 --dataset_name=fusing/fill50k \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./conditioning_image_1.png" "./conditioning_image_2.png" \
 --validation_prompt "red circle with blue background" "cyan circle with brown floral background" \
 --train_batch_size=1 \
 --gradient_accumulation_steps=4 \
 --push_to_hub
```

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

```py
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import torch

controlnet = ControlNetModel.from_pretrained("path/to/controlnet", dtype=torch.float16)
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
    "path/to/base/model", controlnet=controlnet, dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

control_image = load_image("./conditioning_image_1.png")
prompt = "pale golden rod circle with old lace background"

generator = torch.manual_seed(0)
image = pipeline(prompt, num_inference_steps=20, generator=generator, image=control_image).images[0]
image.save("./output.png")
```

## Stable Diffusion XL

Stable Diffusion XL (SDXL) 是一个强大的文本生成图像模型，能够生成高分辨率图像，并在其架构中增加了第二个文本编码器。使用 [`train_controlnet_sdxl.py`](https://github.com/huggingface/diffusers/blob/main/examples/controlnet/train_controlnet_sdxl.py) 脚本为 SDXL 模型训练 ControlNet 适配器。

SDXL 训练脚本在 [SDXL 训练](sdxl) 指南中有更详细的介绍。

## 后续步骤

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

- 了解如何在多种任务中使用 [ControlNet](../using-diffusers/controlnet) 进行推理。