企业官网建设流程全解析

Diffusers ControlNet 训练实战指南：脚本参数、训练原理与多显存优化方案

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

ControlNet 是构建在预训练扩散模型之上的适配器结构，通过额外输入图像（如 Canny 边缘、深度图、人体姿态等）对生成过程进行精细条件控制。本文以 🤗 Diffusers 仓库中的 train_controlnet.py 为核心，完整讲解从环境准备、参数配置到训练循环源码级原理的实战流程，并针对 16GB / 12GB / 8GB 不同显存给出可落地的优化方案，读完即可上手训练并推理自己的 ControlNet 模型。

什么是 ControlNet：以额外图像条件约束生成

ControlNet 本质上是一个轻量化的适配器网络，训练时它被附加在另一个预训练模型（例如 Stable Diffusion 的 UNet）之上。其核心思想是：额外提供一张条件输入图（conditioning image），这张图可以是 Canny 边缘、深度图、人体姿态骨架、语义分割图等多种形式，模型在去噪过程中会同时参考文本提示与这张条件图，从而实现对生成结果的精细控制。

从源码结构看，Diffusers 将 ControlNet 实现为独立的模型类ControlNetModel，位于 src/diffusers/models/controlnets/controlnet.py。它的输出由ControlNetOutput数据类承载，包含两部分（见该文件 第 46-60 行）：

• down_block_res_samples：各下采样块的激活元组，用于条件化原 UNet 的下采样激活；
• mid_down_block_re_sample：中间块（最低采样分辨率）的激活，用于条件化原 UNet 的中间块激活。

这两部分输出会被注入 UNet 对应的残差位置，形成"条件控制"通路——这一调用关系我们会在后面的训练循环一节结合代码详细展开。

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

官方训练脚本要求从源码安装最新版 Diffusers（训练脚本中通过check_min_version("0.41.0.dev0")强制校验版本，见 train_controlnet.py，版本过旧会直接报错）：

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

然后进入示例目录并安装训练脚本所需的依赖：

cd examples/controlnet pip install -r requirements.txt

examples/controlnet/requirements.txt 中的依赖项如下：

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

其中accelerate用于多 GPU/TPU 与混合精度训练，datasets用于加载训练数据集，tensorboard用于训练日志与验证图像的可视化，transformers提供 CLIP 文本编码器与分词器。

使用 Accelerate 配置训练环境

🤗 Accelerate 会根据你的硬件和环境自动配置训练方案。初始化交互式配置：

accelerate config

如需跳过交互、直接使用默认配置：

accelerate config default

如果环境不支持交互式终端（例如在 Notebook 中），也可以用 Python API 写入基础配置：

from accelerate.utils import write_basic_config write_basic_config()

如果你的数据集格式与脚本默认不兼容，可参考仓库中的 Create a dataset for training 指南 学习如何制作适配训练脚本的数据集。

脚本参数详解：parse_args 全量解读

训练脚本的所有命令行参数都在parse_args()函数中定义，每个参数都带默认值与帮助说明。parse_args还内置了若干合法性校验，例如：

• --dataset_name与--train_data_dir必须至少指定一个（L570-L571）；
• --proportion_empty_prompts必须在[0, 1]区间内（L573-L574）；
• --validation_prompt与--validation_image必须成对出现（L576-L580）；
• --resolution必须能被 8 整除，以保证 VAE 与 ControlNet 编码后特征尺寸一致（L594-L597）。

例如，要使用 fp16 混合精度加速训练，只需在命令中追加：

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

基础与通用参数（如学习率、调度器、批次大小等）与 Text-to-image 训练指南 中描述的完全一致，这里只重点列出 ControlNet 相关的关键参数：

[!TIP] 在显存受限的 GPU 上训练时，建议优先开启--gradient_checkpointing、--gradient_accumulation_steps与--mixed_precision三个参数；还可以配合 xFormers 内存高效注意力进一步降低内存占用，详见 xFormers 优化指南。

Min-SNR 加权：加速收敛的损失重平衡

Min-SNR 加权策略通过对损失进行重新平衡来加快收敛速度。训练脚本支持预测epsilon（噪声）或v_prediction两种目标类型，而 Min-SNR 与这两种预测类型都兼容；需要说明的是，这一加权策略目前仅 PyTorch 后端支持。该策略在 examples/controlnet/README.md 中同样有明确记录。

启用方式是在训练命令中添加--snr_gamma参数，推荐值5.0：

accelerate launch train_controlnet.py \ --snr_gamma=5.0

