[特殊字符] Diffusers 实战：从零训练 ControlNet，让 Stable Diffusion 输出精确可控的生成结果

🤗 Diffusers 实战：从零训练 ControlNet，让 Stable Diffusion 输出精确可控的生成结果

【免费下载链接】diffusers🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers

本指南以 🤗 Diffusers 官方仓库中的 ControlNet 训练文档为主体，结合examples/controlnet/train_controlnet.py训练脚本与src/diffusers/models/controlnets/controlnet.py模型实现，完整讲解如何安装依赖、配置 🤗 Accelerate、在 fill50k 合成数据集上训练一个圆圈填充 ControlNet，并针对 38GB / 20GB / 16GB / 12GB / 8GB 不同显存给出可复制的优化方案与最终推理代码。读完本文，你将掌握 ControlNet 训练的全流程、关键参数含义、显存优化原理，以及用StableDiffusionControlNetPipeline加载自训模型进行推理的完整闭环。

ControlNet 是什么：给扩散模型加一根"控制线"

ControlNet 出自论文Adding Conditional Control to Text-to-Image Diffusion Models（Lvmin Zhang 与 Maneesh Agrawala），它本质上是一类在预训练模型之上训练的适配器（adapter）。通过在 UNet 的 down-block 与 mid-block 注入额外的残差信号，ControlNet 允许你用一张额外的输入图像来控制生成过程——这张条件图可以是 Canny 边缘图、深度图、人体姿态骨架，也可以像本文示例一样，是一张带有简单几何图形的合成图。

本文的示例源自原始 ControlNet 仓库的训练文档，并在 🤗 Diffusers 中以可复现的脚本形式落地：使用小型合成数据集 fill50k训练一个"圆圈填充" ControlNet——给定一张"圆在背景中"的条件图，模型学会在指定位置、以指定颜色生成圆。基础模型选用Stable Diffusion 1.5（stable-diffusion-v1-5/stable-diffusion-v1-5），这也是原始 ControlNet 系列模型的训练基础。原则上，ControlNet 可以为任何与 Stable Diffusion 兼容的模型（如 CompVis/stable-diffusion-v1-4、stabilityai/stable-diffusion-2-1）追加条件控制能力。

环境准备：从源码安装与依赖

官方强烈推荐从源码安装diffusers 并保持最新，因为示例脚本更新频繁，且会安装针对示例定制的依赖。请在全新虚拟环境中依次执行：

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

然后进入示例目录并安装训练依赖：

cd examples/controlnet pip install -r requirements.txt

requirements.txt 的内容如下：

accelerate>=0.16.0 torchvision transformers>=4.25.1 ftfy tensorboard datasets

接下来初始化 🤗 Accelerate 环境。Accelerate 负责多 GPU / TPU 与混合精度训练的设备编排，它会根据你的硬件自动配置训练环境。三种初始化方式任选其一：

# 交互式配置（会逐项询问硬件环境） accelerate config

# 不回答任何问题，直接使用默认配置 accelerate config default

如果你的环境不支持交互式 shell（例如 Notebook），可以用 Python 方式写入基础配置：

from accelerate.utils import write_basic_config write_basic_config()

如果你打算使用自己的数据集训练，请参考《为训练创建数据集》指南，了解如何构造与训练脚本兼容的数据集（需要包含image、conditioning_image、text三个字段，具体字段名可通过参数自定义）。

训练脚本与核心参数解析

所有可配置参数都定义在 train_controlnet.py 的parse_args()函数（L251 起）中，每个参数都带有默认值。下面是与 ControlNet 训练强相关的核心参数：

此外，英文版指南还介绍了Min-SNR 加权（论文Efficient Diffusion Training via Min-SNR Weighting Strategy）技巧：通过重新平衡损失来加速收敛，推荐--snr_gamma=5.0。需要说明的是，该参数当前在 PyTorch 版 train_controlnet.py 中尚未实现，但其完整实现可见于 Flax 版 train_controlnet_flax.py（L305 参数定义、L922-L924 的 SNR 截断加权逻辑）；若要在 PyTorch 版脚本中使用，需要自行在损失计算处加入相应的加权逻辑。

