视频生成模型深度解析：从 Video DiT 到 Sora 架构

1
图像生成 vs 视频生成:
2

3
  图像生成:
4
    输入: 文本描述
5
    输出: 单张 512×512 图像
6
    空间维度: H × W × C
7
    计算量: O(H × W)
8

9
  视频生成:
10
    输入: 文本描述 (+ 可选首帧/姿态)
11
    输出: T 帧视频，每帧 H × W
12
    空间维度: T × H × W × C
13
    时间维度: 连续的 T 帧必须保持:
14
      → 物体外观一致（同一角色每帧看起来相同）
15
      → 运动连贯（物体移动轨迹平滑）
16
      → 物理合理（重力、光照、遮挡符合物理规律）
17
      → 时序自然（动作、表情变化流畅）
18
    计算量: O(T × H × W) — 通常是图像的 10-100 倍

1
┌──────────────────────────────────────────────────────────────┐
2
│  视频生成的四大挑战                                          │
3
├──────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  1. 时序一致性 (Temporal Consistency)                       │
6
│     ─────────────────────────────────────────────────────  │
7
│     问题: 第 1 帧的狗和第 30 帧的狗要看起来是同一只狗      │
8
│     原因: 扩散模型逐帧独立生成，缺乏全局身份追踪              │
9
│     解法: 全视频序列联合建模（3D Attention / 视频 VAE）      │
10
│                                                              │
11
│  2. 长程依赖 (Long-range Dependency)                        │
12
│     ─────────────────────────────────────────────────────  │
13
│     问题: 开场动作要影响结尾的响应                           │
14
│     原因: 帧间距离越大，相关性越弱                          │
15
│     解法: Transformer 的全局 Attention（比 RNN 更适合）       │
16
│                                                              │
17
│  3. 计算可扩展性 (Computational Scalability)               │
18
│     ─────────────────────────────────────────────────────  │
19
│     问题: 16 帧 1024×1024 = 1670 万像素                    │
20
│     原因: 直接在像素空间扩散计算量爆炸                       │
21
│     解法: 潜空间视频扩散 (Latent Video Diffusion)          │
22
│                                                              │
23
│  4. 运动建模 (Motion Modeling)                             │
24
│     ─────────────────────────────────────────────────────  │
25
│     问题: 生成平滑、合理的运动轨迹                           │
26
│     原因: 运动是高度结构化的，难以从静态数据中学习           │
27
│     解法: 光流监督 / 运动模块 / 物理先验                    │
28
│                                                              │
29
└──────────────────────────────────────────────────────────────┘

1
2019: DDPM 提出 (Ho et al.) — 图像扩散基础奠定
2
      │
3
2020: DDPM 扩展到视频 (Ho et al.) — 首个视频扩散
4
      │    缺点: 逐帧独立生成，时序一致性差
5
      │
6
2022: LDM / Stable Diffusion — 潜空间扩散图像
7
      │    缺点: 只支持图像
8
      │
9
2022: VideoLDM (Blattmann et al.) — 首个 Latent Video Diffusion
10
      │    引入视频 VAE，扩展 SD 到视频
11
      │
12
2023: Make-A-Video (Singer et al., Meta) — 无需文本视频的预训练
13
      │    分离空间（图像扩散）+ 时间（插帧）
14
      │
15
2023: Align Your Latents (Zhao et al.) — 高质量视频潜空间
16
      │
17
2024: DiT (Peebles & Xie) — Transformer 替代 UNet
18
      │
19
2024: Sora (OpenAI, Feb) — Video DiT，DiT 的视频扩展
20
      │    关键技术: 时空 patchify + 3D VAE + DiT 主干
21
      │
22
2024: CogVideoX (Tsinghua) — 开源视频 DiT，3D 注意力
23
      │
24
2024: Wan2.1 (ByteDance) — 开源视频 DiT，业界最佳
25
      │
26
2024: Open-Sora (Vision SE) — 开源复现 Sora
27
      │
28
2025: Sora 2.0 / Stable Video Diffusion 2 / HunyuanVideo

1
class VideoVAE:
2
    """
3
    视频 VAE: 压缩视频到低维潜空间。
4
    继承自图像 VAE，增加了时间压缩维度。
5
    """
6

7
    def __init__(self, spatial_downsample=8, temporal_downsample=4):
8
        """
9
        参数:
10
            spatial_downsample: 空间压缩比 (H → H/8)
11
            temporal_downsample: 时间压缩比 (T → T/4)
12
        """
13
        self.spatial_factor = spatial_downsample    # 8
14
        self.temporal_factor = temporal_downsample  # 4
15

16
    def encode(self, video):
17
        """
18
        编码视频到潜空间。
19

20
        参数:
21
            video: (B, T, C, H, W) 原始视频
22
        返回:
23
            z: (B, T', C', H', W') 潜变量序列
24
        """
