Diffusers documentation
ControlNet
ControlNet
ControlNet 是一种基于预训练模型的适配器架构。它通过额外输入的条件图像（如边缘检测图、深度图、人体姿态图等），实现对生成图像的精细化控制。
在显存有限的GPU上训练时，建议启用训练命令中的 gradient_checkpointing（梯度检查点）、gradient_accumulation_steps（梯度累积步数）和 mixed_precision（混合精度）参数。还可使用 xFormers 的内存高效注意力机制进一步降低显存占用。虽然JAX/Flax训练支持在TPU和GPU上高效运行，但不支持梯度检查点和xFormers。若需通过Flax加速训练，建议使用显存大于30GB的GPU。
本指南将解析 train_controlnet.py 训练脚本，帮助您理解其逻辑并适配自定义需求。
运行脚本前，请确保从源码安装库：
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .然后进入包含训练脚本的示例目录，安装所需依赖：
cd examples/controlnet
pip install -r requirements.txt🤗 Accelerate 是一个支持多GPU/TPU训练和混合精度的库，它能根据硬件环境自动配置训练方案。参阅 🤗 Accelerate 快速入门 了解更多。
初始化🤗 Accelerate环境：
accelerate config
若要创建默认配置（不进行交互式选择）：
accelerate config default
若环境不支持交互式shell（如notebook），可使用：
from accelerate.utils import write_basic_config
write_basic_config()最后，如需训练自定义数据集，请参阅 创建训练数据集 指南了解数据准备方法。
下文重点解析脚本中的关键模块，但不会覆盖所有实现细节。如需深入了解，建议直接阅读 脚本源码，如有疑问欢迎反馈。
脚本参数
训练脚本提供了丰富的可配置参数，所有参数及其说明详见 parse_args() 函数。虽然该函数已为每个参数提供默认值（如训练批大小、学习率等），但您可以通过命令行参数覆盖这些默认值。
例如，使用fp16混合精度加速训练, 可使用--mixed_precision参数
accelerate launch train_controlnet.py \
--mixed_precision="fp16"基础参数说明可参考 文生图 训练指南，此处重点介绍ControlNet相关参数：
--max_train_samples: 训练样本数量，减少该值可加快训练，但对超大数据集需配合--streaming参数使用--gradient_accumulation_steps: 梯度累积步数，通过分步计算实现显存受限情况下的更大批次训练
Min-SNR加权策略
Min-SNR 加权策略通过重新平衡损失函数加速模型收敛。虽然训练脚本支持预测 epsilon（噪声）或 v_prediction，但Min-SNR对两种预测类型均兼容。该策略仅适用于PyTorch版本，Flax训练脚本暂不支持。
推荐值设为5.0：
accelerate launch train_controlnet.py \ --snr_gamma=5.0
训练脚本
与参数说明类似，训练流程的通用解析可参考 文生图 指南。此处重点分析ControlNet特有的实现。
脚本中的 make_train_dataset 函数负责数据预处理，除常规的文本标注分词和图像变换外，还包含条件图像的特效处理：
在TPU上流式加载数据集时，🤗 Datasets库可能成为性能瓶颈（因其未针对图像数据优化）。建议考虑 WebDataset、TorchData 或 TensorFlow Datasets 等高效数据格式。
conditioning_image_transforms = transforms.Compose(
[
transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
transforms.CenterCrop(args.resolution),
transforms.ToTensor(),
]
)在 main() 函数中，代码会加载分词器、文本编码器、调度器和模型。此处也是ControlNet模型的加载点（支持从现有权重加载或从UNet随机初始化）：
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参数进行更新：
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的下采样和中层模块：
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 为模型保存路径。
下载训练用的条件图像：
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型号，可能需要启用特定优化。默认配置需要约38GB显存。若使用多GPU训练，请在 accelerate launch 命令中添加 --multi_gpu 参数。
16GB显卡可使用bitsandbytes 8-bit优化器和梯度检查点：
pip install bitsandbytes
训练命令添加以下参数：
accelerate launch train_controlnet.py \ --gradient_checkpointing \ --use_8bit_adam \
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训练完成后即可进行推理：
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训练！如需进一步了解模型应用，以下指南可能有所帮助：
- 学习如何 使用ControlNet 进行多样化任务的推理