启动训练：fill50k 圆圈填充示例

先下载两张官方提供的验证条件图（训练过程中会定期用它们生成对比图，跟踪收敛进度）：

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

设置MODEL_DIR（Hub 模型 ID 或本地模型目录，对应--pretrained_model_name_or_path）与OUTPUT_DIR（模型保存目录，对应--output_dir），然后启动训练：

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=4 \ --push_to_hub

注意：该默认配置需要约 38GB 显存。训练脚本会在你的输出目录中生成并保存diffusion_pytorch_model.bin文件。

默认情况下，训练日志写入 TensorBoard；想改用 Weights & Biases，只需在命令中追加--report_to wandb。

用梯度累积把显存降到 ~20GB

显存不足时，可以缩小批大小并配合梯度累积：等效总批大小不变（1 × 4 = 4），但峰值显存显著下降：

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

--gradient_accumulation_steps=4意味着每 4 个 mini-batch 才执行一次参数更新，等效批大小仍为 4。在 train_controlnet.py 的训练循环中，accelerator.accumulate(controlnet)（L1045）正是负责这一累积语义的底层实现。

多 GPU 分布式训练

🤗 Accelerate 对多 GPU 训练是"无缝"的。在accelerate launch时追加--multi_gpu与--mixed_precision="fp16"即可，训练脚本本身无需任何改动：

export MODEL_DIR="stable-diffusion-v1-5/stable-diffusion-v1-5" export OUTPUT_DIR="path to save model" accelerate launch --mixed_precision="fp16" --multi_gpu 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=4 \ --mixed_precision="fp16" \ --tracker_project_name="controlnet-demo" \ --report_to=wandb \ --push_to_hub

--tracker_project_name会作为Accelerator.init_trackers的project_name传入（见 train_controlnet.py L991），用于在 wandb/TensorBoard 中区分不同实验。

低显存 GPU 优化方案

官方按显存档位给出了三套经过验证的优化组合，核心思想都是"梯度检查点 + 低精度优化器 + 高效注意力 + 梯度置 None"的组合拳。

16GB GPU：梯度检查点 + 8-bit Adam

需要先安装 bitsandbytes（pip install bitsandbytes），然后追加两个参数：

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 \ --gradient_checkpointing \ --use_8bit_adam \ --push_to_hub

在源码中，--use_8bit_adam会把优化器切换为bnb.optim.AdamW8bit（train_controlnet.py L898-L908），显著压缩优化器状态占用的显存；--gradient_checkpointing则调用controlnet.enable_gradient_checkpointing()（L873-L874），以少量计算换大量显存。

12GB GPU：再加上 xFormers 与梯度置 None

在 16GB 方案基础上，追加--enable_xformers_memory_efficient_attention与--set_grads_to_none：

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 \ --gradient_checkpointing \ --use_8bit_adam \ --enable_xformers_memory_efficient_attention \ --set_grads_to_none \ --push_to_hub

使用前请务必安装 xFormers：pip install xformers。xFormers 在推理与训练中都能加速注意力计算并降低内存占用，安装与使用说明可参考《xFormers 安装指南》。值得注意的是，训练脚本对 xFormers 版本有检查：0.0.16 版本在部分 GPU 上无法用于训练，会打印警告建议升级到 0.0.17+（见 train_controlnet.py L859-L871）。--set_grads_to_none则把optimizer.zero_grad(set_to_none=True)传入训练循环（L1103），比置零更省显存，但会改变部分行为，若出现问题可关闭。

8GB GPU：DeepSpeed stage 2 + CPU 卸载 + fp16

官方声明尚未对 ControlNet 的 DeepSpeed 支持做充分测试——配置确实能省显存，但并未确认该配置一定能成功完成训练，你可能需要自行调整配置。DeepSpeed 可以把张量从显存卸载到 CPU 或 NVMe，代价是需要约 25GB 系统内存。

先用accelerate config配置环境并选择 DeepSpeed stage 2，生成的配置文件应类似：

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