25
        B, T, C, H, W = video.shape
26

27
        # 空间压缩: H/8, W/8
28
        # 时间压缩: T/4
29
        # 通道扩展: C=3 → C'=16 或 16+16=32 (3D VAE)
30

31
        z = self.encoder(video)  # (B, T/4, 16, H/8, W/8)
32
        return z
33

34
    def decode(self, z):
35
        """
36
        从潜空间解码到视频。
37

38
        参数:
39
            z: (B, T', C', H', W') 潜变量序列
40
        返回:
41
            video: (B, T, C, H, W) 重建视频
42
        """
43
        return self.decoder(z)
44

45
    def compress_ratio(self):
46
        """
47
        压缩比计算。
48
        以 16 帧 512×512 视频为例:
49
        """
50
        T, H, W, C = 16, 512, 512, 3
51

52
        # 像素数量
53
        pixels = T * H * W * C  # 16 × 512 × 512 × 3 = 12,582,912
54

55
        # 潜空间大小
56
        T_z = T // 4             # 4
57
        H_z = H // 8             # 64
58
        W_z = W // 8             # 64
59
        C_z = 16                 # 潜通道数
60

61
        latent_pixels = T_z * H_z * W_z * C_z  # 4 × 64 × 64 × 16 = 262,144
62

63
        ratio = pixels / latent_pixels  # ≈ 48 倍（与图像 LDM 相同）
64
        return ratio

1
def vae_comparison():
2
    """
3
    视频 VAE vs 图像 VAE。
4
    """
5
    return {
6
        "维度": {
7
            "图像 VAE": "3D (H, W, C)",
8
            "视频 VAE": "5D (T, H, W, C)",
9
        },
10
        "编码器": {
11
            "图像 VAE": "2D Conv (H, W)",
12
            "视频 VAE": "3D Conv (T, H, W) — 同时压缩时空",
13
        },
14
        "时间压缩": {
15
            "图像 VAE": "无",
16
            "视频 VAE": "T/4 或 T/8 — 减少时序冗余",
17
        },
18
        "重建质量": {
19
            "图像 VAE": "高（静态纹理）",
20
            "视频 VAE": "中（时间闪烁是主要问题）",
21
        },
22
        "时序一致性": {
23
            "图像 VAE": "无（逐帧编码）",
24
            "视频 VAE": "高（联合编码）",
25
        },
26
        "代表模型": {
27
            "图像 VAE": "SD 的 VAE (4 ch latent)",
28
            "视频 VAE": "Sora 的 3D VAE, CogVideoX VAE, Wan VAE",
29
        },
30
    }

1
class VideoVAE3D(nn.Module):
2
    """
3
    3D 视频 VAE。
4
    用于 Sora、CogVideoX、Wan2.1 等视频 DiT 模型。
5
    """
6

7
    def __init__(self, in_channels=3, latent_channels=16):
8
        super().__init__()
9

10
        # ============ 3D 编码器 ============
11
        self.encoder = nn.Sequential(
12
            # 块 1: T×512×512 → T×256×256
13
            nn.Conv3d(in_channels, 128, kernel_size=3, padding=1),
14
            ResBlock3D(128, 128),
15
            ResBlock3D(128, 128),
16
            nn.Conv3d(128, 128, kernel_size=3, stride=(1, 2, 2), padding=1),  # 只压缩空间
17

18
            # 块 2: T×256×256 → T×128×128
19
            nn.Conv3d(128, 256, kernel_size=3, padding=1),
20
            ResBlock3D(256, 256),
21
            ResBlock3D(256, 256),
22
            nn.Conv3d(256, 256, kernel_size=3, stride=(1, 2, 2), padding=1),
23

24
            # 块 3: T×128×128 → T×64×64
25
            nn.Conv3d(256, 512, kernel_size=3, padding=1),
26
            ResBlock3D(512, 512),
27
            ResBlock3D(512, 512),
28
            nn.Conv3d(512, 512, kernel_size=3, stride=(1, 2, 2), padding=1),
29

30
            # 块 4: T×64×64 → T/2×64×64 (时间压缩开始)
31
            nn.Conv3d(512, 512, kernel_size=3, stride=(2, 1, 1), padding=1),
32

33
            # 最终投影到潜空间
34
            nn.GroupNorm(32, 512),
35
            nn.SiLU(),
36
            nn.Conv3d(512, latent_channels, kernel_size=3, padding=1),
37
        )
38

39
        # ============ 量化/正则化 ============
40
        # 视频 VAE 通常使用 VQ 或 EMA 量化（而非 KL）
41
        self.quant_conv = nn.Conv3d(latent_channels, latent_channels, 1)
42

43
        # ============ 3D 解码器 ============
44
        self.decoder = nn.Sequential(
45
            # 潜空间 → T/2×64×64
46
            nn.Conv3d(latent_channels, 512, kernel_size=3, padding=1),
47
            ResBlock3D(512, 512),
48

49
            # 时间上采样: T/2×64×64 → T×64×64
50
            nn.Conv3d(512, 512, kernel_size=4, stride=(2, 1, 1), padding=1),
51

52
            # 空间上采样: T×64×64 → T×512×512
53
            nn.Conv3d(512, 512, kernel_size=3, padding=1),
54
            ResBlock3D(512, 512),
55
            ResBlock3D(512, 512),
56
            nn.ConvTranspose3d(512, 256, kernel_size=4, stride=(1, 2, 2), padding=1),
57

58
            nn.Conv3d(256, 256, kernel_size=3, padding=1),
59
            ResBlock3D(256, 256),
60
            nn.ConvTranspose3d(256, 128, kernel_size=4, stride=(1, 2, 2), padding=1),
61

62
            nn.Conv3d(128, 128, kernel_size=3, padding=1),
63
            nn.Conv3d(128, in_channels, kernel_size=3, padding=1),
64
        )
65

66
    def encode(self, video):
67
        """
68
        参数:
69
            video: (B, C, T, H, W) — PyTorch 格式
70
        返回:
71
            z: (B, latent_channels, T', H', W')
72
        """
73
        return self.encoder(video)
74

75
    def decode(self, z):
76
        return self.decoder(z)
77

78

79
class ResBlock3D(nn.Module):
80
    """3D 残差块。"""
81
    def __init__(self, in_channels, out_channels):
82
        super().__init__()
83
        self.norm1 = nn.GroupNorm(32, in_channels)
84
        self.conv1 = nn.Conv3d(in_channels, out_channels, 3, padding=1)
85
        self.norm2 = nn.GroupNorm(32, out_channels)
86
        self.conv2 = nn.Conv3d(out_channels, out_channels, 3, padding=1)
87

88
        if in_channels != out_channels:
89
            self.shortcut = nn.Conv3d(in_channels, out_channels, 1)
90
        else:
91
            self.shortcut = nn.Identity()
92

93
    def forward(self, x):
94
        h = F.silu(self.norm1(x))
95
        h = self.conv1(h)
96
        h = F.silu(self.norm2(h))
97
        h = self.conv2(h)
98
        return h + self.shortcut(x)

1
DiT (2D 图像):
2
  ┌────────────────────────────────────┐
3
  │  图像: 512×512×3                   │
4
  │  ↓ Patchify (p=2, 空间)           │
5
  │  每帧: (512/8)×(512/8) = 64×64 patches │
6
  │  每 patch: 2×2×3×C = 12C 维度      │
7
  │  Patch 数量: 64×64 = 4096          │
8
  └────────────────────────────────────┘
9

10
Video DiT (3D 视频):
11
  ┌────────────────────────────────────┐
12
  │  视频: T×512×512×3                 │
13
  │  ↓ Patchify (p_t=2, p_h=2, p_w=2)│
14
  │  时空: (T/4)×(512/8)×(512/8)      │
15
  │       = T_s × 64 × 64 patches     │
16
  │  每 patch: 2×2×2×3×C = 24C 维度   │
17
  │  Patch 数量: T_s × 64 × 64        │
18
  │  (T_s = T/4, 时间压缩 4 倍)        │
19
  └────────────────────────────────────┘
20

21
关键: Video DiT 把 3D 时空立方体展平成序列
22
  → 输入: (T_s × 64 × 64) 个 tokens
23
  → 每个 token 维度: 2×2×2×3×C = 24C
24
  → 类似 DiT，通过线性投影到 D 维

1
class SpatioTemporalPatchify(nn.Module):
2
    """
3
    时空 Patchify 层。
4
    将视频 (B, C, T, H, W) 转为 tokens 序列 (B, N, D)。
5
    """
6

7
    def __init__(self, patch_size_t=2, patch_size_h=2, patch_size_w=2,
8
                 in_channels=16, hidden_size=1024):
9
        super().__init__()
10
        self.pt = patch_size_t    # 时间 patch 大小
11
        self.ph = patch_size_h    # 空间 patch 大小
12
        self.pw = patch_size_w    # 空间 patch 大小
13
        self.in_channels = in_channels  # VAE 潜通道数
14

15
        # 每个 patch 展平后的维度
16
        patch_dim = in_channels * pt * ph * pw  # 16 × 2 × 2 × 2 = 128
17

18
        # 线性投影到 Transformer 维度
19
        self.proj = nn.Linear(patch_dim, hidden_size)
20

21
    def forward(self, z):
22
        """
23
        参数:
24
            z: (B, C, T_s, H_s, W_s) 潜变量视频
25
               例如: (B, 16, 4, 64, 64)
26
        返回:
27
            tokens: (B, N, D) Transformer 输入序列
28
               其中 N = T_s/pt × H_s/ph × W_s/pw
29
        """
30
        B, C, T, H, W = z.shape
31

32
        # 1) 划分 patches
33
        # z: (B, C, T, H, W) → (B, T/pt, C, pt, H/ph, W/pw, ph, pw)
34
        z = z.view(
35
            B,
36
            T // self.pt, self.pt,
37
            H // self.ph, self.ph,
38
            W // self.pw, self.pw,
39
            C,
40
        )
41
        # 转置: (B, T/pt, H/ph, W/pw, pt, ph, pw, C)
42
        z = z.permute(0, 1, 3, 5, 2, 4, 6, 7).contiguous()
43

44
        # 2) 展平每个 patch: (B, T_s, H_s, W_s, pt×ph×pw×C)
45
        N_t = T // self.pt
46
        N_h = H // self.ph
47
        N_w = W // self.pw
48
        z = z.view(B, N_t * N_h * N_w, self.pt * self.ph * self.pw * C)
49

50
        # 3) 线性投影到 hidden_size
51
        tokens = self.proj(z)  # (B, N, D)
52

53
        return tokens
54

55

56
class VideoDiT(nn.Module):
57
    """
58
    视频 DiT。
59
    """
60

61
    def __init__(self, hidden_size=1024, num_layers=16,
62
                 num_heads=16, patch_size_t=2, patch_size_h=2, patch_size_w=2,
63
                 in_channels=16):
64
        super().__init__()
65

66
        # Patchify
67
        self.patchify = SpatioTemporalPatchify(
68
            patch_size_t, patch_size_h, patch_size_w,
69
            in_channels, hidden_size
70
        )
71

72
        # Patch Embedding (与 DiT 相同)
73
        num_patches = self.patchify.num_patches
74
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size))
75

76
        # Transformer
77
        self.transformer = TransformerBlock(
78
            hidden_size=hidden_size,
79
            num_heads=num_heads,
80
            mlp_ratio=4.0,
81
        )
82

83
        # 输出投影
84
        self.proj_out = nn.Linear(hidden_size,
85
                                   patch_size_t * patch_size_h * patch_size_w * in_channels)
86

87
        # 时间步嵌入
88
        self.time_embed = TimeEmbedding(hidden_size)
89

90
    def forward(self, z, t, context):
91
        """
92
        参数:
93
            z: (B, C, T_s, H_s, W_s) 潜变量视频
94
            t: (B,) 时间步
95
            context: (B, seq_len, d) 文本 embedding
96
        """
97
        # 1) Patchify
98
        x = self.patchify(z)  # (B, N, D)
99

100
        # 2) 添加位置编码
101
        x = x + self.pos_embed
102

103
        # 3) 时间步嵌入
104
        t_emb = self.time_embed(t)
105

106
        # 4) Transformer 前向 (带 AdaLN-Zero)
107
        for block in self.transformer.blocks:
108
            x = block(x, t_emb, context)
109

110
        # 5) 逆 Patchify
111
        noise_pred = self.proj_out(x)  # (B, N, patch_dim)
112

113
        return noise_pred

1
Sora 的三大核心组件:
2
  ┌──────────────────────────────────────────────────────────────┐
3
  │                                                              │
4
  │  1. Video SpatioTemporal Patchification                     │
5
  │     ─────────────────────────────────────────────────────  │
6
  │     将视频压缩为时空 patches                                │
7
  │     • 空间: 16×16 pixel patches                            │
8
  │     • 时间: 可变长度 (由视频时长决定)                        │
9
  │     • 每个 patch → 潜向量 (通过视频 VAE)                    │
10
  │     • 类比 DiT: 将 patch 视为 "visual sentence"            │
11
  │                                                              │
12
  │  2. Video DiT (Diffusion Transformer)                     │
13
  │     ─────────────────────────────────────────────────────  │
14
  │     用 Transformer 处理时空 patches                         │
15
  │     • 替代 UNet (之前的方法)                                │
16
  │     • 更好的 scaling 特性                                  │
17
  │     • Transformer 的全局 Attention                         │
18
  │     • 基于 Diffusion (非自回归)                            │
19
  │                                                              │
20
  │  3. Versatile World Simulator (视频作为世界模型)            │
21
  │     ─────────────────────────────────────────────────────  │
22
  │     Sora 不仅是视频生成器                                   │
23
  │     • 可以作为视频条件的视频编辑                            │
24
  │     • 可以做图像到视频                                       │
25
  │     • 可以做视频扩展 (Video Inpainting)                    │
26
  │     • 可以做视频到视频风格迁移                               │
27
  │                                                              │
28
  └──────────────────────────────────────────────────────────────┘

1
def sora_capabilities():
2
    """
3
    Sora 支持的视频生成任务。
4
    """
