| """A tiny text-to-video diffusion transformer, trained from random init with rectified flow.
|
|
|
| Video [B,3,T,H,W] is cut into (1 x p x p) patches. Each block runs spatial attention
|
| within a frame, temporal attention across frames at the same patch position,
|
| cross-attention to the caption, and an MLP, all modulated (adaLN-Zero) by the
|
| diffusion time and pooled caption. The caption encoder is a 2-layer transformer
|
| over the small word vocabulary in data.py.
|
| """
|
| from __future__ import annotations
|
|
|
| import copy
|
| import math
|
| from dataclasses import asdict, dataclass
|
|
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| from torch.utils.checkpoint import checkpoint
|
|
|
| from .data import NULL, PAD, VOCAB, MAX_TOKENS
|
| from .world import BYTE_VOCAB, HYBRID_VOCAB, MAX_BYTES, MAX_HYBRID, NULL as BYTE_NULL, PAD as BYTE_PAD
|
|
|
|
|
| def specials(text: str) -> tuple[int, int]:
|
| """(pad id, null id) for the caption tokenizer a model was trained with."""
|
| return (BYTE_PAD, BYTE_NULL) if text in ("bytes", "hybrid") else (PAD, NULL)
|
|
|
|
|
| @dataclass
|
| class ModelConfig:
|
| frames: int = 16
|
| size: int = 48
|
| patch: int = 4
|
| dim: int = 384
|
| depth: int = 8
|
| heads: int = 6
|
| text_depth: int = 2
|
| vocab_size: int = len(VOCAB)
|
| max_tokens: int = MAX_TOKENS
|
| mlp_ratio: float = 4.0
|
| attn: str = "factorized"
|
| rope: bool = False
|
| text: str = "words"
|
| grad_checkpoint: bool = False
|
|
|
|
|
|
|
|
|
| def modulate(x, shift, scale):
|
| return x * (1 + scale) + shift
|
|
|
|
|
| class Attention(nn.Module):
|
| def __init__(self, dim: int, heads: int, cross: bool = False):
|
| super().__init__()
|
| self.heads = heads
|
| self.q = nn.Linear(dim, dim)
|
| self.kv = nn.Linear(dim, 2 * dim)
|
| self.out = nn.Linear(dim, dim)
|
| hd = dim // heads
|
| self.q_norm = nn.RMSNorm(hd)
|
| self.k_norm = nn.RMSNorm(hd)
|
|
|
| def forward(self, x, ctx=None, mask=None, rope=None):
|
| B, N, D = x.shape
|
| ctx = x if ctx is None else ctx
|
| q = self.q(x).view(B, N, self.heads, -1).transpose(1, 2)
|
| k, v = self.kv(ctx).view(B, ctx.shape[1], 2, self.heads, -1).permute(2, 0, 3, 1, 4)
|
| q, k = self.q_norm(q), self.k_norm(k)
|
| if rope is not None:
|
| q, k = apply_rope(q, rope), apply_rope(k, rope)
|
| o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
|
| return self.out(o.transpose(1, 2).reshape(B, N, D))
|
|
|
|
|
| def rope_table(t, y, x, head_dim: int, base: float = 100.0):
|
| """cos/sin [L, head_dim/2] for integer positions t, y, x (each [L]). Head dims split 1/4 t, 3/8 y, 3/8 x."""
|
| pairs = head_dim // 2
|
| pt = pairs // 4
|
| py = (pairs - pt) // 2
|
| angles = []
|
| for pos, n in ((t, pt), (y, py), (x, pairs - pt - py)):
|
| freqs = base ** (-torch.arange(n, device=pos.device, dtype=torch.float32) / max(n, 1))
|
| angles.append(pos.float()[:, None] * freqs[None])
|
| a = torch.cat(angles, dim=1)
|
| return a.cos(), a.sin()
|
|
|
|
|
| def apply_rope(x, rope):
|
| """Rotate consecutive channel pairs of x [B,H,L,hd] by the angles in rope."""
|
| cos, sin = rope
|
| x1, x2 = x.float().unflatten(-1, (-1, 2)).unbind(-1)
|
| out = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2)
|
| return out.to(x.dtype)
|
|
|
|
|
| class Mlp(nn.Sequential):
|
| def __init__(self, dim: int, ratio: float):
|
| hidden = int(dim * ratio)
|
| super().__init__(nn.Linear(dim, hidden), nn.GELU(approximate="tanh"), nn.Linear(hidden, dim))
|
|
|
|
|
| class Block(nn.Module):
|
| def __init__(self, cfg: ModelConfig):
|
| super().__init__()
|
| D = cfg.dim
|
| self.norms = nn.ModuleList([nn.LayerNorm(D, elementwise_affine=False, eps=1e-6) for _ in range(4)])
|
| self.spatial = Attention(D, cfg.heads)
|
| self.temporal = Attention(D, cfg.heads)
|
| self.cross = Attention(D, cfg.heads, cross=True)
|
| self.mlp = Mlp(D, cfg.mlp_ratio)
|
| self.ada = nn.Sequential(nn.SiLU(), nn.Linear(D, 11 * D))
|
| assert cfg.attn in ("factorized", "joint"), cfg.attn
|
| self.joint = cfg.attn == "joint"
|
|
|
| def forward(self, x, c, ctx, ctx_mask, ropes=None):
|
|
|
| B, T, N, D = x.shape
|
| (s1, sc1, g1, s2, sc2, g2, g3, s4, sc4, g4, sc3) = self.ada(c)[:, None, None].chunk(11, dim=-1)
|
|
|
| r_joint, r_space, r_time = ropes or (None, None, None)
|
| if self.joint:
|
| h = modulate(self.norms[0](x), s1, sc1).reshape(B, T * N, D)
|
| x = x + g1 * self.spatial(h, rope=r_joint).view(B, T, N, D)
|
| else:
|
| h = modulate(self.norms[0](x), s1, sc1).reshape(B * T, N, D)
|
| x = x + g1 * self.spatial(h, rope=r_space).view(B, T, N, D)
|
|
|
| h = modulate(self.norms[1](x), s2, sc2).transpose(1, 2).reshape(B * N, T, D)
|
| x = x + g2 * self.temporal(h, rope=r_time).view(B, N, T, D).transpose(1, 2)
|
|
|
| h = (self.norms[2](x) * (1 + sc3)).reshape(B, T * N, D)
|
| x = x + g3 * self.cross(h, ctx, ctx_mask).view(B, T, N, D)
|
|
|
| h = modulate(self.norms[3](x), s4, sc4)
|
| return x + g4 * self.mlp(h)
|
|
|
|
|
| class TextEncoder(nn.Module):
|
| def __init__(self, cfg: ModelConfig):
|
| super().__init__()
|
| self.pad = specials(cfg.text)[0]
|
| self.embed = nn.Embedding(cfg.vocab_size, cfg.dim)
|
| self.pos = nn.Parameter(torch.randn(cfg.max_tokens, cfg.dim) * 0.02)
|
| layer = nn.TransformerEncoderLayer(cfg.dim, cfg.heads, int(cfg.dim * cfg.mlp_ratio), dropout=0.0,
|
| activation="gelu", batch_first=True, norm_first=True)
|
| self.layers = nn.TransformerEncoder(layer, cfg.text_depth, enable_nested_tensor=False)
|
| self.norm = nn.LayerNorm(cfg.dim)
|
|
|
| def forward(self, tokens):
|
| pad = tokens == self.pad
|
| h = self.embed(tokens) + self.pos[: tokens.shape[1]]
|
| h = self.norm(self.layers(h, src_key_padding_mask=pad))
|
| keep = (~pad).unsqueeze(-1).to(h.dtype)
|
| pooled = (h * keep).sum(1) / keep.sum(1).clamp(min=1)
|
| return h, pooled, (~pad)[:, None, None, :]
|
|
|
|
|
| def timestep_embedding(t, dim: int):
|
| half = dim // 2
|
| freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device, dtype=torch.float32) / half)
|
| args = (t.float() * 1000)[:, None] * freqs[None]
|
| return torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
|
|
|
|
|
|
|
|
| class VideoDiT(nn.Module):
|
| def __init__(self, cfg: ModelConfig | None = None):
|
| super().__init__()
|
| self.cfg = cfg = cfg or ModelConfig()
|
| assert cfg.size % cfg.patch == 0 and cfg.dim % cfg.heads == 0 and cfg.attn in ("factorized", "joint")
|
| D, p = cfg.dim, cfg.patch
|
| self.grid = cfg.size // p
|
| self.patch_embed = nn.Conv3d(3, D, kernel_size=(1, p, p), stride=(1, p, p))
|
| self.pos_space = nn.Parameter(torch.randn(self.grid * self.grid, D) * 0.02)
|
| self.pos_time = nn.Parameter(torch.randn(cfg.frames, D) * 0.02)
|
| self.t_embed = nn.Sequential(nn.Linear(256, D), nn.SiLU(), nn.Linear(D, D))
|
| self.text = TextEncoder(cfg)
|
| self.text_proj = nn.Linear(D, D)
|
| self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.depth)])
|
| self.final_norm = nn.LayerNorm(D, elementwise_affine=False, eps=1e-6)
|
| self.final_ada = nn.Sequential(nn.SiLU(), nn.Linear(D, 2 * D))
|
| self.final = nn.Linear(D, p * p * 3)
|
| self._init()
|
|
|
| def _init(self):
|
| def basic(m):
|
| if isinstance(m, nn.Linear):
|
| nn.init.xavier_uniform_(m.weight)
|
| if m.bias is not None:
|
| nn.init.zeros_(m.bias)
|
| self.apply(basic)
|
| w = self.patch_embed.weight
|
| nn.init.xavier_uniform_(w.view(w.shape[0], -1))
|
| nn.init.normal_(self.text.embed.weight, std=0.02)
|
| for b in self.blocks:
|
| nn.init.zeros_(b.ada[-1].weight)
|
| nn.init.zeros_(b.ada[-1].bias)
|
| for m in (self.final_ada[-1], self.final):
|
| nn.init.zeros_(m.weight)
|
| nn.init.zeros_(m.bias)
|
|
|
| def _ropes(self, T: int, device):
|
| g, hd = self.grid, self.cfg.dim // self.cfg.heads
|
| tt, yy, xx = torch.meshgrid(torch.arange(T, device=device), torch.arange(g, device=device),
|
| torch.arange(g, device=device), indexing="ij")
|
| joint = rope_table(tt.flatten(), yy.flatten(), xx.flatten(), hd)
|
| zeros = torch.zeros(g * g, device=device, dtype=torch.long)
|
| space = rope_table(zeros, yy[0].flatten(), xx[0].flatten(), hd)
|
| time = rope_table(torch.arange(T, device=device), zeros[:T], zeros[:T], hd)
|
| return joint, space, time
|
|
|
| def forward(self, x, t, tokens):
|
| """x [B,3,T,H,W], t [B] in [0,1], tokens [B,L] -> predicted velocity [B,3,T,H,W]."""
|
| B, _, T, H, W = x.shape
|
| h = self.patch_embed(x)
|
| h = h.flatten(3).permute(0, 2, 3, 1)
|
| h = h + self.pos_space[None, None] + self.pos_time[None, :T, None]
|
| ctx, pooled, ctx_mask = self.text(tokens)
|
| c = self.t_embed(timestep_embedding(t, 256)) + self.text_proj(pooled)
|
| ropes = self._ropes(T, x.device) if self.cfg.rope else None
|
| for blk in self.blocks:
|
| if self.cfg.grad_checkpoint and self.training and torch.is_grad_enabled():
|
| h = checkpoint(blk, h, c, ctx, ctx_mask, ropes, use_reentrant=False)
|
| else:
|
| h = blk(h, c, ctx, ctx_mask, ropes)
|
| shift, scale = self.final_ada(c)[:, None, None].chunk(2, dim=-1)
|
| h = self.final(modulate(self.final_norm(h), shift, scale))
|
| p, g = self.cfg.patch, self.grid
|
| h = h.view(B, T, g, g, p, p, 3).permute(0, 6, 1, 2, 4, 3, 5)
|
| return h.reshape(B, 3, T, H, W)
|
|
|
|
|
|
|
|
|
| def null_tokens(tokens: torch.Tensor, text: str = "words") -> torch.Tensor:
|
| pad, null = specials(text)
|
| out = torch.full_like(tokens, pad)
|
| out[:, 0] = null
|
| return out
|
|
|
|
|
| def flow_loss(model: VideoDiT, x0: torch.Tensor, tokens: torch.Tensor, p_uncond: float = 0.1) -> torch.Tensor:
|
| """Rectified flow: x_t = (1-t) x0 + t eps, regress v = eps - x0. Logit-normal t."""
|
| B = x0.shape[0]
|
| t = torch.sigmoid(torch.randn(B, device=x0.device))
|
| eps = torch.randn_like(x0)
|
| tt = t.view(B, 1, 1, 1, 1)
|
| xt = (1 - tt) * x0 + tt * eps
|
| drop = torch.rand(B, device=x0.device) < p_uncond
|
| tokens = torch.where(drop[:, None], null_tokens(tokens, model.cfg.text), tokens)
|
| v = model(xt, t, tokens)
|
| return F.mse_loss(v.float(), (eps - x0).float())
|
|
|
|
|
| @torch.no_grad()
|
| def sample(model: VideoDiT, tokens: torch.Tensor, steps: int = 40, cfg_scale: float = 4.0,
|
| generator: torch.Generator | None = None, shift: float = 2.0) -> torch.Tensor:
|
| """Euler integration of the flow from noise (t=1) to data (t=0). Returns [B,3,T,H,W] in [-1,1].
|
|
|
| `shift` > 1 spends more steps at high noise, where the layout and motion are decided.
|
| """
|
| cfg = model.cfg
|
| device = next(model.parameters()).device
|
| tokens = tokens.to(device)
|
| B = tokens.shape[0]
|
| x = torch.randn(B, 3, cfg.frames, cfg.size, cfg.size, device=device, generator=generator)
|
| u = torch.linspace(1, 0, steps + 1, device=device)
|
| ts = shift * u / (1 + (shift - 1) * u)
|
| cond = torch.cat([tokens, null_tokens(tokens, cfg.text)]) if cfg_scale != 1 else tokens
|
| for i in range(steps):
|
| t = ts[i].expand(cond.shape[0])
|
| xin = torch.cat([x, x]) if cfg_scale != 1 else x
|
| with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
|
| v = model(xin, t, cond).float()
|
| if cfg_scale != 1:
|
| vc, vu = v.chunk(2)
|
| v = vu + cfg_scale * (vc - vu)
|
| x = x + (ts[i + 1] - ts[i]) * v
|
| return x.clamp(-1, 1)
|
|
|
|
|
| def to_uint8(x: torch.Tensor) -> torch.Tensor:
|
| """[B,3,T,H,W] in [-1,1] -> uint8 [B,T,H,W,3] on CPU."""
|
| return ((x.float().clamp(-1, 1) + 1) * 127.5).round().to(torch.uint8).permute(0, 2, 3, 4, 1).cpu()
|
|
|
|
|
| class EMA:
|
| def __init__(self, model: nn.Module, decay: float = 0.999):
|
| self.decay = decay
|
| self.model = copy.deepcopy(model).eval().requires_grad_(False)
|
|
|
| @torch.no_grad()
|
| def update(self, model: nn.Module):
|
| for e, p in zip(self.model.parameters(), model.parameters()):
|
| e.lerp_(p.detach(), 1 - self.decay)
|
|
|
|
|
| def save_checkpoint(path, model: VideoDiT, opt=None, step: int = 0, extra: dict | None = None, ema: EMA | None = None):
|
| torch.save({
|
| "config": asdict(model.cfg),
|
| "model": model.state_dict(),
|
| "ema": ema.model.state_dict() if ema else None,
|
| "opt": opt.state_dict() if opt else None,
|
| "step": step,
|
| "extra": extra or {},
|
| }, path)
|
|
|
|
|
| def load_checkpoint(path, device="cpu", use_ema: bool = True):
|
| ckpt = torch.load(path, map_location=device, weights_only=False)
|
| model = VideoDiT(ModelConfig(**ckpt["config"]))
|
| state = ckpt["ema"] if use_ema and ckpt.get("ema") else ckpt["model"]
|
| model.load_state_dict(state)
|
| return model.to(device).eval(), ckpt
|
|
|
|
|
| def tokenizer_for(cfg: ModelConfig):
|
| """The caption tokenizer matching a model's config (v1 word vocabulary or v2 raw bytes)."""
|
| if cfg.text == "bytes":
|
| from .world import ByteTokenizer
|
| return ByteTokenizer(cfg.max_tokens)
|
| if cfg.text == "hybrid":
|
| from .world import HybridTokenizer
|
| return HybridTokenizer(cfg.max_tokens)
|
| from .data import Tokenizer
|
| return Tokenizer(cfg.max_tokens)
|
|
|
|
|
| def byte_model_config(**kw) -> ModelConfig:
|
| """Defaults for world-v2 models: byte captions, joint space-time attention, 3D RoPE."""
|
| base = dict(frames=16, size=32, patch=4, dim=384, depth=8, heads=6, text_depth=4, vocab_size=BYTE_VOCAB,
|
| max_tokens=MAX_BYTES, attn="joint", rope=True, text="bytes")
|
| return ModelConfig(**{**base, **kw})
|
|
|
|
|
| def hybrid_model_config(**kw) -> ModelConfig:
|
| """byte_model_config with hashed-word tokens added to the caption (text="hybrid")."""
|
| return byte_model_config(**{"text": "hybrid", "vocab_size": HYBRID_VOCAB, "max_tokens": MAX_HYBRID, **kw})
|
|
|
|
|
| def load_matching(model: nn.Module, state: dict) -> dict:
|
| """Copy every tensor whose name and shape match; for embedding tables that grew, copy the overlapping rows.
|
| Returns {"loaded": [...], "partial": [...], "skipped": [...]} for provenance."""
|
| own = model.state_dict()
|
| report = {"loaded": [], "partial": [], "skipped": []}
|
| with torch.no_grad():
|
| for k, v in state.items():
|
| if k not in own:
|
| report["skipped"].append(k)
|
| elif own[k].shape == v.shape:
|
| own[k].copy_(v)
|
| report["loaded"].append(k)
|
| elif own[k].ndim == v.ndim and all(a >= b for a, b in zip(own[k].shape, v.shape)):
|
| own[k][tuple(slice(0, n) for n in v.shape)].copy_(v)
|
| report["partial"].append(k)
|
| else:
|
| report["skipped"].append(k)
|
| return report
|
|
|
|
|