更多 DeepSpeed 配置项可查阅 🤗 Accelerate 的 DeepSpeed 使用指南。然后启动训练：

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 \ --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 \ --gradient_checkpointing \ --enable_xformers_memory_efficient_attention \ --set_grads_to_none \ --mixed_precision fp16 \ --push_to_hub

两个来自官方的额外提示：

• 把默认 Adam 换成 DeepSpeed 的deepspeed.ops.adam.DeepSpeedCPUAdam能获得显著加速，但要求系统的 CUDA 工具链版本与 PyTorch 一致；
• bitsandbytes 的 8-bit 优化器目前与 DeepSpeed不兼容，8GB 方案中请勿同时使用--use_8bit_adam。

训练脚本源码级剖析

理解了参数与命令，再深入 train_controlnet.py 的源码，看训练究竟如何发生。

数据预处理：条件图走独立的变换管线。make_train_dataset（L602）中，目标图像使用"缩放 + 中心裁剪 + ToTensor + Normalize([0.5],[0.5])"的标准管线；而条件图则使用不含 Normalize 的独立变换（L688-L694）：

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

条件图与目标图尺寸必须对齐，因为 ControlNet 的条件 embedding 会与 UNet 中间特征逐元素相加。此外，tokenize_captions支持--proportion_empty_prompts将部分 caption 置空（类似 CFG 训练中的空提示词策略）。

ControlNet 的初始化：复用 UNet 结构。在main()（L734）中，脚本加载 tokenizer、text encoder、DDPMScheduler、VAE 与 UNet 之后，ControlNet 有两种来源（L810-L815）：

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)

ControlNetModel.from_unet的实现位于 src/diffusers/models/controlnets/controlnet.py（L445），它会复制 UNet 的 encoder 结构并随机初始化条件控制分支——这正是 ControlNet 论文"可复用的卷积块"思想的体现：冻结原 UNet，训练一个结构相同但带条件输入的副本。

冻结主干，只训练 ControlNet。L854-L857 将 VAE、UNet、text encoder 全部requires_grad_(False)，仅controlnet.train()；优化器也只接收 ControlNet 参数（L911-L918）：

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

训练循环：条件信号注入 UNet。核心迭代在 L1045-L1103。流程为：VAE 编码目标图为 latent → 采样随机噪声与时间步 → 前向扩散加噪 → text encoder 得到文本嵌入 →ControlNet 接收带噪 latent、时间步、文本嵌入与条件图，输出 down-block 与 mid-block 残差→ 残差注入 UNet 预测噪声：

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, ) model_pred = unet( noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states, down_block_additional_residuals=[sample.to(dtype=weight_dtype) for sample in down_block_res_samples], mid_block_additional_residual=mid_block_res_sample.to(dtype=weight_dtype), return_dict=False, )[0]

损失为目标噪声（epsilon预测型）或速度（v_prediction预测型）与预测值的 MSE（L1089-L1095）。训练中还通过accelerator.clip_grad_norm_做梯度裁剪（L1099-L1100），并支持--checkpointing_steps保存检查点、--resume_from_checkpoint断点续训（L1008-L1032）。

条件编码网络的原生实现。在模型侧，src/diffusers/models/controlnets/controlnet.py 的ControlNetConditioningEmbedding（L66-L108）负责把图像条件映射到 64×64 的 latent 特征空间，其 docstring 直接引用了论文原话：Stable Diffusion 会把 512×512 图像压缩为 64×64 latent，因此 ControlNet 需要一个由 4 个卷积层（4×4 卷积核、2×2 步长、ReLU 激活、通道数 16/32/64/128）组成的微型网络 E(·) 来编码图像条件。该仓库实现使用 3×3 卷积、SiLU 激活，输出层以zero_module零初始化（L94-L96），保证训练初期 ControlNet 输出为零残差、不影响预训练 UNet 的行为——这是 ControlNet 稳定训练的另一个关键设计。

用训练好的 ControlNet 做推理