训练脚本源码解析：数据、模型与训练循环

1. 数据预处理：条件图像同样需要变换

训练脚本通过make_train_dataset函数完成数据集的预处理，包括图像变换与标注分词。除了常规的分词和图像变换外，脚本还专门为条件图像定义了独立的变换管线（L688-L694）：

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

注意条件图像的变换没有Normalize([0.5], [0.5])这一步——这是有意的：目标图像image_transforms会做归一化以匹配 VAE 编码器输入分布，而条件图最终直接作为controlnet_cond传入 ControlNet，不需要归一化到相同的统计分布。

数据预处理还包含几个值得注意的细节：

• tokenize_captions（L660-L677）按--proportion_empty_prompts概率将提示词替换为空串，为后续可能的无分类器引导训练留出空间；标注既可以是单个字符串，也可以是字符串列表（此时随机取一条）。
• --max_train_samples会在主进程内先shuffle(seed=args.seed)再截断样本（L710-L711），保证每次截断取到的样本可复现。
• collate_fn（L718-L731）将pixel_values、conditioning_pixel_values、input_ids分别堆叠为连续内存格式的浮点张量，供 DataLoader 批量取用。

[!TIP] 如果要在 TPU 上流式加载数据集，性能可能受限于 🤗 Datasets 库——它对图像并不友好。为了保证最大吞吐，官方建议考虑 WebDataset、TorchData、TensorFlow Datasets 等其他数据集格式。

2. 模型加载：从已有权重还是从 UNet 初始化

在main()函数中，脚本依次加载分词器、文本编码器（AutoTokenizer+import_model_class_from_model_name_or_path动态识别CLIPTextModel或RobertaSeriesModelWithTransformation）、噪声调度器（DDPMScheduler）、VAE、UNet，然后加载 ControlNet：

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)

两种加载路径对应两类常见场景：

• 从已有权重加载（from_pretrained）：继续训练已训练的 ControlNet，或在其基础上微调；
• 从 UNet 随机初始化（from_unet）：从零训练一个全新条件类型的 ControlNet，from_unet类方法定义在 src/diffusers/models/controlnets/controlnet.py，它复用 UNet 的编码器结构作为 ControlNet 的骨干。

加载完成后，脚本将 VAE、UNet、文本编码器全部冻结（requires_grad_(False)），仅将 ControlNet 置于训练模式（controlnet.train()，见 L854-L857），这是适配器训练的典型策略：只更新适配器参数，基座模型保持不动。

3. 优化器：只更新 ControlNet 参数

优化器通过 get_scheduler 与 Adam 系列配合构建，其参数集合params_to_optimize明确限定为controlnet.parameters()（L910-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, )

optimizer_class根据--use_8bit_adam在torch.optim.AdamW与bnb.optim.AdamW8bit之间切换（L898-L908）。若指定--scale_lr，学习率还会乘以gradient_accumulation_steps × train_batch_size × num_processes进行等比例放大（L892-L895）。

4. 训练循环：ControlNet 输出如何注入 UNet

训练循环位于 train_controlnet.py 第 1043-1154 行，其核心逻辑分五步：

第一步：图像编码到潜空间并加噪。目标图像经 VAE 编码并乘以scaling_factor，再按随机时间步添加高斯噪声：

latents = vae.encode(batch["pixel_values"].to(dtype=weight_dtype)).latent_dist.sample() latents = latents * vae.config.scaling_factor noise = torch.randn_like(latents) timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device) noisy_latents = noise_scheduler.add_noise(latents.float(), noise.float(), timesteps).to(dtype=weight_dtype)

第二步：文本编码与条件图准备。提示词经文本编码器得到encoder_hidden_states，条件图经conditioning_image_transforms后作为controlnet_image：

encoder_hidden_states = text_encoder(batch["input_ids"], return_dict=False)[0] controlnet_image = batch["conditioning_pixel_values"].to(dtype=weight_dtype)

第三步：ControlNet 前向，产出残差样本。噪声潜变量、时间步、文本嵌入与条件图一同送入 ControlNet，得到下采样块与中间块的残差样本：

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

第四步：残差注入 UNet 完成预测。这些残差样本作为down_block_additional_residuals与mid_block_additional_residual注入冻结的 UNet，UNet 在原有文本条件之外"看到"了来自 ControlNet 的空间条件：

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]

第五步：按预测类型计算 MSE 损失。目标值根据调度器配置的prediction_type在噪声epsilon与速度v_prediction之间切换（get_velocity），然后计算 MSE 损失并反传：