5
    return {
6
        "文生视频 (T2V)": {
7
            "输入": "一段文本描述",
8
            "输出": "最长 60 秒的连贯视频",
9
            "示例": "prompt = 'A stylish woman walking...'"
10
        },
11
        "图生视频 (I2V)": {
12
            "输入": "一张图像 + 文本描述",
13
            "输出": "从图像开始延续的视频",
14
            "示例": "用 DALL-E 3 生成的图像 → 视频"
15
        },
16
        "视频编辑": {
17
            "输入": "现有视频 + 文本指令",
18
            "输出": "符合指令的修改后视频",
19
            "示例": "'让狗戴帽子' → 视频中狗戴帽子"
20
        },
21
        "视频扩展 (Inpainting/Outpainting)": {
22
            "输入": "部分视频 + 新内容描述",
23
            "输出": "扩展后的完整视频",
24
            "示例": "给定开头 5 秒，生成后续 55 秒"
25
        },
26
        "视频到视频": {
27
            "输入": "视频 + 风格描述",
28
            "输出": "相同内容但不同风格",
29
            "示例": "写实风格 → 动漫风格"
30
        },
31
        "动态封面": {
32
            "输入": "静态图像",
33
            "输出": "包含轻微运动的动态版本",
34
            "示例": "DALL-E 图像 → 轻微动画"
35
        },
36
        "连接视频": {
37
            "输入": "两个视频片段",
38
            "输出": "在两片段间平滑过渡的视频",
39
            "示例": "视频 A → 自然过渡 → 视频 B"
40
        },
41
    }

1
def sora_limitations():
2
    """
3
    Sora 的技术限制（来自 OpenAI 报告）。
4
    """
5
    return {
6
        "物理交互": {
7
            "问题": "复杂物体交互的物理规律不准确",
8
            "示例": "杯子倒下时水不一定符合重力",
9
            "原因": "模型从视频数据中学习，缺乏物理先验",
10
        },
11
        "长程一致性": {
12
            "问题": "视频越长，物体身份可能丢失",
13
            "示例": "第 30 秒的人物外貌可能与第 1 秒不同",
14
            "原因": "即使有 Transformer Attention，长序列仍困难",
15
        },
16
        "精确计数": {
17
            "问题": "难以准确生成指定数量的事物",
18
            "示例": "prompt '5 只狗' → 可能生成 4 或 6 只",
19
            "原因": "视频中的数量信息不够精确",
20
        },
21
        "空间细节": {
22
            "问题": "文字渲染和精细空间布局不准确",
23
            "示例": "生成正确文本的招牌困难",
24
            "原因": "视频空间分辨率有限",
25
        },
26
        "相机轨迹": {
27
            "问题": "复杂相机运动可能不连贯",
28
            "示例": "360 度旋转镜头可能中途变形",
29
            "原因": "3D 一致性难以保证",
30
        },
31
    }

1
def cogvideoX_architecture():
2
    """
3
    CogVideoX 架构特点。
4
    """
5
    return {
6
        "团队": "Tsinghua & Zhipu AI",
7
        "发布时间": "2024年7月",
8
        "模型规模": "5B 参数 (DiT) + 600M (Text Encoder)",
9
        "关键创新": "Expert Transformer (专家 Transformer)",
10
        "注意力机制": {
11
            "方法": "3D Full Attention",
12
            "特点": "所有 tokens 之间的时空 Attention",
13
            "问题": "O(N^2) 复杂度，N 是时空 tokens 数",
14
        },
15
        "文本编码器": "T5Encoder (来自 Google Flan-T5)",
16
        "视频 VAE": "自研视频 VAE (4x 空间压缩，4x 时间压缩)",
17
        "分辨率": "480×480, 16 帧",
18
        "训练数据": "约 10M 文本-视频对",
19
        "开源": "权重完全开源，推理代码开源",
20
    }
21

22

23
def cogvideoX_transformer():
24
    """
25
    CogVideoX 的 Transformer 结构。
26
    """
27
    def forward(self, x, t, context, attention_mask=None):
28
        # x: (B, N, D) — N = T_s × H_s × W_s
29

30
        # 每个 Transformer Block:
31
        # 1) Self-Attention (Q=K=V=x)
32
        x = x + self.attn1(self.norm1(x), attention_mask=attention_mask)
33

34
        # 2) Cross-Attention (Q=x, K=V=context)
35
        x = x + self.attn2(self.norm2(x), context, attention_mask=attention_mask)
36

37
        # 3) FFN
38
        x = x + self.mlp(self.norm3(x))
39

40
        return x

1
def wan21_architecture():
2
    """
3
    Wan2.1 架构特点。
4
    """
5
    return {
6
        "团队": "ByteDance (字节跳动)",
7
        "发布时间": "2025年1月",
8
        "模型规模": "14B 参数 (DiT) + 双文本编码器",
9
        "关键创新": "Causal 3D Attention (因果 3D 注意力)",
10
        "注意力机制": {
11
            "方法": "Causal 3D Attention",
12
            "特点": "时间维度因果 (只 attend 到之前和当前的帧)",
13
            "空间维度": "全空间 Attention",
14
            "优势": "比 Full Attention 更高效，适合长视频",
15
        },
16
        "文本编码器": "T5-XXL (1.8B) + CLIP ViT-L (300M)",
17
        "视频 VAE": "自研 Wan VAE",
18
        "分辨率": "480P / 720P / 1080P",
19
        "帧数": "最长 5 分钟 (3000 帧 @ 10fps)",
20
        "长视频生成": "通过自回归方式生成长视频",
21
        "开源": "14B 模型完全开源 (HuggingFace)",
22
        "显存需求": "FP16 需要 ~28GB (单卡 A100)",
23
    }

1
def opensora_architecture():
2
    """
3
    Open-Sora 架构特点。
4
    """
5
    return {
6
        "团队": "Vision SE (社区主导)",
7
        "发布时间": "2024年3月",
8
        "目标": "开源复现 Sora 的完整流程",
9
        "关键设计": "ST-DiT (Spatio-Temporal DiT)",
10
        "注意力机制": {
11
            "方法": "Causal 3D Attention",
12
            "空间: 全 Attention",
13
            "时间: Causal (只向前看)",
14
        },
15
        "视频 VAE": "支持多种 VAE (自研, CogVideoX, Wan)",
16
        "特点": "完全开源，民主化视频生成技术",
17
        "社区": "活跃的贡献者生态",
18
    }

1
def opensource_comparison():
2
    """
3
    CogVideoX vs Wan2.1 vs Open-Sora 对比。
4
    """
5
    return {
6
        "模型规模": {
7
            "CogVideoX": "5B 参数",
8
            "Wan2.1": "14B 参数",
9
            "Open-Sora": "2B-9B (多个版本)",
10
        },
11
        "注意力机制": {
12
            "CogVideoX": "3D Full Attention (全局)",
13
            "Wan2.1": "Causal 3D Attention (时间因果)",
14
            "Open-Sora": "Causal 3D Attention",
15
        },
16
        "分辨率": {
17
            "CogVideoX": "480×480",
18
            "Wan2.1": "480P - 1080P",
19
            "Open-Sora": "16:9, 9:16, 1:1",
20
        },
21
        "帧数": {
22
            "CogVideoX": "最多 32 帧",
23
            "Wan2.1": "最长 5 分钟",
24
            "Open-Sora": "最多 240 帧",
25
        },
26
        "视频质量": {
27
            "CogVideoX": "★★★☆☆ (早期开源)",
28
            "Wan2.1": "★★★★★ (最佳开源)",
29
            "Open-Sora": "★★★★☆ (持续改进)",
30
        },
31
        "开源程度": {
32
            "CogVideoX": "权重 + 代码完全开源",
33
            "Wan2.1": "权重 + 代码完全开源",
34
            "Open-Sora": "权重 + 代码完全开源",
35
        },
36
    }

1
VideoLDM 架构:
2
  ┌──────────────────────────────────────────────────────────────┐
3
  │                                                              │
4
  │  输入视频: T × 512 × 512 × 3                                │
5
  │           ↓                                                  │
6
  │  视频 VAE (3D): T × 512 × 512 × 3                          │
7
  │           → T/4 × 64 × 64 × 4 (潜空间)                    │
8
  │           ↓                                                  │
9
  │  时间层 (Temporal Layers):                                   │
10
  │    在潜空间添加时间建模层 (3D Conv / Temporal Attention)      │
11
  │           ↓                                                  │
12
  │  空间层 (Spatial Layers):                                    │
13
  │    SD 的 UNet 保持不变 (处理每帧空间)                       │
14
  │           ↓                                                  │
15
  │  扩散过程: 在 T/4 × 64 × 64 潜空间执行 DDPM                │
16
  │           ↓                                                  │
17
  │  视频 VAE 解码: 潜变量 → 视频                              │
18
  │                                                              │
19
  │  关键设计:                                                   │
20
  │    • 时间层插入 SD UNet 的 ResBlock 之间                    │
21
  │    • 时间层在潜空间的 T 维度上建模                          │
22
  │    • 冻结空间层（使用预训练 SD 权重）                       │
23
  │    • 只训练时间层（高效微调）                               │
24
  │                                                              │
25
  └──────────────────────────────────────────────────────────────┘

1
class TemporalLayer(nn.Module):
2
    """
3
    VideoLDM 的时间层。
4
    插入到 SD UNet 的 ResBlock 之间。
5
    """
6
    def __init__(self, channels, num_frames):
7
        super().__init__()
8
        self.channels = channels
9
        self.num_frames = num_frames
10

11
        # 3D 卷积 (时间 + 空间)
12
        # 作用: 在时间维度上建模帧间关系
13
        self.temporal_conv = nn.Sequential(
14
            nn.Conv3d(channels, channels, kernel_size=3, padding=1),
15
            nn.GroupNorm(32, channels),
16
            nn.SiLU(),
17
            nn.Conv3d(channels, channels, kernel_size=3, padding=1),
18
        )
19

20
        # 残差连接
21
        self.residual_scale = 0.1  # 小权重开始，避免破坏预训练
22

23
    def forward(self, x):
24
        """
25
        参数:
26
            x: (B*T, C, H, W) — UNet 的特征
27
        返回:
28
            融合了时间信息的特征
29
        """
30
        BT, C, H, W = x.shape
31

32
        # 恢复时间维度: (B*T, C, H, W) → (B, C, T, H, W)
33
        B = BT // self.num_frames
34
        x = x.view(B, self.num_frames, C, H, W)
35
        x = x.permute(0, 2, 1, 3, 4)  # (B, C, T, H, W)
36

37
        # 3D 卷积 (同时处理时间 + 空间)
38
        h = self.temporal_conv(x)
39

40
        # 残差连接
41
        h = x + self.residual_scale * h
42

43
        # 展平回: (B, C, T, H, W) → (B*T, C, H, W)
44
        h = h.permute(0, 2, 1, 3, 4).reshape(BT, C, H, W)
45
        return h
46

47

48
class VideoLDMUNet(nn.Module):
49
    """
50
    带时间层的 VideoLDM UNet。
51
    """
52

53
    def __init__(self, base_unet, num_frames=16):
54
        super().__init__()
55
        self.unet = base_unet
56
        self.num_frames = num_frames
57

58
        # 在每个 ResBlock 后插入时间层
59
        self.temporal_layers = nn.ModuleList([
60
            TemporalLayer(ch, num_frames)
61
            for ch in base_unet.channel_list
62
        ])
63

64
    def forward(self, x_t, t, context):
65
        """
66
        参数:
67
            x_t: (B*T, C, H, W) — 展平的时空潜变量
68
            t: (B,) — 时间步（所有帧相同）
69
            context: (B, seq_len, d) — 文本 embedding
70
        """
71
        # 通过 UNet（处理空间信息）
72
        h = self.unet.forward_spatial(x_t, t, context)
73

74
        # 通过时间层（处理时间信息）
75
        for temp_layer in self.temporal_layers:
76
            h = temp_layer(h)
77

78
        return h

1
def lvdm_improvements():
2
    """
3
    LVDM 相比 VideoLDM 的改进。
4
    """
5
    return {
6
        "长视频生成": {
7
            "VideoLDM": "固定长度 (如 16 帧)",
8
            "LVDM": "通过分层生成支持可变长度",
9
        },
10
        "时间建模": {
11
            "VideoLDM": "3D Conv（局部时间）",
12
            "LVDM": "Temporal Attention（全局时间）",
13
        },
14
        "训练效率": {
15
            "VideoLDM": "需要逐帧编码",
16
            "LVDM": "批量视频编码，更高效",
17
        },
18
        "无条件生成": {
19
            "VideoLDM": "困难",
20
            "LVDM": "通过分层设计改善",
21
        },
22
    }

1
def attention_3d():
2
    """
3
    3D 时空注意力的计算方式。
4
    """
5

6
    def compute_full_3d_attention(tokens_3d, num_heads=16):
7
        """
8
        全 3D 注意力（所有 tokens 互相关注）。
9

10
        参数:
11
            tokens_3d: (B, T, H, W, D) — 时空 tokens
12
        返回:
13
            输出: (B, T, H, W, D) — 更新后的 tokens
14
        """
15
        B, T, H, W, D = tokens_3d.shape
16
        N = T * H * W  # tokens 总数
17

18
        # 展平: (B, T, H, W, D) → (B, N, D)
19
        x = tokens_3d.view(B, N, D)
20

21
        # 分头
22
        H_heads = num_heads
23
        d_head = D // H_heads
24

25
        # Q, K, V: (B, N, D)
26
        Q = compute_Q(x)  # (B, N, D)
27
        K = compute_K(x)
28
        V = compute_V(x)
29

30
        # Reshape: (B, N, D) → (B, H, N, d_head)
31
        Q = Q.view(B, N, H_heads, d_head).transpose(1, 2)
32
        K = K.view(B, N, H_heads, d_head).transpose(1, 2)
33
        V = V.view(B, N, H_heads, d_head).transpose(1, 2)
34

35
        # 注意力: (B, H, N, N) — O(N^2) 复杂度
36
        attn = (Q @ K.transpose(-2, -1)) / (d_head ** 0.5)
37
        attn = F.softmax(attn, dim=-1)
38

39
        # 输出: (B, H, N, d_head) → (B, N, D)
40
        out = (attn @ V).transpose(1, 2).reshape(B, N, D)
41
        return out.view(B, T, H, W, D)
42

43
    def compute_causal_3d_attention(tokens_3d, num_heads=16):
44
        """
45
        因果 3D 注意力（时间维度因果）。
46

47
        参数:
48
            tokens_3d: (B, T, H, W, D)
49
        返回:
50
            输出: (B, T, H, W, D)
51
        """
52
        B, T, H, W, D = tokens_3d.shape
53
        N = T * H * W
54
        x = tokens_3d.view(B, N, D)
55

56
        # Q, K, V (同上)
57
        Q, K, V = compute_QKV(x)
58

59
        # 掩码: 时间维度因果
60
        # 帧 t 只能 attend 到帧 0, 1, ..., t
61
        # 不能 attend 到帧 t+1, ..., T-1
62
        mask = torch.zeros(T, T, device=tokens_3d.device)
63
        for i in range(T):
64
            for j in range(i + 1, T):
65
                mask[i, j] = float("-inf")
66

67
        # 扩展掩码到 (B, H, N, N)
68
        mask_3d = mask.unsqueeze(0).unsqueeze(0)  # (1, 1, T, T)
69
        mask_3d = mask_3d.unsqueeze(-1).unsqueeze(-1)  # (1, 1, T, T, 1, 1)
70
        mask_3d = mask_3d.expand(B, H_heads, T, T, H * W, H * W)
71

72
        # 应用掩码
73
        attn = (Q @ K.transpose(-2, -1)) / (d_head ** 0.5)
74
        attn = attn + mask_3d  # 因果掩码
75
        attn = F.softmax(attn, dim=-1)
76

77
        return (attn @ V).view(B, T, H, W, D)
78

79
    return {"full_3d": compute_full_3d_attention, "causal_3d": compute_causal_3d_attention}

1
def attention_complexity():
2
    """
3
    注意力机制的计算复杂度对比。
4
    """
5
    return {
6
        "2D Full Attention (DiT, 图像)": {
7
            "图像尺寸": "512×512",
8
            "Patch 数 (p=2, 压缩8x)": "64×64 = 4096",
9
            "复杂度": "O(4096²) = 16.8M",
10
        },
11
        "3D Full Attention (视频)": {
12
            "视频尺寸": "16×512×512",
13
            "Patch 数": "4×64×64 = 16384",
14
            "复杂度": "O(16384²) = 268M",
15
        },
16
        "3D Causal Attention": {
17
            "视频尺寸": "16×512×512",
18
            "有效计算": "每帧只 attend 到之前帧",
19
            "复杂度": "O(T² × H² × W²) = 约 50M",
20
            "节省": "约 5x",
21
        },
22
        "Sparse Attention (可选)": {
23
            "方法": "空间局部 + 时间稀疏",
24
            "复杂度": "O(T × H × W × k²)",
25
            "k": "局部窗口大小 (如 7)",
26
        },
27
    }

1
def train_video_dit(video_model, vae, text_encoder, dataloader,
2
                   optimizer, device, num_frames=16):
3
    """
4
    视频 DiT 训练循环。
5
    """
6

7
    def forward_batch(videos, prompts):
8
        """
9
        参数:
10
            videos: (B, T, C, H, W) 原始视频
11
            prompts: list[str] 文本描述
12
        """
13
        B, T, C, H, W = videos.shape
14

15
        # ============ 1) 视频 VAE 编码 ============
16
        with torch.no_grad():
17
            # (B, T, C, H, W) → (B, C, T, H, W) (PyTorch 格式)
18
            videos = videos.permute(0, 1, 2, 3, 4)
19
            z = vae.encode(videos)  # (B, C', T', H', W')
20

21
        # ============ 2) 文本编码 ============
22
        context = text_encoder.encode(prompts)  # (B, seq_len, d)
23

24
        # ============ 3) 加噪 ============
25
        t = torch.randint(0, 1000, (B,), device=device)
26
        noise = torch.randn_like(z)
27
        alpha_bar = diffusion_schedule(t)
28

29
        # 对所有帧加相同噪声（时序一致性关键）
30
        alpha_bar = alpha_bar.view(B, 1, 1, 1, 1)
31
        z_noisy = alpha_bar.sqrt() * z + (1 - alpha_bar).sqrt() * noise
32

33
        # ============ 4) Patchify ============
34
        noise_pred = video_model(z_noisy, t, context)
35

36
        # ============ 5) 损失 ============
37
        loss = F.mse_loss(noise_pred, noise)
38

39
        return loss
40

41
    # 训练循环
42
    video_model.train()
43
    for batch in tqdm(dataloader):
44
        videos, prompts = batch["video"].to(device), batch["prompt"]
45

46
        loss = forward_batch(videos, prompts)
47

48
        optimizer.zero_grad()
49
        loss.backward()
50
        optimizer.step()
51

52
    return loss.item()