CogVideoX: Text-to-Video Diffusion with Expert Transformer.

Implements the CogVideoX architecture from "CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer" (Yang et al., ZHIPU AI 2024). A state-of-the-art open-source video generation model combining 3D causal VAE compression with an expert transformer that routes text and video tokens to specialized FFN experts.

Key Innovations

1. 3D Causal VAE

Compresses video to latent space while preserving temporal causality:

• Spatial compression: 8× downsample in H and W via 2D convolutions
• Temporal compression: 4× downsample in time via causal 1D convolutions
• Causal constraint: Frame t can only depend on frames ≤ t (enables streaming)
• 3D convolutions: Factorized as spatial 2D conv + temporal 1D conv

Video [B, T, C, H, W]  →  Latent [B, T/4, C', H/8, W/8]
     (49 frames)              (13 latent frames)

2. Expert Transformer (DiT + MoE FFN)

Full 3D attention over space-time with modality-specific experts:

• 3D patchification: Video latent → space-time tokens
• Full 3D attention: Every token attends to all other tokens (expensive but high quality)
• Expert FFN routing: Text tokens → text experts, video tokens → video experts
• RoPE3D: Rotary position embeddings over (time, height, width)

Text tokens ────┐
                ├──→ Shared 3D Attention ──→ Expert FFN ──→ ...
Video tokens ───┘         │                     │
                     (all-to-all)          (routed by modality)

Architecture Details

Input: Video [batch, frames, 3, H, W] + Text embeddings [batch, text_len, dim]
       |
       v
+---------------------------+
| 3D Causal VAE Encoder     |  Compress video to latent space
+---------------------------+
       |
       v
Latent [batch, T', latent_dim, H', W']
       |
       v
+---------------------------+
| 3D Patchify + RoPE3D      |  Convert to space-time tokens
+---------------------------+
       |
       v
Concat with text tokens → [batch, text_len + video_tokens, hidden]
       |
       v
+---------------------------+
| Expert Transformer Block  |  × num_layers
|  • Full 3D Self-Attention |
|  • Expert FFN (text/video)|
+---------------------------+
       |
       v
| Unpatchify + VAE Decode   |
       |
       v
Output: Generated Video [batch, frames, 3, H, W]

Usage

# Build the 3D causal VAE
vae = CogVideoX.build_vae(
  in_channels: 3,
  latent_channels: 16,
  num_frames: 49
)

# Build the expert transformer
transformer = CogVideoX.build_transformer(
  patch_size: [1, 2, 2],
  hidden_size: 1920,
  num_heads: 48,
  num_layers: 42,
  num_frames: 49
)

# Or build the full pipeline
model = CogVideoX.build(
  hidden_size: 1920,
  num_heads: 48,
  num_layers: 42
)

References

• Paper: "CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer"
• Authors: Yang et al., ZHIPU AI
• Year: 2024
• Code: https://github.com/THUDM/CogVideo