if noise_scheduler.config.prediction_type == "epsilon": target = noise elif noise_scheduler.config.prediction_type == "v_prediction": target = noise_scheduler.get_velocity(latents, noise, timesteps) loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")

反传后若sync_gradients为真，则对 ControlNet 参数做梯度裁剪（max_grad_norm默认 1.0），随后optimizer.step()、lr_scheduler.step()并执行optimizer.zero_grad(set_to_none=args.set_grads_to_none)。循环内还按--checkpointing_steps保存 checkpoint（支持--resume_from_checkpoint续训），按--validation_steps调用log_validation将验证图像上报到 TensorBoard 或 wandb（L80-L186）。

若想深入理解去噪过程与管线/模型/调度器之间的关系，可阅读 Understanding pipelines, models and schedulers 教程。

启动训练：以 fill50k 数据集为例

本指南使用fusing/fill50k数据集（50k 张"圆/方/背景颜色"构成的填充图像对，非常适合快速验证 ControlNet 训练流程），当然你也可以按 Create a dataset for training 指南 制作并使用自己的数据集。

首先设置环境变量：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

启动训练前还需注意显存问题：脚本默认配置大约需要 38GB 显存。如果使用多 GPU，需要在accelerate launch命令中添加--multi_gpu参数。

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

上述命令中--train_batch_size=1配合--gradient_accumulation_steps=4实现等效 batch size 为 4 的训练，同时把单步显存需求降到最低；--validation_image与--validation_prompt各提供两张/两个，训练中每--validation_steps（默认 100）步就会生成并记录一次验证图像。

不同显存规模的优化配置

16GB 显存：8-bit 优化器 + 梯度检查点

在 16GB GPU 上，可以使用 bitsandbytes 8-bit 优化器配合梯度检查点。先安装 bitsandbytes：

pip install bitsandbytes

然后向训练命令追加以下参数：

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

12GB 显存：叠加 xFormers 与 grads-to-None

在 12GB GPU 上，需要在 16GB 方案基础上再叠加 xFormers 内存高效注意力，并将梯度置为None而非零值以进一步降低内存占用：

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

8GB 显存：DeepSpeed 阶段 2 + CPU 卸载

在 8GB GPU 上，需要用 DeepSpeed 将部分张量从显存卸载到 CPU 或 NVME。先运行accelerate config重新配置环境，在配置过程中确认使用DeepSpeed stage 2。组合 DeepSpeed stage 2、fp16 混合精度、将模型参数与优化器状态卸载到 CPU，即可在 8GB 以下显存训练（代价是需要约 25GB 系统内存）。配置文件大致如下：

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

进一步的配置选项可参考 Accelerate 的 DeepSpeed 使用指南。完成 DeepSpeed 配置后不需要在训练命令中追加任何额外参数。两个补充说明：

• 建议将默认 Adam 优化器替换为 DeepSpeed 优化的deepspeed.ops.adam.DeepSpeedCPUAdam以显著提速，但启用它要求系统 CUDA 工具链版本与 PyTorch 自带版本一致；
• bitsandbytes 8-bit 优化器目前与 DeepSpeed 不兼容。

推理验证：加载训练好的 ControlNet

训练完成后即可用训练好的 ControlNet 进行推理。将path/to/controlnet替换为--output_dir实际路径：

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

推理时把条件图（此处为验证用的conditioning_image_1.png）通过image=参数传入管线，模型就会在参考该条件结构的同时按新提示词生成结果。训练脚本在结束后还会自动生成模型卡片（含示例图），并在--push_to_hub开启时通过upload_folder将权重与模型卡一并推送到 Hub（见 L1178-L1190）。

训练 SDXL 版 ControlNet

Stable Diffusion XL（SDXL）是能生成高分辨率图像的强大文生图模型，其架构中增加了第二个文本编码器。如需为 SDXL 训练 ControlNet 适配器，使用仓库中的 train_controlnet_sdxl.py 脚本，其详细训练流程参见 SDXL 训练指南。此外，仓库还提供了针对 FLUX 的 train_controlnet_flux.py 与 train_control_flux.py 等脚本，可满足不同基座模型的训练需求。

下一步学习

训练好自己的 ControlNet 之后，可以进一步学习如何将其用于各类推理任务（边缘/深度/姿态条件生成、图像编辑等），参见 使用 ControlNet 进行推理。控制条件类型（canny、depth、pose 等）的选择、conditioning_scale的调节，都是在推理阶段进一步掌控生成结果的关键手段。

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