Consistency Models 深度解析：一步生成扩散模型

1
┌──────────────────────────────────────────────────────────────┐
2
│  扩散模型加速采样的三条路线                                   │
3
├──────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  路线 1: 更少的采样步数 (DDIM / PNDM)                      │
6
│  ────────────────────────────────────────────────────────   │
7
│  思想: 跳过部分中间步，通过 ODE 插值重建                     │
8
│  效果: 50-100 步仍可接受，20 步开始明显质量下降             │
9
│  缺点: 仍有固定步数，延迟有下限                              │
10
│                                                              │
11
│  路线 2: 预测更少的变量 (Latent Consistency Models / LCM)   │
12
│  ────────────────────────────────────────────────────────   │
13
│  思想: 预测 x_0 或一致性函数，减少步数                       │
14
│  效果: 2-8 步生成，质量接近全步                              │
15
│  缺点: 需要额外引导，对某些模型结构有依赖                    │
16
│                                                              │
17
│  路线 3: 一步直接生成 (Consistency Models)                 │
18
│  ────────────────────────────────────────────────────────   │
19
│  思想: 训练一个模型，直接从噪声一步映射到数据                 │
20
│  效果: 1 步生成，质量通过蒸馏提升                           │
21
│  缺点: 需要从预训练扩散模型蒸馏或从零训练                   │
22
│                                                              │
23
└──────────────────────────────────────────────────────────────┘

1
def diffusion_ode():
2
    """
3
    扩散模型的随机微分方程 (SDE) 视角。
4

5
    关键洞察:
6
    前向过程和反向过程都在同一条轨迹上移动。
7
    """
8

9
    return {
10
        "前向 SDE (添加噪声)": {
11
            "方程": "dx = σ_max^t · dW_t",
12
            "解释": " Wiener 过程 W_t 驱动 x 趋向噪声",
13
            "边界": "x_0 → x_T (T → ∞ 时，x_T ≈ N(0, σ_max²))",
14
        },
15
        "反向 ODE (生成采样)": {
16
            "方程": "dx = [x - ∇_x log p_θ(x, t)] dt",
17
            "解释": "得分函数 ∇ log p 指向数据密度梯度",
18
            "关键": "这是确定性路径，从噪声流向数据",
19
        },
20
        "轨迹一致性": {
21
            "观察": "反向 ODE 定义了一条从 x_T 到 x_0 的路径",
22
            "性质": "轨迹上任意点 x_t 都是同一个扩散过程的"快照"",
23
            "需求": "我们希望一个函数 f 满足: f(x_t) = x_0 对所有 t",
24
        },
25
    }
26

27

28
def consistency_mapping():
29
    """
30
    一致性映射。
31

32
    目标: 学习一个函数 f_θ
33
    使得: f_θ(x_t) = x_0 对所有 t ∈ [0, T]
34
    其中 x_0 和 x_t 是同一轨迹上的点。
35

36
    这意味着:
37
    - f_θ 沿着 ODE 轨迹是"一致的"
38
    - 轨迹上的所有点映射到同一点 (轨迹起点)
39
    - 一旦学会，从任意中间步 x_t 一步到 x_0
40
    """
41

42
    return {
43
        "数学定义": "f_θ(x_t) = x_0",
44
        "物理含义": "轨迹上任意时刻的状态都"记住"了起点",
45
        "推理优势": "从任意 x_t，直接一步到达 x_0",
46
        "训练挑战": "x_0 是未知的 (我们只有噪声样本)",
47
    }

1
def consistency_function_derivation():
2
    """
3
    一致性函数的数学推导。
4

5
    核心问题: 如何训练 f_θ 使得 f_θ(x_t) ≈ x_0？
6

7
    方法 1: 自监督学习
8
    ─────────────────────────────────────────────
9
    1. 采样 x_0 ~ p_data (真实数据)
10
    2. 采样 t_1, t_2 ~ Uniform(ε, T), 其中 t_1 < t_2
11
    3. 通过前向过程: x_{t_1} = α_{t_1} x_0 + σ_{t_1} ε_1
12
                      x_{t_2} = α_{t_2} x_0 + σ_{t_2} ε_2
13
    4. 损失: |f_θ(x_{t_1}) - x_{t_2}|²
14

15
    关键观察:
16
    - 如果 t_2 → 0，则 x_{t_2} ≈ x_0
17
    - f_θ(x_{t_1}) 应该接近 x_{t_2} (都接近 x_0)
18
    - 不需要知道真实的 x_0！
19

20
    方法 2: 蒸馏 (从预训练扩散模型)
21
    ─────────────────────────────────────────────
22
    1. 使用预训练的 DDPM 作为教师
23
    2. 采样 x_{t+ε}，用教师模型计算 x_0^teacher
24
    3. 损失: |f_θ(x_t) - x_0^teacher|²
25

26
    方法 3: TTCT (Two-Track Consistency Training)
27
    ─────────────────────────────────────────────
28
    结合方法 1 和 2，用更强的教师信号
29
    """
30

31
    return {
32
        "TTCT 损失": {
33
            "公式": "L = λ_1 · |f_θ(x_t) - f_φ(x_{t+ε})|² + λ_2 · |f_θ(x_t) - x_0^teacher|²",
34
            "第一项": "自监督一致性 (不依赖教师)",
35
            "第二项": "蒸馏一致性 (利用预训练知识)",
36
            "优点": "既能从零训练，又保留蒸馏的效率",
37
        },
38
    }

1
def edm_integration():
2
    """
3
    Consistency Models 与 EDM 的整合。
4

5
    EDM 的核心设计:
6
    - 使用网络输出 c_out = σ_data · F_θ(x, t)
7
    - 其中 F_θ 是 UNet，σ_data 是数据标准差
8
    - 采样使用 EDM 的 ODE 求解器
9

10
    Consistency Model 的 EDM 版本:
11
    - f_θ(x_t, t) 直接输出 x_0
12
    - 网络: F_θ(x, t) 输出 c_skip(t)·x + c_out(t)·F_θ(x, t)
13
    - 其中 c_skip 和 c_out 是时间依赖的跳跃连接权重
14

15
    优势:
16
    - 与 SD 的 EDM 实现兼容
17
    - 可以直接加载 SD 的预训练权重进行蒸馏
18
    """
19

20
    return {
21
        "网络输出": {
22
            "标准 EDM": "c_out = σ_data · F_θ(x, t)",
23
            "Consistency Model": "x_0 = c_skip(t)·x + c_out(t)·F_θ(x, t)",
24
        },
25
        "时间依赖权重": {
26
            "c_skip(t)": "控制输入 x 的残差比例",
27
            "c_out(t)": "控制网络输出的比例",
28
            "设计": "与 EDM 相同，通过 MLP 从 t 计算",
29
        },
30
    }

1
TTCT 训练流程:
2

3
  轨道 1: 自监督一致性
4
  ───────────────────────────────────────────────────────
5
  x_0 ~ p_data
6
      ↓
7
  t_1, t_2 ~ Uniform(ε, T), t_1 < t_2
8
      ↓
9
  x_{t_1} = α_{t_1} x_0 + σ_{t_1} ε_1
10
  x_{t_2} = α_{t_2} x_0 + σ_{t_2} ε_2
11
      ↓
12
  损失: L_1 = |f_θ(x_{t_1}) - x_{t_2}|²
13

14
  轨道 2: 蒸馏一致性 (从预训练教师)
15
  ───────────────────────────────────────────────────────
16
  x_ε ~ N(0, σ²)
17
      ↓
18
  用预训练教师模型计算 x_0^teacher
19
      ↓
20
  损失: L_2 = |f_θ(x_ε) - x_0^teacher|²
21

22
  总损失: L = λ_1 · L_1 + λ_2 · L_2

1
def ttct_loss():
2
    """
3
    TTCT 损失函数实现。
4
    """
5
    def compute_ttct_loss(student_model, teacher_model, p_data,
6
                          optimizer, config):
7
        """
8
        参数:
9
            student_model: 要训练的 Consistency Model
10
            teacher_model: 预训练的 DDPM (用于蒸馏)
11
            p_data: 真实数据分布
12
            config: 训练配置
13
        """
14
        # 采样一个 batch
15
        x_0 = sample_batch(p_data, batch_size=config.batch_size)
16

17
        # ============ 轨道 1: 自监督一致性 ============
18
        # 采样两个时间步
19
        epsilon = config.epsilon  # 最小时间步
20
        T = config.T             # 最大时间步
21

22
        t_1 = torch.rand(batch_size) * (T - epsilon) + epsilon
23
        t_2 = torch.rand(batch_size) * (T - t_1) + t_1  # t_2 > t_1
24

25
        # 加噪
26
        noise_1 = torch.randn_like(x_0)
27
        noise_2 = torch.randn_like(x_0)
28

29
        x_t1 = alpha(t_1).view(-1, 1, 1, 1) * x_0 + \
30
               sigma(t_1).view(-1, 1, 1, 1) * noise_1
31

32
        x_t2 = alpha(t_2).view(-1, 1, 1, 1) * x_0 + \
33
               sigma(t_2).view(-1, 1, 1, 1) * noise_2
34

35
        # 学生模型预测
36
        pred_t1 = student_model(x_t1, t_1)
37

38
        # 自监督损失: f_θ(x_{t_1}) 应该接近 x_{t_2}
39
        # (因为两者都接近 x_0)
40
        loss_1 = F.mse_loss(pred_t1, x_t2.detach())
41

42
        # ============ 轨道 2: 蒸馏一致性 ============
43
        # 采样噪声
44
        x_epsilon = torch.randn_like(x_0) * sigma(epsilon)
45

46
        # 教师模型预测 (全步 DDPM)
47
        with torch.no_grad():
48
            x_0_teacher = teacher_model(x_epsilon, epsilon)
49

50
        # 学生模型预测
51
        pred_eps = student_model(x_epsilon, epsilon)
52

53
        # 蒸馏损失
54
        loss_2 = F.mse_loss(pred_eps, x_0_teacher.detach())
55

56
        # ============ 总损失 ============
57
        loss = config.lambda_1 * loss_1 + config.lambda_2 * loss_2
58

59
        return loss, {"self_supervised": loss_1, "distillation": loss_2}
60

61
    return compute_ttct_loss

1
def training_modes():
2
    """
3
    Consistency Models 的两种训练模式。
4
    """
5
    return {
6
        "Training-from-Scratch": {
7
            "方法": "不使用预训练教师，只用自监督一致性",
8
            "损失": "L = |f_θ(x_{t_1}) - x_{t_2}|²",
9
            "优点": "不需要预训练扩散模型",
10
            "缺点": "需要更多数据和训练时间",
11
            "适用": "从头构建新模型",
12
            "收敛": "慢，但最终质量可达",
13
        },
14
        "SCT (Sampling Consistency Training)": {
15
            "方法": "从零训练 + 使用单步学生作为教师",
16
            "损失": "L = |f_θ(x_t) - f_φ(x_{t+Δ})|²",
17
            "优点": "渐进提升，不需要预训练教师",
18
            "缺点": "需要迭代训练教师-学生",
19
            "适用": "持续改进模型",
20
        },
21
        "CFG (Consistency Fine-tuning)": {
22
            "方法": "在预训练 CM 基础上微调，加入 CFG",
23
            "损失": "L = |f_θ(x_t, c) - x_0^teacher(x_t, c)|²",
24
            "优点": "继承预训练 CM + 支持条件生成",
25
            "缺点": "需要额外的条件编码器",
26
            "适用": "快速适配下游任务",
27
        },
28
        "CD (Consistency Distillation)": {
29
            "方法": "从预训练扩散模型蒸馏",
30
            "损失": "L = |f_θ(x_t) - x_0^teacher(x_t)|²",
31
            "优点": "快速收敛，质量高",
32
            "缺点": "需要预训练扩散模型",
33
            "适用": "已有 SD 模型，想加速采样",
34
        },
35
    }

1
class ConsistencyModelEDM(nn.Module):
2
    """
3
    基于 EDM UNet 的 Consistency Model。
4

5
    核心修改:
6
    - 网络输出改为预测 x_0 (而非噪声或 v-prediction)
7
    - 添加时间依赖的 c_skip 和 c_out 跳跃连接
8
    """
9

10
    def __init__(self, unet_backbone):
11
        super().__init__()
12
        self.unet = unet_backbone
13

14
        # 时间嵌入 → 跳跃连接权重
15
        self.c_skip_net = nn.Sequential(
16
            nn.Linear(256, 1),
17
            nn.Sigmoid(),
18
        )
19
        self.c_out_net = nn.Sequential(
20
            nn.Linear(256, 1),
21
            nn.Sigmoid(),
22
        )
23

24
    def forward(self, x, t):
25
        """
26
        参数:
27
            x: (B, C, H, W) 输入 (可以是任意时间步的 x_t)
28
            t: (B,) 时间步 (0 到 T)
29
        返回:
30
            x_0: (B, C, H, W) 预测的原始数据
31
        """
32
        # 时间嵌入
33
        t_emb = get_timestep_embedding(t, 128)
34
        t_emb = self.unet.time_embed(t_emb)  # (B, 256)
35

36
        # EDM UNet 前向
37
        h = self.unet.prep(x)
38
        h = self.unet.downsample(h)
39
        h = self.unet.attention(h, t_emb)
40
        h = self.unet.upsample(h)
41

42
        # EDM 输出: F_θ(x, t)
43
        f_theta = self.unet.final(h)  # (B, C, H, W)
44

45
        # 时间依赖跳跃连接
46
        c_skip = self.c_skip_net(t_emb)        # (B, 1) ∈ (0, 1)
47
        c_out = self.c_out_net(t_emb)          # (B, 1) ∈ (0, 1)
48

49
        # Consistency Model 输出: x_0 = c_skip · x + c_out · F_θ
50
        x_0 = c_skip.view(-1, 1, 1, 1) * x + \
51
              c_out.view(-1, 1, 1, 1) * f_theta
52

53
        return x_0

1
def boundary_conditions():
2
    """
3
    Consistency Models 的边界条件。
4

5
    t = 0 时:
6
    - x_t = x_0 (无噪声)
7
    - f_θ(x_0, 0) 应该 = x_0 (恒等映射)
8

9
    t = T 时:
10
    - x_T ≈ 纯噪声
11
    - f_θ(x_T, T) 应该 ≈ x_0 (仍然预测数据)
12
    """
13

14
    return {
15
        "t = 0 (边界)": {
16
            "条件": "f_θ(x, 0) = x",
17
            "实现": "强制 c_skip(0) = 1, c_out(0) = 0",
18
            "原因": "此时 x 就是真实数据，不需要预测",
19
        },
20
        "t = T (边界)": {
21
            "条件": "f_θ(x, T) = E[x_0 | x_T = x]",
22
            "实现": "网络在高噪声时预测数据分布的均值",
23
            "原因": "从纯噪声中重建数据",
24
        },
25
        "中间 t": {
26
            "条件": "f_θ(x_t, t) 应该介于 x_t 和 x_0 之间",
27
            "实现": "c_skip(t) 和 c_out(t) 通过 MLP 学习",
28
            "观察": "t 越大，c_skip 越小（x_t 越偏离 x_0）",
29
        },
30
    }

1
@torch.no_grad()
2
def cm_sample(model, x_T, num_steps=1, guide_scale=1.0):
3
    """
4
    Consistency Model 推理。
5

6
    参数:
7
        model: 训练好的 Consistency Model
8
        x_T: (B, C, H, W) 初始噪声
9
        num_steps: 采样步数 (1 = 一步生成)
10
        guide_scale: CFG 引导强度
11
    返回:
12
        x_0: (B, C, H, W) 生成的图像
13
    """
14

15
    if num_steps == 1:
16
        # ============ 一步生成 (核心) ============
17
        t = torch.full((x_T.shape[0],), 1.0, device=x_T.device)
18

19
        if guide_scale == 1.0:
20
            # 无条件生成
21
            x_0 = model(x_T, t)
22
        else:
23
            # CFG 引导
24
            # 需要条件和无条件输入
25
            x_0_cond = model(x_T, t, condition)      # 条件输入
26
            x_0_uncond = model(x_T, t, uncondition)  # 无条件输入
27
            x_0 = x_0_uncond + guide_scale * (x_0_cond - x_0_uncond)
28

29
        return x_0
30

31
    else:
32
        # ============ 多步生成 (TTA 增强) ============
33
        return cm_sample_tta(model, x_T, num_steps)
34

35

36
@torch.no_grad()
37
def cm_sample_tta(model, x_T, num_steps=4):
38
    """
39
    TTST (Test-Time Self-Taught) 多步增强。
40

41
    核心思想:
42
    - 多步生成比单步更稳定
43
    - 但步数仍然远少于 DDPM
44
    - 通过 self-distillation 持续改进
45
    """
46

47
    # 初始化
48
    x = x_T.clone()
49
    T = 1.0
50

51
    # 时间步调度 (从大到小)
52
    timesteps = torch.linspace(T, 0, num_steps + 1).tolist()
53

54
    for i in range(num_steps):
55
        t = timesteps[i]
56
        t_next = timesteps[i + 1]
57

58
        # 单步预测
59
        x_0_pred = model(x, torch.full((x.shape[0],), t, device=x.device))
60

61
        # 如果 t_next = 0，则这就是最终输出
62
        if t_next == 0:
63
            return x_0_pred
64

65
        # 否则，用预测的 x_0 重建下一步
66
        # x_{t_next} = α_{t_next} · x_0 + σ_{t_next} · noise
67
        # 其中 noise = (x - α_t · x_0) / σ_t
68
        alpha_t = alpha(torch.tensor(t))
69
        sigma_t = sigma(torch.tensor(t))
70
        alpha_next = alpha(torch.tensor(t_next))
71
        sigma_next = sigma(torch.tensor(t_next))
72

73
        noise = (x - alpha_t * x_0_pred) / sigma_t.clamp(min=1e-6)
74
        x = alpha_next * x_0_pred + sigma_next * noise
75

76
    return x

1
def ttst_explained():
2
    """
3
    TTST (Test-Time Self-Taught) 推理详解。
4

5
    目标: 用多步推理提升单步 CM 的质量。
6

7
    原理:
8
    1. 从噪声 x_T 开始
9
    2. 用 CM 预测 x_0
10
    3. 用预测的 x_0 重建中间步 x_{t-Δ}
11
    4. 再从 x_{t-Δ} 预测 x_0
12
    5. 重复 N 次，质量逐步提升
13

14
    为什么有效:
15
    - 每一步都在"纠正"上一步的误差
16
    - 与 ODE 求解器的思想类似
17
    - 但不需要梯度，计算量小
18
    """
19

20
    return {
21
        "单步 TTST": {
22
            "步骤": "x_T → f_θ(x_T) = x_0",
23
            "质量": "★★★☆☆",
24
            "速度": "★★★★★",
25
        },
26
        "4步 TTST": {
27
            "步骤": "x_T → x_{0.75} → x_{0.5} → x_{0.25} → x_0",
28
            "质量": "★★★★☆",
29
            "速度": "★★★★☆",
30
        },
31
        "8步 TTST": {
32
            "步骤": "更细粒度的多步重建",
33
            "质量": "★★★★★",
34
            "速度": "★★★☆☆",
35
        },
36
        "对比 DDPM": {
37
            "DDPM 100步": "质量 ★★★★★, 速度 ★☆☆☆☆",
38
            "CM 1步": "质量 ★★★☆☆, 速度 ★★★★★",
39
            "CM 4步 TTST": "质量 ★★★★☆, 速度 ★★★★☆",
40
        },
41
    }

1
def cfg_consistency_model():
2
    """
3
    CFG Consistency Model。
4

5
    方法: 微调预训练 CM，添加条件信息。
6
    损失: L = |f_θ(x_t, c) - x_0^teacher(x_t, c)|²
7
    """
8

9
    class CFGConsistencyModel(nn.Module):
10
        def __init__(self, base_cm, text_encoder):
11
            super().__init__()
12
            self.base_cm = base_cm
13
            self.text_encoder = text_encoder
14

15
            # 额外的条件 embedding 投影
16
            self.cfg_head = nn.Sequential(
17
                nn.Linear(768 + 768, 256),  # cond + uncond → time_embed dim
18
                nn.SiLU(),
19
                nn.Linear(256, 256),
20
            )
21

22
        def forward(self, x, t, cond_embed=None):
23
            """
24
            参数:
25
                x: (B, C, H, W) 输入
26
                t: (B,) 时间步
27
                cond_embed: (B, D) 条件 embedding (如文本)
28
            """
29
            # 文本编码
30
            if cond_embed is None:
31
                cond_embed = torch.zeros(x.shape[0], 768, device=x.device)
32

33
            # 无条件 (全零 embedding)
34
            uncond_embed = torch.zeros_like(cond_embed)
35

36
            # 条件预测
37
            x_cond = self.base_cm(x, t, cond_embed)
38

39
            # 无条件预测
40
            x_uncond = self.base_cm(x, t, uncond_embed)
41

42
            # CFG 合并
43
            cfg_scale = 7.5
44
            x_0 = x_uncond + cfg_scale * (x_cond - x_uncond)
45

46
            return x_0
47

48
    return CFGConsistencyModel
49

50

51
def cfg_training():
52
    """
53
    CFG Consistency Model 训练。
54
    """
55

56
    def train_cfg_cm(student_cm, teacher_cm, p_data, p_text, optimizer, config):
57
        """
58
        参数:
59
            student_cm: 要训练的 CFG CM
60
            teacher_cm: 预训练的 CFG 扩散模型
61
            p_data: 图像分布
62
            p_text: 文本分布
63
        """
64
        # 采样
65
        x_0 = sample_batch(p_data)
66
        text = sample_batch(p_text)
67

68
        # 文本编码
69
        cond_embed = student_cm.text_encoder(text)
70

71
        # 加噪
72
        t = torch.rand(x_0.shape[0]) * (config.T - config.epsilon) + config.epsilon
73
        noise = torch.randn_like(x_0)
74
        x_t = alpha(t).view(-1, 1, 1, 1) * x_0 + sigma(t).view(-1, 1, 1, 1) * noise
75

76
        # 学生预测
77
        pred = student_cm(x_t, t, cond_embed)
78

79
        # 教师预测 (使用相同的条件)
80
        with torch.no_grad():
81
            target = teacher_cm(x_t, t, cond_embed)
82

83
        # CFG 损失
84
        loss = F.mse_loss(pred, target)
85

86
        return loss

1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as F
4
import math
5
from typing import Optional
6

7

8
def alpha_sigma(t, sigma_data=0.5, p_mean=-1.2, p_std=1.2):
9
    """
10
    EDM 风格的 alpha 和 sigma 调度。
11
    """
12
    # t 是 [0, 1] 归一化的时间步
13
    # sigma(t) 使用指数调度
14
    sigma = sigma_data ** (1 + 0.5 * (t - 0.5))  # 非线性调度
15

16
    # alpha 从 t 计算
17
    # 使用 sigmoid 中心化
18
    c = -2.5 * (t - 0.5)
19
    alpha = 1 / (1 + math.e ** c)
20

21
    return alpha, sigma
22

23

24
class ConsistencyModel(nn.Module):
25
    """
26
    Consistency Model 完整实现。
27

28
    核心设计:
29
    - 输入任意时间步的 x_t，输出 x_0
30
    - 使用 EDM UNet 作为 backbone
31
    - 时间依赖的 c_skip 和 c_out 跳跃连接
32
    """
33

34
    def __init__(self, unet: nn.Module, sigma_data=0.5):
35
        super().__init__()
36
        self.unet = unet
37
        self.sigma_data = sigma_data
38

39
        # 时间嵌入维度 (与 UNet 匹配)
40
        time_embed_dim = 256
41

42
        # c_skip(t) 和 c_out(t) 的 MLP
43
        self.c_skip_net = nn.Sequential(
44
            nn.Linear(time_embed_dim, 256),
45
            nn.SiLU(),
46
            nn.Linear(256, 1),
47
        )
48
        self.c_out_net = nn.Sequential(
49
            nn.Linear(time_embed_dim, 256),
50
            nn.SiLU(),
51
            nn.Linear(256, 1),
52
        )
53

54
        # 时间步嵌入
55
        self.time_embed = nn.Sequential(
56
            nn.Linear(128, time_embed_dim),
57
            nn.SiLU(),
58
            nn.Linear(time_embed_dim, time_embed_dim),
59
        )
60

61
    def get_time_embedding(self, t):
62
        """获取时间步的正弦位置嵌入。"""
63
        half_dim = 64
64
        emb = math.log(10000) / (half_dim - 1)
65
        emb = torch.exp(torch.arange(half_dim, device=t.device) * -emb)
66
        emb = t[:, None] * emb[None, :]
67
        emb = torch.cat([emb.sin(), emb.cos()], dim=-1)
68
        return self.time_embed(emb)
69

70
    def forward(self, x, t):
71
        """
72
        前向传播。
73

74
        参数:
75
            x: (B, C, H, W) 输入 (任意时间步的 x_t)
76
            t: (B,) 时间步 (0 到 1)
77
        返回:
78
            x_0: (B, C, H, W) 预测的原始数据
79
        """
80
        # 时间嵌入
81
        t_emb = self.get_time_embedding(t)  # (B, 256)
82

83
        # EDM UNet 前向
84
        h = self.unet.prep(x)
85
        h = self.unet.downsample(h, t_emb)
86
        h = self.unet.attention(h, t_emb)
87
        h = self.unet.upsample(h, t_emb)
88

89
        # 网络输出 F_θ(x, t)
90
        f_theta = self.unet.final(h)
91

92
        # 时间依赖跳跃连接
93
        c_skip = self.c_skip_net(t_emb)  # (B, 1)
94
        c_out = self.c_out_net(t_emb)    # (B, 1)
95

96
        # 边界条件: t=0 时，c_skip=1, c_out=0
97
        # 边界条件: t=1 时，c_skip=0, c_out=1
98
        # 通过 Sigmoid 实现自动插值
99
        c_skip = torch.sigmoid(c_skip)
100
        c_out = torch.sigmoid(c_out)
101

102
        # Consistency Model 输出
103
        x_0 = c_skip.view(-1, 1, 1, 1) * x + c_out.view(-1, 1, 1, 1) * f_theta
104

105
        return x_0
106

107

108
class ConsistencyModelTrainer:
109
    """Consistency Model 训练器。"""
110

111
    def __init__(self, model: ConsistencyModel, teacher_model: Optional[nn.Module] = None,
112
                 device="cuda", lambda_1=1.0, lambda_2=1.0):
113
        self.model = model.to(device)
114
        self.teacher = teacher_model.to(device) if teacher_model else None
115
        self.device = device
116
        self.lambda_1 = lambda_1
117
        self.lambda_2 = lambda_2
118
        self.step = 0
119

120
    def train_step(self, x_0: torch.Tensor, cfg_emb: Optional[torch.Tensor] = None):
121
        """
122
        单步训练。
123

124
        参数:
125
            x_0: (B, C, H, W) 真实数据
126
            cfg_emb: (B, D) 条件 embedding (可选)
127
        """
128
        B = x_0.shape[0]
129

130
        # ============ 轨道 1: 自监督一致性 ============
131
        epsilon = 1e-3
132
        T = 1.0
133

134
        # 采样两个时间步 (t_1 < t_2)
135
        t_1 = torch.rand(B, device=self.device) * (T - epsilon) + epsilon
136
        t_2 = torch.rand(B, device=self.device) * (T - t_1) + t_1
137

138
        # 加噪
139
        noise_1 = torch.randn_like(x_0)
140
        noise_2 = torch.randn_like(x_0)
141

142
        alpha_1, sigma_1 = alpha_sigma(t_1)
143
        alpha_2, sigma_2 = alpha_sigma(t_2)
144

145
        x_t1 = alpha_1.view(-1, 1, 1, 1) * x_0 + sigma_1.view(-1, 1, 1, 1) * noise_1
146
        x_t2 = alpha_2.view(-1, 1, 1, 1) * x_0 + sigma_2.view(-1, 1, 1, 1) * noise_2
147

148
        # 学生预测
149
        pred_1 = self.model(x_t1, t_1)
150

151
        # 自监督损失: f_θ(x_{t_1}) ≈ x_{t_2}
152
        loss_1 = F.mse_loss(pred_1, x_t2.detach())
153

154
        # ============ 轨道 2: 蒸馏一致性 ============
155
        if self.teacher is not None:
156
            # 用教师模型计算目标
157
            t_eps = torch.full((B,), epsilon, device=self.device)
158
            alpha_eps, sigma_eps = alpha_sigma(t_eps)
159
            x_eps = alpha_eps * x_0 + sigma_eps * torch.randn_like(x_0)
160

161
            with torch.no_grad():
162
                # 教师预测 x_0
163
                target = self.teacher(x_eps, t_eps)
164

165
            # 学生预测
166
            pred_eps = self.model(x_eps, t_eps)
167

168
            # 蒸馏损失
169
            loss_2 = F.mse_loss(pred_eps, target.detach())
170

171
            total_loss = self.lambda_1 * loss_1 + self.lambda_2 * loss_2
172
        else:
173
            total_loss = loss_1
174
            loss_2 = torch.tensor(0.0)
175

176
        # 反向传播
177
        total_loss.backward()
178

179
        return {"total": total_loss.item(), "self_sup": loss_1.item(), "distill": loss_2.item()}
180

181
    @torch.no_grad()
182
    def sample(self, batch_size=1, steps=1, guide_scale=1.0):
183
        """
184
        推理采样。
185
        """
186
        # 初始化噪声
187
        x = torch.randn(batch_size, 3, 32, 32, device=self.device)
188

189
        if steps == 1:
190
            # 一步生成
191
            t = torch.ones(batch_size, device=self.device)
192
            x_0 = self.model(x, t)
193
        else:
194
            # TTST 多步生成
195
            timesteps = torch.linspace(1.0, 0.0, steps + 1, device=self.device)
196
            for i in range(steps):
197
                t = timesteps[i]
198
                t_next = timesteps[i + 1]
199

200
                pred = self.model(x, t.expand(batch_size))
201

202
                if t_next == 0:
203
                    return pred
204

205
                # 重建下一步
206
                alpha_t, sigma_t = alpha_sigma(t)
207
                alpha_next, sigma_next = alpha_sigma(t_next)
208

209
                noise = (x - alpha_t * pred) / sigma_t.clamp(min=1e-6)
210
                x = alpha_next * pred + sigma_next * noise
211

212
            x_0 = x
213

214
        return x_0

1
LCM vs 标准 CM:
2

3
  Consistency Model:
4
    - 在像素空间操作
5
    - 需要处理高分辨率 (512×512×3)
6
    - 计算量大
7

8
  LCM (Latent Consistency Model):
9
    - 在潜空间操作 (使用 SD 的 VAE)
10
    - 处理低维潜变量 (64×64×4)
11
    - 计算量大幅降低
12
    - 可以直接加载 SD 权重进行蒸馏
13

14
  核心公式:
15
    x_0 = c_skip(t) · x + c_out(t) · ε_θ(x, t)
16

17
  LCM 的特殊处理:
18
    - 使用 SD 的 UNet 作为 backbone
19
    - 直接蒸馏 ε_θ 的预测
20
    - 输出 x_0 而非噪声 ε

1
def lcm_training():
2
    """
3
    LCM 训练 (与 Consistency Model 相同，但操作在潜空间)。
4
    """
5

6
    def train_lcm(sd_unet, vae, text_encoder, dataloader, optimizer):
7
        """
8
        LCM 训练。
9
        """
10
        for batch in dataloader:
11
            images, texts = batch
12

13
            # VAE 编码到潜空间
14
            with torch.no_grad():
15
                z_0 = vae.encode(images)
16

17
            # 采样时间步
18
            t = torch.rand(B) * 0.998 + 0.002  # 避免 t=0, t=1
19

20
            # 加噪
21
            noise = torch.randn_like(z_0)
22
            z_t = alpha(t) * z_0 + sigma(t) * noise
23

24
            # 文本编码
25
            context = text_encoder(texts)
26

27
            # LCM 预测 x_0
28
            # 使用 SD UNet，但输出头替换为 CM 的 c_skip/c_out
29
            epsilon_pred = sd_unet(z_t, t, context)
30
            x_0_pred = (z_t - sigma(t) * epsilon_pred) / alpha(t)
31

32
            # 教师目标
33
            with torch.no_grad():
34
                target = x_0_pred  # 可以用更好的教师
35

36
            loss = F.mse_loss(x_0_pred, target)
37

38
            optimizer.zero_grad()
39
            loss.backward()
40
            optimizer.step()
41

42
    return train_lcm
43

44

45
@torch.no_grad()
46
def lcm_sample(lcm_model, text_encoder, prompt,
47
               num_steps=4, guidance_scale=7.0):
48
    """
49
    LCM 推理 — 4 步生成 SD 图像。
50
    """
51
    # 文本编码
52
    context = text_encoder([prompt])
53
    uncond = text_encoder([""])
54

55
    # 初始化噪声
56
    z = torch.randn(1, 4, 64, 64, device=device)
57

58
    # 调度
59
    timesteps = [0.9, 0.7, 0.4, 0.0]
60

61
    for i, t in enumerate(tqdm(timesteps)):
62
        t_tensor = torch.tensor([t], device=device)
63

64
        # CFG
65
        z_cond = lcm_model(z, t_tensor, context)
66
        z_uncond = lcm_model(z, t_tensor, uncond)
67
        z_0 = z_uncond + guidance_scale * (z_cond - z_uncond)
68

69
        if t == 0.0:
70
            break
71

72
        # 重建下一步
73
        alpha_t = alpha_of_t(t)
74
        alpha_next = alpha_of_t(timesteps[i + 1])
75
        sigma_t = sigma_of_t(t)
76
        sigma_next = sigma_of_t(timesteps[i + 1])
77

78
        noise = (z - alpha_t * z_0) / sigma_t.clamp(min=1e-6)
79
        z = alpha_next * z_0 + sigma_next * noise
80

81
    # VAE 解码
82
    image = vae.decode(z)
83

84
    return image

1
def comparison_all():
2
    """
3
    扩散模型加速方法的全面对比。
4
    """
5
    return {
6
        "DDIM": {
7
            "步数": "20-50 步",
8
            "质量": "★★★★☆ (20步) - ★★★★★ (50步)",
9
            "速度": "★★★☆☆",
10
            "方法": "确定性 ODE 采样",
11
            "优点": "无需训练，通用",
12
            "缺点": "仍有固定步数延迟",
13
        },
14
        "LCM": {
15
            "步数": "2-8 步",
16
            "质量": "★★★★☆",
17
            "速度": "★★★★★",
18
            "方法": "潜空间一致性模型",
19
            "优点": "步数极少，兼容 SD",
20
            "缺点": "需要蒸馏训练",
21
        },
22
        "Consistency Models": {
23
            "步数": "1 步 (可达 ★★★★☆)",
24
            "质量": "★★★★☆ (TTST) - ★★★☆☆ (单步)",
25
            "速度": "★★★★★",
26
            "方法": "ODE 轨迹一致性蒸馏",
27
            "优点": "一步生成，理论优美",
28
            "缺点": "单步质量有限，需 TTST 增强",
29
        },
30
        "LCM + TTST": {
31
            "步数": "4-8 步",
32
            "质量": "★★★★★",
33
            "速度": "★★★★☆",
34
            "方法": "LCM + 多步自改进",
35
            "优点": "质量和速度的最佳平衡",
36
            "缺点": "比单步慢",
37
        },
38
        "对抗蒸馏": {
39
            "步数": "1 步",
40
            "质量": "★★★★☆",
41
            "速度": "★★★★★",
42
            "方法": "GAN + 扩散模型蒸馏",
43
            "优点": "一步高质量",
44
            "缺点": "训练不稳定，mode collapse",
45
        },
46
    }