ControlNet 深度剖析：条件控制扩散模型的结构化控制系统

1
方案 1: 在 UNet 中加入额外输入通道
2
  思想: 把条件图像拼接到潜变量通道
3
  问题: ❌ 需要重新训练整个 UNet，显存爆炸
4
  参数量: +86M 额外参数需要从头训练
5

6
方案 2: 只微调额外输入层
7
  思想: 新增输入分支，只训练新层
8
  问题: ❌ 新层参数随机初始化，微调效果差
9
  原因: 随机初始化的层在前向传播中产生噪声，破坏预训练 UNet 的特征
10

11
方案 3: 注意力机制
12
  思想: 用 cross-attention 注入条件
13
  问题: ❌ 文本注意力不够强，结构条件信息丢失
14
  原因: 文本 token 数量有限，无法精细编码空间信息
15

16
方案 4: 分类器引导
17
  思想: 训练一个条件分类器，引导生成
18
  问题: ❌ 需要额外的分类器，数据标注成本高
19
  原因: 分类器的质量直接限制生成效果

1
┌────────────────────────────────────────────────────────────────┐
2
│                                                                │
3
│   [条件图像 c] → 条件编码器                                     │
4
│      (Canny边缘/深度图/姿态)    ↓                              │
5
│                              可训练的副本分支                  │
6
│                              (ControlNet Block)               │
7
│                                        ↓                       │
8
│   [噪声潜变量 z_t] ───────────────────┐                        │
9
│           │                          ↓                         │
10
│   [时间步 t] ─→ UNet 编码器 ─→ [冻结 ❄️] ──────────────────┐ │
11
│                              ↓                               │ │
12
│                          UNet 解码器                         │ │
13
│                              ↑                               │ │
14
│                      跳跃连接 ← ← ← ← ← ← ← ← ← ← ← ← ← ← ┘ │
15
│                      (零卷积 ZeroConv)                        │
16
│                              ↑                                │
17
│                   ControlNet 副本编码器 ← 条件图像 c          │
18
│                   (训练 🔥)                                   │
19
│                                                                │
20
│   [文本条件] ──→ CLIP Text Encoder ──→ 交叉注意力注入        │
21
│                                                                │
22
└────────────────────────────────────────────────────────────────┘
23

24
关键设计:
25
  1. UNet 编码器完整复制 → 冻结 → 预训练权重
26
  2. 条件编码器: 与 UNet 编码器结构相同的副本 → 训练
27
  3. 零卷积: 跳跃连接处的 1×1 卷积，初始为零 → 渐进增长

1
class ControlNetBlock(nn.Module):
2
    """
3
    ControlNet 的基本模块。
4
    与 UNet 的 ResBlock 结构相同，但包含跳跃连接的零卷积。
5
    """
6
    def __init__(self, in_channels, out_channels, context_dim=None):
7
        super().__init__()
8

9
        # 主体: ResBlock
10
        self.norm1 = nn.GroupNorm(32, in_channels)
11
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
12

13
        if context_dim is not None:
14
            # 如果需要，添加交叉注意力（ControlNet 不需要，这里预留给扩展）
15
            self.has_attn = False  # ControlNet 主体不用 attention
16

17
        self.norm2 = nn.GroupNorm(32, out_channels)
18
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
19

20
        # 残差连接
21
        if in_channels != out_channels:
22
            self.shortcut = nn.Conv2d(in_channels, out_channels, 1)
23
        else:
24
            self.shortcut = nn.Identity()
25

26
        # ===== 关键: 跳跃连接的零卷积 =====
27
        # 初始为零，训练时渐进增长
28
        self.zero_conv = ZeroConv2d(in_channels, out_channels)
29

30

31
class ZeroConv2d(nn.Module):
32
    """
33
    零卷积层。
34
    权重初始化为零 → 前向传播输出为零 → 不破坏预训练网络。
35
    训练时通过梯度更新逐渐增长。
36
    """
37
    def __init__(self, in_channels, out_channels):
38
        super().__init__()
39
        self.conv = nn.Conv2d(in_channels, out_channels, 1, padding=0)
40

41
        # 关键: 权重初始化为零
42
        nn.init.constant_(self.conv.weight, 0.0)
43
        nn.init.constant_(self.conv.bias, 0.0)
44

45
    def forward(self, x):
46
        return self.conv(x)

1
def zero_conv_guarantee():
2
    """
3
    零卷积的数学保证。
4
    """
5
    analysis = {
6
        "初始化时": {
7
            "权重": "W = zero",
8
            "偏置": "b = zero",
9
            "输出": "y = W·x + b = zero",
10
            "效果": "ControlNet 对主网络的影响 = 0",
11
            "主网络行为": "完全等价于原始 SD，生成质量不变",
12
        },
13
        "训练过程中": {
14
            "梯度": "∂L/∂W ≠ 0 → W 开始更新",
15
            "更新方向": "由条件损失驱动，向条件方向调整",
16
            "增长速度": "初期慢（因为 W 接近零，梯度被放大）→ 后期稳定",
17
            "避免破坏": "W 从零增长，不会产生大的破坏性跳变",
18
        },
19
        "收敛时": {
20
            "权重": "W → W* (最优)",
21
            "效果": "ControlNet 学会了精确的条件控制",
22
            "主网络": "仍然冻结，预训练知识完整保留",
23
        },
24
    }
25
    return analysis

1
class ControlNet(nn.Module):
2
    """
3
    完整的 ControlNet。
4
    复制 SD 的 UNet 编码器为可训练的副本。
5
    """
6

7
    def __init__(self, sd_unet):
8
        super().__init__()
9
        self.sd_unet = sd_unet
10

11
        # 冻结 SD UNet
12
        for param in sd_unet.parameters():
13
            param.requires_grad_(False)
14

15
        # 复制编码器结构为 ControlNet
16
        self.control_scales = [1.0] * 12  # 12 层控制强度
17

18
        # 每一层对应一个可训练的 ControlNet Block
19
        self.control_blocks = nn.ModuleList([
20
            self._make_control_block(ch) for ch in sd_unet.channel_list
21
        ])
22

23
        # 零卷积跳跃连接
24
        self.zero_convs = nn.ModuleList([
25
            ZeroConv2d(sd_unet.channel_list[i], sd_unet.channel_list[i])
26
            for i in range(len(sd_unet.channel_list))
27
        ])
28

29
        # 额外的条件图像处理
30
        self.condition_encoder = nn.Sequential(
31
            nn.Conv2d(3, 16, 3, padding=1),  # 假设条件图为 RGB
32
            nn.SiLU(),
33
            nn.Conv2d(16, 32, 3, stride=2, padding=1),
34
            nn.SiLU(),
35
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
36
            nn.SiLU(),
37
            nn.Conv2d(64, 64, 3, padding=1),
38
        )
39

40
    def _make_control_block(self, channels):
41
        """创建一个 ControlNet Block。"""
42
        return nn.Sequential(
43
            ResBlock(channels, channels),
44
            ResBlock(channels, channels),
45
        )
46

47
    def forward(self, x_t, t, cond_image, context):
48
        """
49
        前向传播。
50

51
        参数:
52
            x_t:       (B, 4, H, W) 噪声潜变量
53
            t:         (B,) 时间步
54
            cond_image: (B, 3, H, W) 条件图像（边缘/深度等）
55
            context:   (B, seq_len, d) 文本 embedding
56
        """
57
        # 1) 编码条件图像
58
        cond = self.condition_encoder(cond_image)  # (B, 64, H/4, W/4)
59

60
        # 2) 获取 SD UNet 编码器特征
61
        #    (冻结的编码器用于提供特征，ControlNet 副本用于学习条件)
62
        features = self.sd_unet.get_encoder_features(x_t, t)
63

64
        # 3) 条件注入
65
        controlled_features = []
66
        for i, (feat, ctrl_block, zero_conv) in enumerate(zip(
67
            features, self.control_blocks, self.zero_convs
68
        )):
69
            # ControlNet 副本处理特征
70
            ctrl_feat = ctrl_block(feat + cond)  # 条件信息注入
71

72
            # 零卷积跳跃连接
73
            # 初始时 zero_conv 输出为 0 → 不影响主网络
74
            # 训练后: 输出控制信号
75
            skip = zero_conv(ctrl_feat)
76
            controlled_features.append(skip)
77

78
        # 4) SD UNet 解码器
79
        output = self.sd_unet.decode_with_skip(
80
            x_t, t, context, skip_connections=controlled_features
81
        )
82

83
        return output
84

85
    def unfix_layers(self):
86
        """
87
        解冻部分 SD UNet 层（可选的高级用法）。
88
        渐进解冻: 先训练零卷积，再解冻更多层。
89
        """
90
        pass  # ControlNet 主体设计为冻结 SD

1
┌──────────────────────────────────────────────────────────────┐
2
│  阶段 1: 固定阶段 (Locked Stage)                              │
3
│                                                              │
4
│  SD UNet:  冻结 ❄️ (权重不变)                              │
5
│  ZeroConv: 冻结 🔒 (梯度计算，但权重不更新)                  │
6
│                                                              │
7
│  效果:                                                       │
8
│  - 零卷积从零开始，输出 = 0                                   │
9
│  - 只有 ControlNet Block 更新                               │
10
│  - 解码器收到的跳跃连接 = 0                                  │
11
│  - 主网络完全保持预训练状态                                  │
12
│                                                              │
13
├──────────────────────────────────────────────────────────────┤
14
│  阶段 2: 解冻阶段 (Unlocked Stage)                          │
15
│                                                              │
16
│  ZeroConv: 解冻 🔥 (零卷积开始学习)                         │
17
│                                                              │
18
│  效果:                                                       │
19
│  - 零卷积从零开始增长                                        │
20
│  - ControlNet 的控制信号逐渐注入                             │
21
│  - 避免突然的大跳变                                          │
22
│                                                              │
23
│  为什么需要解冻 ZeroConv?                                    │
24
│  - ZeroConv 需要学习最优的控制强度缩放                       │
25
│  - 某些层可能需要更强的控制，某些层需要更弱                   │
26
│  - 解冻后 ZeroConv 可以独立调整每层的贡献                    │
27
└──────────────────────────────────────────────────────────────┘

1
def train_controlnet(controlnet, sd_unet, vae, text_encoder, dataloader,
2
                      optimizer, device, cfg_drop_prob=0.1):
3
    """
4
    ControlNet 训练循环。
5
    """
6
    controlnet.train()
7
    sd_unet.eval()  # 冻结
8

9
    for batch in tqdm(dataloader):
10
        images, cond_images, texts = (
11
            batch["image"].to(device),
12
            batch["condition"].to(device),  # Canny边缘/深度图/姿态
13
            batch["text"],
14
        )
15

16
        # ============ 1) 编码图像到潜空间 ============
17
        with torch.no_grad():
18
            z = vae.encode(images)  # (B, 4, 64, 64)
19

20
        # ============ 2) 采样时间步 ============
21
        batch_size = z.shape[0]
22
        t = torch.randint(0, 1000, (batch_size,), device=device)
23

24
        # ============ 3) 加噪 ============
25
        noise = torch.randn_like(z)
26
        alpha_bar = get_diffusion_schedule(t)  # 预设噪声调度
27
        z_t = alpha_bar.sqrt().view(-1, 1, 1, 1) * z + \
28
              (1 - alpha_bar).sqrt().view(-1, 1, 1, 1) * noise
29

30
        # ============ 4) 文本编码 + CFG Drop ============
31
        context = text_encoder.encode(texts)
32
        uncond_context = text_encoder.encode([""] * batch_size)
33

34
        # 随机丢弃条件（Classifier-Free Guidance）
35
        drop_mask = torch.rand(batch_size, device=device) < cfg_drop_prob
36
        context = context.masked_fill(drop_mask.unsqueeze(1), 0.0)
37

38
        # ============ 5) SD UNet 前向（冻结，提供特征） ============
39
        with torch.no_grad():
40
            # SD 的无条件预测
41
            noise_pred_uncond = sd_unet(z_t, t, uncond_context)
42

43
        # ============ 6) ControlNet 前向（训练） ============
44
        noise_pred_cond = controlnet(z_t, t, cond_images, context)
45

46
        # ============ 7) CFG 合并 ============
47
        # 训练时也用 CFG: 预测的噪声 = uncond + w * (cond - uncond)
48
        guidance_scale = 7.5
49
        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
50

51
        # ============ 8) 计算损失 ============
52
        loss = F.mse_loss(noise_pred, noise)
53

54
        # ============ 9) 反向传播 ============
55
        optimizer.zero_grad()
56
        loss.backward()
57
        torch.nn.utils.clip_grad_norm_(controlnet.parameters(), max_norm=1.0)
58
        optimizer.step()
59

60
    return loss.item()

1
class ConditionPreprocessor:
2
    """
3
    条件图像预处理。
4
    把各种类型的条件图像标准化到模型需要的格式。
5
    """
6
    @staticmethod
7
    def canny_edge(image, low_threshold=100, high_threshold=200):
8
        """
9
        Canny 边缘检测 → 二值边缘图。
10
        """
11
        import cv2
12
        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
13
        edges = cv2.Canny(gray, low_threshold, high_threshold)
14
        # 转为 RGB (3通道，与 SD 输入对齐)
15
        edges_rgb = np.stack([edges, edges, edges], axis=-1)
16
        return edges_rgb / 255.0
17

18
    @staticmethod
19
    def depth_map(image, model="MiDaS"):
20
        """
21
        单目深度估计 → 深度图。
22
        """
23
        # 使用 MiDaS 或 ZoeDepth 估计深度
24
        # 返回归一化的深度图
25
        pass
26

27
    @staticmethod
28
    def human_pose(image, model="OpenPose"):
29
        """
30
        人体姿态估计 → 关键点骨架图。
31
        """
32
        # 使用 OpenPose 或 DWPose 检测 17/25 个关键点
33
        # 绘制骨架连接
34
        pass
35

36
    @staticmethod
37
    def normal_map(image, model="MiDaS"):
38
        """
39
        表面法线估计 → 法线图。
40
        """
41
        # 估计表面法线，用于几何控制
42
        pass
43

44
    @staticmethod
45
    def semantic_seg(image, model="U2Net"):
46
        """
47
        语义分割 → 分割掩码图。
48
        """
49
        # 像素级类别分割
50
        pass
51

52

53
def normalize_condition(cond_image, target_size=(512, 512)):
54
    """
55
    标准化条件图像到 [0, 1] 范围，并 resize 到目标尺寸。
56
    """
57
    import torchvision.transforms as T
58

59
    transform = T.Compose([
60
        T.ToTensor(),
61
        T.Resize(target_size),
62
        T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
63
    ])
64

65
    return transform(cond_image)

1
def control_types():
2
    """
3
    ControlNet 支持的八种控制类型。
4
    """
5
    return {
6
        "Canny Edge": {
7
            "输入": "任意图像",
8
            "预处理": "Canny 边缘检测",
9
            "效果": "精确的边缘轮廓控制",
10
            "适用": "草图到图像、线稿上色",
11
            "预训练模型大小": "~1.4GB",
12
        },
13
        "Depth": {
14
            "输入": "任意图像",
15
            "预处理": "单目深度估计 (MiDaS)",
16
            "效果": "空间深度结构控制",
17
            "适用": "深度感知生成、3D 场景构建",
18
            "预训练模型大小": "~1.4GB",
19
        },
20
        "Normal Map": {
21
            "输入": "任意图像",
22
            "预处理": "表面法线估计",
23
            "效果": "物体表面几何控制",
24
            "适用": "纹理映射、几何细节",
25
            "预训练模型大小": "~1.4GB",
26
        },
27
        "OpenPose": {
28
            "输入": "人物图像",
29
            "预处理": "人体姿态估计 (17 关键点)",
30
            "效果": "人物姿态控制",
31
            "适用": "角色动作生成、舞蹈/运动",
32
            "预训练模型大小": "~1.4GB",
33
        },
34
        "Semantic Segmentation": {
35
            "输入": "任意图像",
36
            "预处理": "语义分割 (ADE20K)",
37
            "效果": "物体类别和位置控制",
38
            "适用": "场景布局、物体放置",
39
            "预训练模型大小": "~1.4GB",
40
        },
41
        "M-LSD": {
42
            "输入": "任意图像",
43
            "预处理": "直线检测",
44
            "效果": "建筑/室内线条结构控制",
45
            "适用": "建筑草图、室内设计",
46
            "预训练模型大小": "~1.4GB",
47
        },
48
        "HED Soft Edge": {
49
            "输入": "任意图像",
50
            "预处理": "HED 边缘检测 (软边缘)",
51
            "效果": "柔和的边缘轮廓，比 Canny 更自然",
52
            "适用": "艺术风格边缘、风格迁移",
53
            "预训练模型大小": "~1.4GB",
54
        },
55
        "Scribble": {
56
            "输入": "手绘涂鸦",
57
            "预处理": "涂鸦转二值图",
58
            "效果": "手绘草图控制",
59
            "适用": "从涂鸦生成完整图像",
60
            "预训练模型大小": "~1.4GB",
61
        },
62
    }

1
def canny_control_workflow():
2
    """
3
    Canny Edge 控制的工作流程。
4
    """
5
    return {
6
        "输入图像": "任意照片或生成图",
7
        "↓": "",
8
        "Canny 边缘检测": "低阈值 100, 高阈值 200",
9
        "↓": "得到精确的边缘轮廓图 (512×512×3)",
10
        "ControlNet (Canny)": "学习"边缘 → 图像内容"的映射",
11
        "↓": "输出条件噪声预测",
12
        "SD UNet + CFG": "结合文本条件生成图像",
13
        "↓": "",
14
        "输出图像": "保留原始边缘结构的生成图",
15
    }
16

17

18
def canny_training_data():
19
    """
20
    Canny ControlNet 的训练数据构建。
21
    """
22
    return {
23
        "数据来源": "LAION-5B / COCO / ADE20K",
24
        "样本数": "约 300 万对",
25
        "条件图像": "原始图像经过 Canny 边缘检测",
26
        "目标图像": "原始图像本身",
27
        "文本": "图像的描述 (BLIP-2 生成)",
28
        "训练技巧": "使用 ImageNet 预训练 Canny 检测器",
29
    }

1
OpenPose 骨架:
2
  ┌─────────────────────────────────────────────────┐
3
  │                                                 │
4
  │   [0] 鼻子                                      │
5
  │        │                                        │
6
  │   [1]─[2]─[3]─[4]─[15]─[14]─[16]─[17]        │
7
  │        │                                        │
8
  │   [5]  [6]  [7]  [8]                           │
9
  │        │                                        │
10
  │   [9]  [10] [11]                                │
11
  │        │                                        │
12
  │   [12] [13]                                     │
13
  │                                                 │
14
  │  17 个关键点:                                    │
15
  │  0: 鼻子   1-2: 眼睛   3-4: 耳朵               │
16
  │  5-8: 肩膀、手肘、手腕 (右)                     │
17
  │  9-12: 髋、膝盖、脚踝 (右)                      │
18
  │  13-16: 肩膀、手肘、手腕 (左)                   │
19
  │  17-20: 髋、膝盖、脚踝 (左)                     │
20
  │                                                 │
21
  └─────────────────────────────────────────────────┘
22

23
关键点连线 (骨骼):
24
  头部: 0-1, 0-2, 1-3, 2-4     (脸)
25
  躯干: 5-6, 6-7, 7-8          (右臂)
26
  躯干: 5-11, 11-12, 12-13     (右腿)
27
  躯干: 5-9, 9-10, 10-11       (左臂)
28
  躯干: 5-14, 14-15, 15-16     (左腿)

1
class OpenPoseVisualizer:
2
    """
3
    OpenPose 关键点可视化。
4
    """
5
    def __init__(self, keypoint_confidence=0.3):
6
        self.confidence = keypoint_confidence
7

8
        # 骨骼连接
9
        self.pairs = [
10
            (0, 1), (0, 2), (1, 3), (2, 4),  # 脸
11
            (5, 6), (6, 7), (7, 8),           # 右臂
12
            (5, 9), (9, 10), (10, 11),        # 左臂
13
            (5, 11), (11, 12), (12, 13),      # 右腿
14
            (5, 14), (14, 15), (15, 16),      # 左腿
15
        ]
16

17
        # 颜色
18
        self.keypoint_color = (255, 0, 85)    # 红色
19
        self.limb_color = (255, 85, 0)        # 橙色
20

21
    def draw_pose(self, keypoints, image_size):
22
        """
23
        在空白图上绘制姿态骨架。
24

25
        参数:
26
            keypoints: (17, 3) [x, y, confidence] 关键点
27
            image_size: (H, W) 输出图像大小
28
        """
29
        canvas = np.zeros((image_size[0], image_size[1], 3), dtype=np.uint8)
30

31
        # 绘制关键点
32
        for kp in keypoints:
33
            x, y, conf = kp
34
            if conf > self.confidence:
35
                cv2.circle(canvas, (int(x), int(y)), 4, self.keypoint_color, -1)
36

37
        # 绘制骨骼连线
38
        for i, j in self.pairs:
39
            if keypoints[i][2] > self.confidence and keypoints[j][2] > self.confidence:
40
                x1, y1 = int(keypoints[i][0]), int(keypoints[i][1])
41
                x2, y2 = int(keypoints[j][0]), int(keypoints[j][1])
42
                cv2.line(canvas, (x1, y1), (x2, y2), self.limb_color, 2)
43

44
        return canvas

1
def multi_control_inference(controlnet_models, sd_unet, vae, text_encoder,
2
                             z_t, t, conditions, texts, guidance_scale=7.5):
3
    """
4
    多控制组合推理。
5
    多个 ControlNet 可以组合使用。
6
    """
7
    noise_preds = []
8

9
    for ctrl_model, cond_img in zip(controlnet_models, conditions):
10
        # 每个 ControlNet 预测
11
        ctrl_pred = ctrl_model(z_t, t, cond_img, text_encoder.encode(texts))
12
        noise_preds.append(ctrl_pred)
13

14
    # 平均所有 ControlNet 的预测
15
    noise_pred_cond = torch.stack(noise_preds).mean(dim=0)
16

17
    # 无条件预测
18
    noise_pred_uncond = sd_unet(z_t, t, text_encoder.encode([""]))
19

20
    # CFG 合并
21
    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
22

23
    return noise_pred

1
@torch.no_grad()
2
def controlnet_inference(controlnet, sd_unet, vae, text_encoder,
3
                          prompt, cond_image, num_steps=50, guidance_scale=7.5):
4
    """
5
    ControlNet 推理。
6

7
    参数:
8
        prompt:       str 文本描述
9
        cond_image:   PIL.Image 条件图像（边缘/深度/姿态等）
10
        num_steps:    推理步数 (DDIM)
11
        guidance_scale: CFG 强度
12
    """
13
    device = next(controlnet.parameters()).device
14

15
    # ============ 1) 预处理条件图像 ============
16
    cond_tensor = preprocess_condition(cond_image)  # (1, 3, 512, 512)
17
    cond_tensor = cond_tensor.to(device)
18

19
    # ============ 2) 编码文本 ============
20
    context = text_encoder.encode([prompt])     # 条件文本
21
    uncond_context = text_encoder.encode([""])   # 无条件文本
22

23
    # ============ 3) 初始化噪声 ============
24
    z_T = torch.randn(1, 4, 64, 64, device=device)
25

26
    # ============ 4) DDIM 采样 ============
27
    timesteps = get_ddim_timesteps(num_steps)  # 递减时间步
28

29
    for i, t in enumerate(tqdm(timesteps)):
30
        t_tensor = torch.tensor([t], device=device)
31

32
        # 条件预测 (ControlNet)
33
        ctrl_pred = controlnet(z_T, t_tensor, cond_tensor, context)
34

35
        # 无条件预测 (SD)
36
        uncond_pred = sd_unet(z_T, t_tensor, uncond_context)
37

38
        # CFG 合并
39
        noise_pred = uncond_pred + guidance_scale * (ctrl_pred - uncond_pred)
40

41
        # DDIM 去噪步
42
        z_T = ddim_step(z_T, noise_pred, t, timesteps[i - 1] if i > 0 else 0)
43

44
    # ============ 5) VAE 解码 ============
45
    image = vae.decode(z_T)
46

47
    return image

1
def control_strength():
2
    """
3
    控制强度参数 (control_scale) 的效果。
4
    """
5
    results = {
6
        "0.0 (无控制)": {
7
            "效果": "完全由文本决定，忽略条件图像",
8
            "生成": "与普通 SD 无异",
9
        },
10
        "0.3-0.5 (弱控制)": {
11
            "效果": "条件提供参考，但允许灵活变化",
12
            "生成": "保留条件的部分结构，加入创意",
13
        },
14
        "0.7-0.9 (标准控制)": {
15
            "效果": "条件主导结构，细节和风格由文本决定",
16
            "生成": "精确的结构 + 丰富的细节",
17
        },
18
        "1.0 (强控制)": {
19
            "效果": "条件主导一切，文本仅影响颜色/风格",
20
            "生成": "结构与条件高度一致",
21
        },
22
        "1.2+ (过强)": {
23
            "效果": "可能产生 artifacts 或过度平滑",
24
            "生成": "结构过拟合，缺乏多样性",
25
        },
26
    }
27
    return results
28

29

30
def set_control_scale(controlnet, scale=0.8):
31
    """
32
    设置 ControlNet 的控制强度。
33
    推理时可以动态调整。
34
    """
35
    for module in controlnet.modules():
36
        if isinstance(module, ZeroConv2d):
37
            # ZeroConv 权重乘以 scale
38
            with torch.no_grad():
39
                module.conv.weight.mul_(scale)
40
    return controlnet

1
class ControlNetXS(nn.Module):
2
    """
3
    ControlNet-XS: 轻量级 ControlNet。
4
    只复制编码器的后半部分。
5
    """
6

7
    def __init__(self, sd_unet, num_layers_to_copy=4):
8
        super().__init__()
9
        self.sd_unet = sd_unet
10
        self.num_layers = num_layers_to_copy  # 只需复制 4 层而非 12 层
11

12
        # 只复制编码器的深层部分
13
        self.control_blocks = nn.ModuleList([
14
            self._make_control_block(ch)
15
            for ch in sd_unet.channel_list[-num_layers_to_copy:]
16
        ])
17

18
        # 轻量级条件编码器
19
        self.cond_encoder = nn.Sequential(
20
            nn.Conv2d(3, 32, 3, padding=1),
21
            nn.SiLU(),
22
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
23
            nn.SiLU(),
24
            nn.Conv2d(64, 64, 3, padding=1),
25
        )
26

27
        # 轻量级跳跃连接
28
        self.zero_convs = nn.ModuleList([
29
            ZeroConv2d(sd_unet.channel_list[-num_layers_to_copy + i],
30
                       sd_unet.channel_list[-num_layers_to_copy + i])
31
            for i in range(num_layers_to_copy)
32
        ])
33

34
    def forward(self, x_t, t, cond_image, context):
35
        # 条件编码
36
        cond = self.cond_encoder(cond_image)
37

38
        # 只在深层注入控制
39
        features = self.sd_unet.get_deep_features(x_t, t)  # 只取后 N 层
40

41
        controlled_skips = []
42
        for i, (feat, ctrl_block, zero_conv) in enumerate(zip(
43
            features, self.control_blocks, self.zero_convs
44
        )):
45
            ctrl_feat = ctrl_block(feat + cond)
46
            controlled_skips.append(zero_conv(ctrl_feat))
47

48
        # UNet 解码（与原始相同）
49
        output = self.sd_unet.decode_with_skip(x_t, t, context, controlled_skips)
50
        return output

1
def controlnet_comparison():
2
    """
3
    ControlNet vs ControlNet-XS 对比。
4
    """
5
    return {
6
        "参数量": {
7
            "ControlNet": "~860M (完整复制)",
8
            "ControlNet-XS": "~300M (部分复制)",
9
            "节省": "65%",
10
        },
11
        "推理速度": {
12
            "ControlNet": "1x (慢)",
13
            "ControlNet-XS": "~2x (快)",
14
        },
15
        "生成质量": {
16
            "ControlNet": "最佳",
17
            "ControlNet-XS": "略有下降，但在可接受范围",
18
            "差距": "FID 差值 < 2%",
19
        },
20
        "适用场景": {
21
            "ControlNet": "追求最高质量",
22
            "ControlNet-XS": "边缘部署、实时应用",
23
        },
24
    }

1
class ControlLoRA(nn.Module):
2
    """
3
    ControlLoRA: 用 LoRA 替代完整复制。
4
    """
5
    def __init__(self, sd_unet, rank=16, alpha=16):
6
        super().__init__()
7
        self.sd_unet = sd_unet
8
        self.rank = rank
9
        self.alpha = alpha
10

11
        # 为每层添加 LoRA
12
        self.lora_layers = nn.ModuleDict()
13
        for name, param in sd_unet.named_parameters():
14
            if "conv" in name and "weight" in name:
15
                # 对卷积层添加 LoRA
16
                self.lora_layers[name] = LoRALinear(param, rank, alpha)
17

18
        # 冻结原始权重
19
        for param in sd_unet.parameters():
20
            param.requires_grad_(False)
21

22
        # 条件处理
23
        self.cond_encoder = CondEncoder(rank)
24

25
    def forward(self, x_t, t, cond_image, context):
26
        # 处理条件
27
        cond = self.cond_encoder(cond_image)  # 低秩条件表示
28

29
        # 用 LoRA 注入条件
30
        for name, lora_layer in self.lora_layers.items():
31
            original = self.sd_unet.state_dict()[name]
32
            lora_weight = lora_layer(cond)  # 基于条件的低秩调整
33
            self.sd_unet.set_param(name, original + lora_weight)
34

35
        # 前向
36
        output = self.sd_unet(x_t, t, context)
37

38
        # 恢复原始权重
39
        for name in self.lora_layers:
40
            original = self.sd_unet.state_dict()[name]
41
            self.sd_unet.set_param(name, original)
42

43
        return output

1
def control_variants():
2
    """
3
    三种控制方案对比。
4
    """
5
    return {
6
        "ControlNet": {
7
            "参数量": "~860M",
8
            "训练方式": "全量微调",
9
            "质量": "最高",
10
            "速度": "慢",
11
            "适用": "追求最佳效果",
12
        },
13
        "ControlNet-XS": {
14
            "参数量": "~300M",
15
            "训练方式": "部分层微调",
16
            "质量": "略低 (~2% FID差距)",
17
            "速度": "中等",
18
            "适用": "效率敏感的部署",
19
        },
20
        "ControlLoRA": {
21
            "参数量": "~10-20M",
22
            "训练方式": "LoRA 微调",
23
            "质量": "低于 XS",
24
            "速度": "快",
25
            "适用": "轻量部署、个性化控制",
26
        },
27
    }

1
class T2IAdapter(nn.Module):
2
    """
3
    T2I-Adapter: 轻量级条件适配器。
4
    """
5
    def __init__(self, channels=[320, 640, 1280, 1280], num_blocks=4):
6
        super().__init__()
7

8
        # 条件图像编码器
9
        self.cond_encoder = nn.Sequential(
10
            nn.Conv2d(3, 64, 3, padding=1),
11
            *[ResBlock(64, 64) for _ in range(num_blocks)],
12
            nn.Conv2d(64, channels[-1], 1),
13
        )
14

15
        # 适配块: 与 UNet 各层对齐
16
        self.adapters = nn.ModuleList([
17
            AdapterBlock(in_ch, out_ch)
18
            for in_ch, out_ch in zip(channels[:-1], channels[1:])
19
        ])
20

21
    def forward(self, cond_image, unet_features):
22
        """
23
        注入条件信息到 UNet 特征。
24

25
        参数:
26
            cond_image:     (B, 3, H, W) 条件图像
27
            unet_features:   List[Tensor] UNet 各层特征
28
        """
29
        # 编码条件
30
        cond = self.cond_encoder(cond_image)
31

32
        # 逐层适配
33
        adapted_features = []
34
        for feat, adapter in zip(unet_features, self.adapters):
35
            adapted = adapter(feat, cond)
36
            adapted_features.append(adapted)
37

38
        return adapted_features

1
def controlnet_vs_adapter():
2
    """
3
    ControlNet vs T2I-Adapter 对比。
4
    """
5
    return {
6
        "架构": {
7
            "ControlNet": "完整复制 UNet 编码器",
8
            "T2I-Adapter": "独立轻量适配器",
9
        },
10
        "参数量": {
11
            "ControlNet": "~860M",
12
            "T2I-Adapter": "~80M",
13
        },
14
        "可组合性": {
15
            "ControlNet": "每个控制类型独立模型",
16
            "T2I-Adapter": "多个 Adapter 可叠加",
17
        },
18
        "质量": {
19
            "ControlNet": "更高（更多参数）",
20
            "T2I-Adapter": "略低，但足够好",
21
        },
22
        "灵活性": {
23
            "ControlNet": "固定与 SD 配合",
24
            "T2I-Adapter": "更灵活的注入点",
25
        },
26
    }