使用Diffusers训练Stable Diffusion 3/3.5的ControlNet模型指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00991/article/details/148360324

使用Diffusers训练Stable Diffusion 3/3.5的ControlNet模型指南

diffusers Diffusers：在PyTorch中用于图像和音频生成的最先进扩散模型。项目地址: https://gitcode.com/gh_mirrors/di/diffusers

前言

ControlNet是一种强大的神经网络结构，能够为扩散模型（如Stable Diffusion）提供额外的控制条件。本文将详细介绍如何使用Diffusers库为Stable Diffusion 3和3.5版本训练ControlNet模型。

环境准备

硬件要求

由于Stable Diffusion 3/3.5是大型模型，训练需要较高的GPU资源：

推荐使用至少80GB显存的GPU
单卡训练即可满足需求

软件依赖

需要安装以下依赖项：

# 安装Diffusers库及其训练依赖
pip install diffusers[training]

# 安装ControlNet训练特定依赖
pip install -r requirements_sd3.txt

# 配置加速环境
accelerate config

数据集准备

本教程使用"Circle filling"数据集，该数据集包含5万张图像，主要用于训练模型根据条件图像生成填充圆形。

数据集应放置在examples/controlnet/fill50k目录下。数据集结构如下：

包含原始图像
对应的条件图像（控制图）
文本描述

模型下载

由于Stable Diffusion 3/3.5是受控模型，使用前需要：

访问模型页面并填写申请表格
接受使用条款
使用Hugging Face CLI登录：

huggingface-cli login

训练流程

基本训练命令

export MODEL_DIR="stabilityai/stable-diffusion-3-medium-diffusers"
export OUTPUT_DIR="sd3-controlnet-out"

accelerate launch train_controlnet_sd3.py \
    --pretrained_model_name_or_path=$MODEL_DIR \
    --output_dir=$OUTPUT_DIR \
    --train_data_dir="fill50k" \
    --resolution=1024 \
    --learning_rate=1e-5 \
    --max_train_steps=15000 \
    --validation_image "./conditioning_image_1.png" "./conditioning_image_2.png" \
    --validation_prompt "red circle with blue background" "cyan circle with brown floral background" \
    --validation_steps=100 \
    --train_batch_size=1 \
    --gradient_accumulation_steps=4

关键参数说明

pretrained_model_name_or_path: 基础模型路径（SD3或SD3.5）
resolution: 训练分辨率，建议1024x1024
learning_rate: 学习率，1e-5是较好的起点
validation_*: 验证相关参数，用于训练过程监控

训练技巧

使用梯度累积（gradient_accumulation_steps）可以在有限显存下实现更大的batch size
启用torch编译模式可以显著加速训练
定期保存检查点以防训练中断

模型推理

训练完成后，可以使用以下代码进行推理：

from diffusers import StableDiffusion3ControlNetPipeline, SD3ControlNetModel
from diffusers.utils import load_image
import torch

# 加载基础模型和训练好的ControlNet
base_model_path = "stabilityai/stable-diffusion-3-medium-diffusers"
controlnet_path = "path/to/trained/controlnet"

controlnet = SD3ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16)
pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
    base_model_path, controlnet=controlnet
)
pipe.to("cuda", torch.float16)

# 准备输入
control_image = load_image("conditioning_image.png").resize((1024, 1024))
prompt = "描述你想要生成的图像"

# 生成图像
generator = torch.manual_seed(0)
image = pipe(
    prompt, 
    num_inference_steps=20, 
    generator=generator, 
    control_image=control_image
).images[0]
image.save("output.png")