Notes/机器学习/深度学习模型/06-多模态模型/6-文生图模型/2-ControlNet.md

ControlNet

概述

ControlNet是斯坦福大学提出的条件控制扩散模型框架，通过添加额外条件输入（如边缘检测图、姿态关键点等）来控制图像生成过程。ControlNet的核心创新在于其并行分支结构和零卷积技术，使得精确控制成为可能，同时保持原有模型的能力。

1. ControlNet核心原理

1.1 问题背景

扩散模型如Stable Diffusion虽然能生成高质量图像，但缺乏对生成过程的精确控制能力。用户只能通过文本描述来间接影响生成结果，无法指定具体的构图、姿态或边缘结构。

1.2 解决方案

ControlNet通过添加并行控制分支来解决这个问题：

原始SD: text → UNet → image
ControlNet: text + condition → UNet(copied) + SD → image
                               ↑
                        零卷积连接

1.3 关键创新

零初始化卷积（Zero Convolution）：

class ZeroConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, 1)
        # 初始化为零
        nn.init.zeros_(self.conv.weight)
        nn.init.zeros_(self.conv.bias)
    
    def forward(self, x):
        return self.conv(x)

零卷积的作用：

• 训练初期：控制分支输出为0，等同于原SD
• 训练过程中：逐渐学习条件信息
• 训练完成：完整条件控制

2. 并行分支结构

2.1 架构设计

输入图像/条件图
       ↓
┌─────────────────────────────────────┐
│  ControlNet分支 (复制UNet编码器)      │
│  输入: 条件图 + 时间步               │
│  输出: 中间特征                      │
└─────────────────────────────────────┘
       ↓ 零卷积连接
┌─────────────────────────────────────┐
│  Stable Diffusion UNet              │
│  输入: 噪声latent + 文本embedding    │
│  输出: 去噪latent                    │
└─────────────────────────────────────┘

2.2 分支详情

ControlNet编码器分支：

class ControlNetBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        self.norm1 = nn.GroupNorm(32, in_channels)
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.norm2 = nn.GroupNorm(32, out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
        
        # 零卷积用于残差连接
        self.zero_conv = ZeroConv(in_channels, out_channels)
    
    def forward(self, x, condition):
        h = self.norm1(x)
        h = SiLU(h)
        h = self.conv1(h)
        
        h = self.norm2(h)
        h = SiLU(h)
        h = self.conv2(h)
        
        # 残差连接 + 条件注入
        return x + self.zero_conv(condition) + h

2.3 条件注入时机

ControlNet将条件信息注入到UNet的多个层级：

# 注入点分布
control_scales = {
    'encoder_mid': 1.0,    # 中间层最强
    'decoder_mid': 1.0,
    'up_blocks': [1.0, 1.0, 1.0],
    'down_blocks': [1.0, 1.0, 1.0, 0.0]  # 最底层不注入
}

# 多尺度控制
for i, down_block in enumerate(down_blocks):
    condition = extract_condition(condition_input, scale=control_scales[i])
    output = down_block(input, condition)

3. 八种控制类型

3.1 Canny边缘检测

输入： RGB图像 处理： Canny边缘检测 → 二值边缘图

# Canny边缘检测流程
def preprocess_canny(image):
    # 1. 灰度化
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    
    # 2. 高斯模糊
    blur = cv2.GaussianBlur(gray, (5, 5), 0)
    
    # 3. Canny边缘检测
    edges = cv2.Canny(blur, 100, 200)
    
    # 4. 边缘图作为条件
    return edges

应用场景： 精确边缘控制、线稿上色、建筑还原

3.2 姿态关键点（Pose）

输入： 人体姿态关键点（18点）

# OpenPose关键点定义
keypoints = {
    0: 'nose',
    1: 'left_eye',
    2: 'right_eye',
    3: 'left_ear',
    4: 'right_ear',
    5: 'left_shoulder',
    6: 'right_shoulder',
    7: 'left_elbow',
    8: 'right_elbow',
    9: 'left_wrist',
    10: 'right_wrist',
    11: 'left_hip',
    12: 'right_hip',
    13: 'left_knee',
    14: 'right_knee',
    15: 'left_ankle',
    16: 'right_ankle',
    17: 'neck'
}

应用场景： 人物动作控制、舞蹈动作生成

3.3 深度图（Depth）

输入： 单目深度估计图 处理： 归一化到0-1范围

# 深度图预处理
def preprocess_depth(image):
    # 使用MiDaS等模型估计深度
    depth = MiDaS.predict(image)
    
    # 归一化
    depth_normalized = (depth - depth.min()) / (depth.max() - depth.min())
    
    # 下采样到目标尺寸
    depth_small = resize(depth_normalized, (H//8, W//8))
    return depth_small

应用场景： 3D场景生成、深度感知合成

3.4 法线图（Normal）

输入： 表面法线图 用途： 提供3D表面方向信息

# 法线图生成
def compute_normal_map(image):
    # 计算梯度
    dx = sobel_x(image)
    dy = sobel_y(image)
    
    # 计算法线
    normal = normalize(cross(dx, dy))
    return normal

3.5 语义分割（Segmentation）

输入： 语义分割图（ADE20K等格式）

# 分割图预处理
def preprocess_segmentation(seg_map):
    # 类别数对应通道
    # 20类分割 → 20通道one-hot
    one_hot = F.one_hot(seg_map, num_classes=20)
    return one_hot.permute(2, 0, 1).float()

应用场景： 场景布局控制、物体放置

3.6 线稿图（Scribble）

输入： 手工绘制的线稿

# 线稿提取
def preprocess_scribble(image):
    # 简单的边缘检测
    edges = edge_detect(image)
    
    # 二值化
    binary = threshold(edges, 128)
    return binary

3.7 霍夫线检测（Hough Line）

输入： 霍夫变换检测出的线条

# 霍夫线检测
def detect_hough_lines(image):
    # 边缘检测
    edges = canny(image)
    
    # 霍夫变换
    lines = cv2.HoughLinesP(edges, 1, np.pi/180, 50, minLineLength=50)
    
    # 绘制线条
    line_image = draw_lines(lines)
    return line_image

3.8 阈值图（Threshold/Blur）

输入： 简单的二值或模糊图像

应用场景： 软控制、区域指定

4. 训练策略

4.1 训练数据准备

# 数据格式
dataset = {
    'condition': canny_edge_image,  # 条件图
    'prompt': "a photo of a person",  # 文本描述
    'gt_image': original_image,  # 目标图像
    ' timestep': sampled_timestep
}

4.2 损失函数

def compute_loss(model, condition, prompt, gt_image, timestep):
    # 1. 添加噪声
    noise = torch.randn_like(gt_latent)
    noisy_latent = add_noise(gt_latent, timestep, noise)
    
    # 2. 预测噪声
    pred_noise = model(
        noisy_latent,
        timestep,
        encoder_hidden_states=text_embed,
        controlnet_output=controlnet_features(condition)
    )
    
    # 3. 计算MSE损失
    loss = F.mse_loss(pred_noise, noise)
    return loss

4.3 训练配置

# 训练超参数
config = {
    'learning_rate': 1e-4,
    'batch_size': 16,
    'gradient_accumulation': 2,
    'max_steps': 100000,
    'warmup_steps': 1000,
    'ema_decay': 0.9999,
    'precision': 'fp16',
    'control_scale': 1.0  # 控制强度
}

5. 推理使用

5.1 基础用法

from controlnet import ControlNetModel
from diffusers import StableDiffusionControlNetPipeline

# 加载ControlNet
controlnet = ControlNetModel.from_pretrained(
    "lllyasviel/sd-controlnet-canny",
    torch_dtype=torch.float16
)

# 创建Pipeline
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    controlnet=controlnet,
    torch_dtype=torch.float16
)

# 生成图像
image = pipe(
    prompt="a person sitting on a chair",
    image=canny_edge_image,  # 条件图
    num_inference_steps=20,
    guidance_scale=7.5
).images[0]

5.2 多ControlNet组合

# 多条件控制
combined_controlnet = MultiControlNet([canny, pose, depth])

# 组合控制强度
output = pipe(
    prompt=prompt,
    image=[canny_img, pose_img, depth_img],
    control_weight=[0.5, 0.8, 0.3]  # 各条件权重
)

6. 应用场景

6.1 建筑与室内设计

# 使用Canny控制建筑外观
architectural_sketch = extract_canny(blueprint_image)
result = pipe(
    prompt="modern architectural building, photo realistic",
    image=architectural_sketch
)

6.2 人物动作控制

# 使用姿态控制人物动作
pose_keypoints = extract_pose(person_image)
result = pipe(
    prompt="person dancing in a club",
    image=pose_keypoints
)

6.3 风格迁移

# 使用深度图保持结构
depth_map = estimate_depth(content_image)
styled = pipe(
    prompt="impressionist painting style",
    image=depth_map,
    control_strength=0.8
)

7. 性能与限制

7.1 计算开销

配置	显存需求	推理时间
SD 1.5 + ControlNet	~8GB	~15s
SDXL + ControlNet	~16GB	~30s
多ControlNet	+4GB/Condition	+50%

7.2 当前限制

1. 条件图质量依赖： 边缘检测、姿态估计的准确性影响结果
2. 分辨率限制： 高分辨率控制需要更多显存
3. 风格漂移： 强条件可能改变原风格
4. 组合难度： 多条件组合需要调参

8. 变体与扩展

8.1 ControlNet-Tile

支持任意尺寸的图像控制，避免分辨率限制

8.2 ControlNet-Inpaint

结合修复功能的ControlNet变体

8.3 T2I-Adapter

轻量级条件控制方案，参数更少

9. 总结

ControlNet的核心价值：

创新	影响
零卷积	实现稳定训练，保留原模型能力
并行分支	灵活添加多种条件
多控制类型	覆盖主流控制需求
开源模型	推动社区发展

ControlNet将扩散模型的生成能力从"模糊描述"提升到"精确控制"，是文生图领域的重要里程碑。