训练结束后，--output_dir中保存的便是 ControlNet 权重（脚本在 L1160 通过controlnet.save_pretrained(args.output_dir)保存）。推理时，把base_model_path设为--pretrained_model_name_or_path的值、controlnet_path设为--output_dir的值，加载方式与原始 ControlNet 完全一致：

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler from diffusers.utils import load_image import torch base_model_path = "path to model" controlnet_path = "path to controlnet" controlnet = ControlNetModel.from_pretrained(controlnet_path, dtype=torch.float16) pipe = StableDiffusionControlNetPipeline.from_pretrained( base_model_path, controlnet=controlnet, dtype=torch.float16 ) # 用更快的调度器与内存优化加速扩散过程 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) # 未安装 xformers 时请删除下面这行 pipe.enable_xformers_memory_efficient_attention() pipe.enable_model_cpu_offload() control_image = load_image("./conditioning_image_1.png") prompt = "pale golden rod circle with old lace background" # 生成图像 generator = torch.manual_seed(0) image = pipe(prompt, num_inference_steps=20, generator=generator, image=control_image).images[0] image.save("./output.png")

这段代码与训练脚本内部验证阶段（log_validation，L80-L186）的用法完全一致：训练脚本每--validation_steps步也会构建同样的 pipeline，用UniPCMultistepScheduler以 20 步推理生成对比图，并记录到 TensorBoard 或 wandb。

训练效果与验证结果

fill50k 数据集包含 5 万对"条件图（背景 + 圆的位置）→ 目标图（填充颜色后的圆）"样本，收敛非常快。官方以 batch size 8 训练后给出的定性结果是：

• 300 步后：模型已经能大致理解"把圆画在条件图指定的位置"，但圆的形状与背景颜色的贴合还比较粗糙；
• 6000 步后：模型能够精确地在蓝色背景中生成红色圆、在棕色花背景中生成青色圆，位置、颜色、形状都与条件图高度一致。

如果你设置了--validation_image与--validation_prompt，训练日志中的对比图（条件图 vs 生成图）会直观展示这一收敛过程，这也是判断"该停训了"的最直接依据。训练完成后若加了--push_to_hub，脚本还会自动生成带示例图片的模型卡片并上传到 Hub。

扩展：SDXL 与更多 ControlNet 训练脚本

除了 SD1.5，examples/controlnet 目录还提供多种变体：

• train_controlnet_sdxl.py：为 Stable Diffusion XL（双文本编码器架构）训练 ControlNet，详见 README_sdxl.md；
• train_controlnet_sd3.py 与 train_controlnet_flux.py：分别面向 SD3 与 Flux 的 ControlNet 训练，各有独立的 README 与 requirements 文件（如 README_sd3.md、README_flux.md）；
• train_controlnet_flax.py：面向 TPU 的 JAX/Flax 版本，支持流式数据集与 Min-SNR 加权，但注意JAX/Flax 支持已在 diffusers v0.40.0 中移除，运行该脚本需要diffusers<=0.39.x。

脚本质量由 test_controlnet.py 保障：例如test_controlnet_checkpointing_checkpoints_total_limit用hf-internal-testing/tiny-stable-diffusion-torch微型模型 +hf-internal-testing/fill10微型数据集验证检查点数量上限控制逻辑，test_controlnet_sdxl、test_controlnet_sd3、test_controlnet_flux则分别验证各变体脚本可以端到端跑通。

小结

本文完整覆盖了从依赖安装、Accelerate 配置、fill50k 训练、显存分档优化到推理部署的 ControlNet 训练全链路，并结合 train_controlnet.py 与 controlnet.py 源码解释了参数背后的实现机制。无论是复现官方的圆圈填充示例，还是用--train_data_dir+--image_column/--conditioning_image_column/--caption_column适配你自己的数据集（数据构造方法见《为训练创建数据集》），本文给出的命令与参数表都能直接落地。更进一步的推理玩法（如 Canny 边缘、深度图控制）可以继续参考仓库内StableDiffusionControlNetPipeline的配套文档。

【免费下载链接】diffusers🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers