431 lines
10 KiB
Markdown
431 lines
10 KiB
Markdown
# 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）：**
|
||
```python
|
||
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编码器分支：**
|
||
```python
|
||
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的多个层级：
|
||
|
||
```python
|
||
# 注入点分布
|
||
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边缘检测 → 二值边缘图
|
||
|
||
```python
|
||
# 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点）
|
||
|
||
```python
|
||
# 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范围
|
||
|
||
```python
|
||
# 深度图预处理
|
||
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表面方向信息
|
||
|
||
```python
|
||
# 法线图生成
|
||
def compute_normal_map(image):
|
||
# 计算梯度
|
||
dx = sobel_x(image)
|
||
dy = sobel_y(image)
|
||
|
||
# 计算法线
|
||
normal = normalize(cross(dx, dy))
|
||
return normal
|
||
```
|
||
|
||
### 3.5 语义分割（Segmentation）
|
||
|
||
**输入：** 语义分割图（ADE20K等格式）
|
||
|
||
```python
|
||
# 分割图预处理
|
||
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）
|
||
|
||
**输入：** 手工绘制的线稿
|
||
|
||
```python
|
||
# 线稿提取
|
||
def preprocess_scribble(image):
|
||
# 简单的边缘检测
|
||
edges = edge_detect(image)
|
||
|
||
# 二值化
|
||
binary = threshold(edges, 128)
|
||
return binary
|
||
```
|
||
|
||
### 3.7 霍夫线检测（Hough Line）
|
||
|
||
**输入：** 霍夫变换检测出的线条
|
||
|
||
```python
|
||
# 霍夫线检测
|
||
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 训练数据准备
|
||
|
||
```python
|
||
# 数据格式
|
||
dataset = {
|
||
'condition': canny_edge_image, # 条件图
|
||
'prompt': "a photo of a person", # 文本描述
|
||
'gt_image': original_image, # 目标图像
|
||
' timestep': sampled_timestep
|
||
}
|
||
```
|
||
|
||
### 4.2 损失函数
|
||
|
||
```python
|
||
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 训练配置
|
||
|
||
```python
|
||
# 训练超参数
|
||
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 基础用法
|
||
|
||
```python
|
||
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组合
|
||
|
||
```python
|
||
# 多条件控制
|
||
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 建筑与室内设计
|
||
|
||
```python
|
||
# 使用Canny控制建筑外观
|
||
architectural_sketch = extract_canny(blueprint_image)
|
||
result = pipe(
|
||
prompt="modern architectural building, photo realistic",
|
||
image=architectural_sketch
|
||
)
|
||
```
|
||
|
||
### 6.2 人物动作控制
|
||
|
||
```python
|
||
# 使用姿态控制人物动作
|
||
pose_keypoints = extract_pose(person_image)
|
||
result = pipe(
|
||
prompt="person dancing in a club",
|
||
image=pose_keypoints
|
||
)
|
||
```
|
||
|
||
### 6.3 风格迁移
|
||
|
||
```python
|
||
# 使用深度图保持结构
|
||
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将扩散模型的生成能力从"模糊描述"提升到"精确控制"，是文生图领域的重要里程碑。 |