
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 ModelsLvmin Zhang 与 Maneesh Agrawala它本质上是一类在预训练模型之上训练的适配器adapter。通过在 UNet 的 down-block 与 mid-block 注入额外的残差信号ControlNet 允许你用一张额外的输入图像来控制生成过程——这张条件图可以是 Canny 边缘图、深度图、人体姿态骨架也可以像本文示例一样是一张带有简单几何图形的合成图。本文的示例源自原始 ControlNet 仓库的训练文档并在 Diffusers 中以可复现的脚本形式落地使用小型合成数据集 fill50k训练一个圆圈填充 ControlNet——给定一张圆在背景中的条件图模型学会在指定位置、以指定颜色生成圆。基础模型选用Stable Diffusion 1.5stable-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.txtrequirements.txt 的内容如下accelerate0.16.0 torchvision transformers4.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 训练强相关的核心参数参数默认值说明--pretrained_model_name_or_path必填基础模型在 Hub 上的仓库 ID或本地模型权重目录的路径--controlnet_model_name_or_pathNone已有 ControlNet 权重的路径不指定时 ControlNet 会从 UNet 随机初始化ControlNetModel.from_unet(unet)--output_dircontrolnet-model模型预测与检查点输出目录--resolution512训练/验证图片统一缩放的目标分辨率必须能被 8 整除否则 VAE 与 ControlNet 编码器输出的 latent 尺寸不一致脚本会直接报错--train_batch_size4单设备训练批大小--gradient_accumulation_steps1反向传播前累积的更新步数等效扩大批大小突破显存限制--learning_rate5e-6初始学习率--lr_schedulerconstant可选linear、cosine、cosine_with_restarts、polynomial、constant、constant_with_warmup--lr_warmup_steps500学习率预热步数--dataset_nameNoneHub 数据集名称也可指向本地数据集目录与--train_data_dir二选一否则脚本报错--image_column/--conditioning_image_column/--caption_columnimage/conditioning_image/text数据集三列字段名自定义数据集时按需修改--validation_image/--validation_promptNone验证条件图路径列表与提示词列表两者要么数量一致要么其中一方为单值自动广播否则parse_args抛出 ValueError--validation_steps100每 N 步运行一次验证并记录生成图--num_validation_images4每个验证 prompt/条件图对生成的图片数--max_train_samplesNone截断训练样本数用于快速调试或流式加载超大数据集--checkpointing_steps500每 N 步保存一次训练状态检查点可用于--resume_from_checkpoint断点续训--gradient_checkpointing关闭以更慢的反向传播换取显存节省--use_8bit_adam关闭使用 bitsandbytes 的 8-bit AdamW 优化器--enable_xformers_memory_efficient_attention关闭使用 xFormers 内存高效注意力--set_grads_to_none关闭将梯度置None而非清零以省显存--mixed_precisionNoneno/fp16/bf16bf16 需 PyTorch ≥ 1.10 且 Nvidia Ampere GPU--report_totensorboard实验记录后端可选tensorboard、wandb、comet_ml、all--push_to_hub关闭训练结束后将模型推送到 Hub--allow_tf32关闭在 Ampere GPU 上允许 TF32 加速训练此外英文版指南还介绍了Min-SNR 加权论文Efficient Diffusion Training via Min-SNR Weighting Strategy技巧通过重新平衡损失来加速收敛推荐--snr_gamma5.0。需要说明的是该参数当前在 PyTorch 版 train_controlnet.py 中尚未实现但其完整实现可见于 Flax 版 train_controlnet_flax.pyL305 参数定义、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_DIRHub 模型 ID 或本地模型目录对应--pretrained_model_name_or_path与OUTPUT_DIR模型保存目录对应--output_dir然后启动训练export MODEL_DIRstable-diffusion-v1-5/stable-diffusion-v1-5 export OUTPUT_DIRpath to save model accelerate launch train_controlnet.py \ --pretrained_model_name_or_path$MODEL_DIR \ --output_dir$OUTPUT_DIR \ --dataset_namefusing/fill50k \ --resolution512 \ --learning_rate1e-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_size4 \ --push_to_hub注意该默认配置需要约 38GB 显存。训练脚本会在你的输出目录中生成并保存diffusion_pytorch_model.bin文件。默认情况下训练日志写入 TensorBoard想改用 Weights Biases只需在命令中追加--report_to wandb。用梯度累积把显存降到 ~20GB显存不足时可以缩小批大小并配合梯度累积等效总批大小不变1 × 4 4但峰值显存显著下降export MODEL_DIRstable-diffusion-v1-5/stable-diffusion-v1-5 export OUTPUT_DIRpath to save model accelerate launch train_controlnet.py \ --pretrained_model_name_or_path$MODEL_DIR \ --output_dir$OUTPUT_DIR \ --dataset_namefusing/fill50k \ --resolution512 \ --learning_rate1e-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_size1 \ --gradient_accumulation_steps4 \ --push_to_hub--gradient_accumulation_steps4意味着每 4 个 mini-batch 才执行一次参数更新等效批大小仍为 4。在 train_controlnet.py 的训练循环中accelerator.accumulate(controlnet)L1045正是负责这一累积语义的底层实现。多 GPU 分布式训练 Accelerate 对多 GPU 训练是无缝的。在accelerate launch时追加--multi_gpu与--mixed_precisionfp16即可训练脚本本身无需任何改动export MODEL_DIRstable-diffusion-v1-5/stable-diffusion-v1-5 export OUTPUT_DIRpath to save model accelerate launch --mixed_precisionfp16 --multi_gpu train_controlnet.py \ --pretrained_model_name_or_path$MODEL_DIR \ --output_dir$OUTPUT_DIR \ --dataset_namefusing/fill50k \ --resolution512 \ --learning_rate1e-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_size4 \ --mixed_precisionfp16 \ --tracker_project_namecontrolnet-demo \ --report_towandb \ --push_to_hub--tracker_project_name会作为Accelerator.init_trackers的project_name传入见 train_controlnet.py L991用于在 wandb/TensorBoard 中区分不同实验。低显存 GPU 优化方案官方按显存档位给出了三套经过验证的优化组合核心思想都是梯度检查点 低精度优化器 高效注意力 梯度置 None的组合拳。16GB GPU梯度检查点 8-bit Adam需要先安装 bitsandbytespip install bitsandbytes然后追加两个参数export MODEL_DIRstable-diffusion-v1-5/stable-diffusion-v1-5 export OUTPUT_DIRpath to save model accelerate launch train_controlnet.py \ --pretrained_model_name_or_path$MODEL_DIR \ --output_dir$OUTPUT_DIR \ --dataset_namefusing/fill50k \ --resolution512 \ --learning_rate1e-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_size1 \ --gradient_accumulation_steps4 \ --gradient_checkpointing \ --use_8bit_adam \ --push_to_hub在源码中--use_8bit_adam会把优化器切换为bnb.optim.AdamW8bittrain_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_noneexport MODEL_DIRstable-diffusion-v1-5/stable-diffusion-v1-5 export OUTPUT_DIRpath to save model accelerate launch train_controlnet.py \ --pretrained_model_name_or_path$MODEL_DIR \ --output_dir$OUTPUT_DIR \ --dataset_namefusing/fill50k \ --resolution512 \ --learning_rate1e-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_size1 \ --gradient_accumulation_steps4 \ --gradient_checkpointing \ --use_8bit_adam \ --enable_xformers_memory_efficient_attention \ --set_grads_to_none \ --push_to_hub使用前请务必安装 xFormerspip 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_noneTrue)传入训练循环L1103比置零更省显存但会改变部分行为若出现问题可关闭。8GB GPUDeepSpeed 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_DIRstable-diffusion-v1-5/stable-diffusion-v1-5 export OUTPUT_DIRpath to save model accelerate launch train_controlnet.py \ --pretrained_model_name_or_path$MODEL_DIR \ --output_dir$OUTPUT_DIR \ --dataset_namefusing/fill50k \ --resolution512 \ --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_size1 \ --gradient_accumulation_steps4 \ --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_datasetL602中目标图像使用缩放 中心裁剪 ToTensor Normalize([0.5],[0.5])的标准管线而条件图则使用不含 Normalize 的独立变换L688-L694conditioning_image_transforms transforms.Compose( [ transforms.Resize(args.resolution, interpolationtransforms.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-L815if 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.pyL445它会复制 UNet 的 encoder 结构并随机初始化条件控制分支——这正是 ControlNet 论文可复用的卷积块思想的体现冻结原 UNet训练一个结构相同但带条件输入的副本。冻结主干只训练 ControlNet。L854-L857 将 VAE、UNet、text encoder 全部requires_grad_(False)仅controlnet.train()优化器也只接收 ControlNet 参数L911-L918params_to_optimize controlnet.parameters() optimizer optimizer_class( params_to_optimize, lrargs.learning_rate, betas(args.adam_beta1, args.adam_beta2), weight_decayargs.adam_weight_decay, epsargs.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_statesencoder_hidden_states, controlnet_condcontrolnet_image, return_dictFalse, ) model_pred unet( noisy_latents, timesteps, encoder_hidden_statesencoder_hidden_states, down_block_additional_residuals[sample.to(dtypeweight_dtype) for sample in down_block_res_samples], mid_block_additional_residualmid_block_res_sample.to(dtypeweight_dtype), return_dictFalse, )[0]损失为目标噪声epsilon预测型或速度v_prediction预测型与预测值的 MSEL1089-L1095。训练中还通过accelerator.clip_grad_norm_做梯度裁剪L1099-L1100并支持--checkpointing_steps保存检查点、--resume_from_checkpoint断点续训L1008-L1032。条件编码网络的原生实现。在模型侧src/diffusers/models/controlnets/controlnet.py 的ControlNetConditioningEmbeddingL66-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, dtypetorch.float16) pipe StableDiffusionControlNetPipeline.from_pretrained( base_model_path, controlnetcontrolnet, dtypetorch.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_steps20, generatorgenerator, imagecontrol_image).images[0] image.save(./output.png)这段代码与训练脚本内部验证阶段log_validationL80-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.5examples/controlnet 目录还提供多种变体train_controlnet_sdxl.py为 Stable Diffusion XL双文本编码器架构训练 ControlNet详见 README_sdxl.mdtrain_controlnet_sd3.py 与 train_controlnet_flux.py分别面向 SD3 与 Flux 的 ControlNet 训练各有独立的 README 与 requirements 文件如 README_sd3.md、README_flux.mdtrain_controlnet_flax.py面向 TPU 的 JAX/Flax 版本支持流式数据集与 Min-SNR 加权但注意JAX/Flax 支持已在 diffusers v0.40.0 中移除运行该脚本需要diffusers0.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创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考