LlamaFactory
算力平台:

ControlNet ​

ControlNet 模型是在另一个预训练模型基础上训练的适配器。它通过使用额外的输入图像对模型进行条件化，从而对图像生成进行更精细的控制。输入图像可以是 Canny 边缘、深度图、人体姿态等。

如果你在 vRAM 有限的 GPU 上进行训练，建议在训练命令中启用 gradient_checkpointing、gradient_accumulation_steps 和 mixed_precision 参数。你还可以通过使用 xFormers 的内存高效注意力机制来减少内存占用。JAX/Flax 训练也支持在 TPUs 和 GPUs 上高效训练，但不支持梯度检查点或 xFormers。如果你希望使用 Flax 进行更快的训练，建议使用内存大于 30GB 的 GPU。

本指南将探讨 train_controlnet.py 训练脚本，帮助你熟悉它，并了解如何根据自己的需求进行调整。

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

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

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

初始化一个 🤗 Accelerate 环境：

bash
accelerate config

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

bash
accelerate config default

或者，如果你的环境不支持交互式 shell，比如笔记本，你可以使用：

py
from accelerate.utils import write_basic_config

write_basic_config()

最后，如果你想在自己的数据集上训练模型，请参阅创建用于训练的数据集指南，了解如何创建与训练脚本兼容的数据集。

脚本参数 ​

训练脚本提供了许多参数，帮助你自定义训练过程。所有参数及其描述都在 parse_args() 函数中。此函数为每个参数提供了默认值，例如训练批次大小和学习率，但你也可以在训练命令中设置自己的值。

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

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

许多基本和重要的参数已在 文本到图像 训练指南中描述，因此本指南仅关注与 ControlNet 相关的参数：

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

Min-SNR 权重 ​

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

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

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

训练脚本 ​

关于脚本参数的详细说明，可以在 文本到图像 训练指南中找到。本指南将重点介绍 ControlNet 脚本的相关部分。

训练脚本中有一个 make_train_dataset 函数，用于对数据集进行预处理，包括图像变换和标题标记化。你将看到，除了常规的标题标记化和图像变换外，该脚本还包含了条件图像的变换。

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

在 main() 函数中，你可以找到加载分词器、文本编码器、调度器和模型的代码。这里也是加载 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)

优化器 被设置为更新 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,
)

最后，在训练循环中，条件文本嵌入和图像被传递给 ControlNet 模型的下采样块和中间块：

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,
)

如果你想了解更多关于训练循环的工作原理，请查看理解管道、模型和调度器教程，该教程详细介绍了去噪过程的基本模式。

启动脚本 ​

现在你已经准备好启动训练脚本了！🚀

本指南使用fusing/fill50k数据集，但请记住，你可以创建并使用自己的数据集（参见为训练创建数据集指南）。

将环境变量 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 参数。

一旦训练完成，你就可以使用新训练的模型进行推理！

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

controlnet = ControlNetModel.from_pretrained("path/to/controlnet", torch_dtype=torch.float16)
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
    "path/to/base/model", controlnet=controlnet, torch_dtype=torch.float16
).to("cuda")

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 脚本来训练 SDXL 模型的 ControlNet 适配器。

SDXL 训练脚本在 SDXL 训练 指南中进行了更详细的讨论。

下一步 ​

恭喜你成功训练了自己的 ControlNet！要了解更多如何使用你的新模型，以下指南可能会有所帮助：

Copyright © 2024-2025 llamafactory.cn