| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| import math |
| from dataclasses import dataclass, asdict |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| |
|
|
| @dataclass |
| class DiTConfig: |
| in_channels: int = 16 |
| frames: int = 4 |
| height: int = 32 |
| width: int = 32 |
| patch: tuple = (1, 2, 2) |
| dim: int = 1024 |
| depth: int = 20 |
| heads: int = 16 |
| mlp_ratio: float = 4.0 |
| text_dim: int = 768 |
| attn: str = "softmax" |
| qk_norm: bool = True |
| frame_cond: bool = True |
| time_shift: float = 3.0 |
|
|
| @property |
| def tokens(self): |
| return (self.frames // self.patch[0]) * (self.height // self.patch[1]) * (self.width // self.patch[2]) |
|
|
| @property |
| def grid(self): |
| return (self.frames // self.patch[0], self.height // self.patch[1], self.width // self.patch[2]) |
|
|
|
|
| PRESETS = { |
| "tiny": dict(dim=128, depth=2, heads=4), |
| "small": dict(dim=512, depth=12, heads=8), |
| "v1": dict(dim=1024, depth=20, heads=16), |
| } |
|
|
|
|
| |
|
|
| def timestep_embedding(t, dim, max_period=10000): |
| """t in [0,1] -> sinusoidal [B, dim]. Scaled by 1000 so the frequencies match the usual 0..1000 diffusion range.""" |
| half = dim // 2 |
| freqs = torch.exp(-math.log(max_period) * torch.arange(half, device=t.device, dtype=torch.float32) / half) |
| args = (t.float() * 1000.0)[:, None] * freqs[None] |
| return torch.cat([args.cos(), args.sin()], dim=-1) |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim, eps=1e-6): |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x): |
| xf = x.float() |
| return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)).to(x.dtype) * self.weight |
|
|
|
|
| def rope_3d(grid, head_dim, theta=10000.0): |
| """cos/sin tables [L, head_dim/2] for a (T,H,W) token grid. head_dim/2 angles split 1:1.5:1.5 between t, h, w.""" |
| T, H, W = grid |
| n = head_dim // 2 |
| nt = (n // 4) if n >= 4 else 1 |
| nh = (n - nt) // 2 |
| nw = n - nt - nh |
| def axis(num, pos): |
| freqs = 1.0 / (theta ** (torch.arange(num, dtype=torch.float32) / num)) |
| return pos.float()[:, None] * freqs[None] |
| t_idx, h_idx, w_idx = torch.meshgrid(torch.arange(T), torch.arange(H), torch.arange(W), indexing="ij") |
| ang = torch.cat([axis(nt, t_idx.flatten()), axis(nh, h_idx.flatten()), axis(nw, w_idx.flatten())], dim=-1) |
| return ang.cos(), ang.sin() |
|
|
|
|
| def apply_rope(x, cos, sin): |
| """x [B, heads, L, d]; pairs (i, i + d/2) are rotated by the angle i.""" |
| x1, x2 = x.float().chunk(2, dim=-1) |
| out = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) |
| return out.to(x.dtype) |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, dim, heads, qk_norm, kind="softmax", ctx_dim=None): |
| super().__init__() |
| self.heads, self.hd, self.kind = heads, dim // heads, kind |
| ctx_dim = ctx_dim or dim |
| self.q = nn.Linear(dim, dim) |
| self.kv = nn.Linear(ctx_dim, 2 * dim) |
| self.o = nn.Linear(dim, dim) |
| self.qn = RMSNorm(self.hd) if qk_norm else nn.Identity() |
| self.kn = RMSNorm(self.hd) if qk_norm else nn.Identity() |
|
|
| def forward(self, x, ctx=None, ctx_mask=None, rope=None): |
| B, L, _ = x.shape |
| ctx = x if ctx is None else ctx |
| q = self.q(x).view(B, L, self.heads, self.hd).transpose(1, 2) |
| k, v = self.kv(ctx).view(B, ctx.shape[1], 2, self.heads, self.hd).permute(2, 0, 3, 1, 4) |
| q, k = self.qn(q), self.kn(k) |
| if rope is not None: |
| q, k = apply_rope(q, *rope), apply_rope(k, *rope) |
| if self.kind == "linear": |
| out = self._linear(q, k, v) |
| else: |
| mask = None |
| if ctx_mask is not None: |
| mask = ctx_mask[:, None, None, :] |
| out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) |
| return self.o(out.transpose(1, 2).reshape(B, L, -1)) |
|
|
| @staticmethod |
| def _linear(q, k, v, eps=1e-6): |
| """ReLU linear attention (SANA). fp32 for the sums; O(L d^2) instead of O(L^2 d).""" |
| q, k, v = F.relu(q.float()), F.relu(k.float()), v.float() |
| kv = torch.einsum("bhld,bhle->bhde", k, v) |
| z = 1.0 / (torch.einsum("bhld,bhd->bhl", q, k.sum(2)) + eps) |
| out = torch.einsum("bhld,bhde->bhle", q, kv) * z[..., None] |
| return out.to(v.dtype) if v.dtype != torch.float32 else out |
|
|
|
|
| class Block(nn.Module): |
| """adaLN-zero DiT block with a shared modulation table (PixArt-alpha): the 6 x dim modulation comes from one global |
| MLP on the timestep plus a per-block learned offset, instead of a 6 x dim^2 Linear per block.""" |
| def __init__(self, cfg): |
| super().__init__() |
| d = cfg.dim |
| self.n1 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6) |
| self.attn = Attention(d, cfg.heads, cfg.qk_norm, cfg.attn) |
| self.n2 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6) |
| self.cross = Attention(d, cfg.heads, cfg.qk_norm, "softmax", ctx_dim=cfg.text_dim) |
| self.n3 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6) |
| hidden = int(d * cfg.mlp_ratio) |
| self.mlp = nn.Sequential(nn.Linear(d, hidden), nn.GELU(approximate="tanh"), nn.Linear(hidden, d)) |
| self.table = nn.Parameter(torch.randn(6, d) / d ** 0.5) |
|
|
| def forward(self, x, t6, text, text_mask, rope): |
| |
| shift1, scale1, gate1, shift2, scale2, gate2 = (self.table[None] + t6).unbind(1) |
| h = self.n1(x) * (1 + scale1[:, None]) + shift1[:, None] |
| x = x + gate1[:, None] * self.attn(h, rope=rope) |
| x = x + self.cross(self.n2(x), ctx=text, ctx_mask=text_mask) |
| h = self.n3(x) * (1 + scale2[:, None]) + shift2[:, None] |
| x = x + gate2[:, None] * self.mlp(h) |
| return x |
|
|
|
|
| |
|
|
| class BoraDiT(nn.Module): |
| def __init__(self, cfg: DiTConfig): |
| super().__init__() |
| self.cfg = cfg |
| self.grad_checkpoint = False |
| d, p = cfg.dim, cfg.patch |
| in_ch = cfg.in_channels + (cfg.in_channels + 1 if cfg.frame_cond else 0) |
| self.patch_embed = nn.Conv3d(in_ch, d, kernel_size=p, stride=p) |
| self.t_embed = nn.Sequential(nn.Linear(256, d), nn.SiLU(), nn.Linear(d, d)) |
| self.t_mod = nn.Sequential(nn.SiLU(), nn.Linear(d, 6 * d)) |
| self.text_proj = nn.Sequential(nn.Linear(cfg.text_dim, cfg.text_dim), nn.GELU(approximate="tanh"), nn.Linear(cfg.text_dim, cfg.text_dim)) |
| 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_mod = nn.Sequential(nn.SiLU(), nn.Linear(d, 2 * d)) |
| self.final = nn.Linear(d, p[0] * p[1] * p[2] * cfg.in_channels) |
| if cfg.attn == "linear": |
| self.pos = nn.Parameter(torch.randn(1, cfg.tokens, d) * 0.02) |
| else: |
| cos, sin = rope_3d(cfg.grid, d // cfg.heads) |
| self.register_buffer("rope_cos", cos, persistent=False) |
| self.register_buffer("rope_sin", sin, persistent=False) |
| 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.data |
| nn.init.xavier_uniform_(w.view(w.shape[0], -1)) |
| nn.init.zeros_(self.patch_embed.bias) |
| nn.init.normal_(self.t_embed[0].weight, std=0.02) |
| nn.init.normal_(self.t_embed[2].weight, std=0.02) |
| nn.init.zeros_(self.t_mod[1].weight); nn.init.zeros_(self.t_mod[1].bias) |
| nn.init.zeros_(self.final_mod[1].weight); nn.init.zeros_(self.final_mod[1].bias) |
| nn.init.zeros_(self.final.weight); nn.init.zeros_(self.final.bias) |
| for b in self.blocks: |
| nn.init.zeros_(b.cross.o.weight); nn.init.zeros_(b.cross.o.bias) |
|
|
| def unpatchify(self, x): |
| B = x.shape[0] |
| T, H, W = self.cfg.grid |
| pt, ph, pw = self.cfg.patch |
| c = self.cfg.in_channels |
| x = x.view(B, T, H, W, pt, ph, pw, c) |
| x = x.permute(0, 7, 1, 4, 2, 5, 3, 6) |
| return x.reshape(B, c, T * pt, H * ph, W * pw) |
|
|
| def forward(self, x_t, t, text, text_mask, cond=None, cond_mask=None): |
| """x_t [B,16,F,H,W], t [B] in [0,1], text [B,L,768], text_mask [B,L] bool, |
| cond [B,16,F,H,W] clean frames (zeros where unknown), cond_mask [B,1,F,1,1] in {0,1}. Returns v [B,16,F,H,W].""" |
| B = x_t.shape[0] |
| if self.cfg.frame_cond: |
| if cond is None: |
| cond = torch.zeros_like(x_t) |
| cond_mask = torch.zeros(B, 1, x_t.shape[2], 1, 1, device=x_t.device, dtype=x_t.dtype) |
| m = cond_mask.expand(B, 1, *x_t.shape[2:]).to(x_t.dtype) |
| x_in = torch.cat([x_t, cond * m, m], dim=1) |
| else: |
| x_in = x_t |
| x = self.patch_embed(x_in).flatten(2).transpose(1, 2) |
| if self.cfg.attn == "linear": |
| x = x + self.pos |
| rope = None |
| else: |
| rope = (self.rope_cos, self.rope_sin) |
| te = self.t_embed(timestep_embedding(t, 256).to(x.dtype)) |
| t6 = self.t_mod(te).view(B, 6, -1) |
| text = self.text_proj(text.to(x.dtype)) |
| for blk in self.blocks: |
| if self.grad_checkpoint and self.training: |
| x = torch.utils.checkpoint.checkpoint(blk, x, t6, text, text_mask, rope, use_reentrant=False) |
| else: |
| x = blk(x, t6, text, text_mask, rope) |
| shift, scale = self.final_mod(te).chunk(2, dim=-1) |
| x = self.final_norm(x) * (1 + scale[:, None]) + shift[:, None] |
| return self.unpatchify(self.final(x)) |
|
|
|
|
| |
|
|
| def shift_t(u, shift): |
| """SD3 timestep shift: pushes samples toward t=1 (more noise). Identity for shift=1.""" |
| return shift * u / (1 + (shift - 1) * u) |
|
|
|
|
| def sample_timesteps(B, shift, device): |
| u = torch.sigmoid(torch.randn(B, device=device)) |
| return shift_t(u, shift) |
|
|
|
|
| def make_frame_cond(x, p_cond, generator=None): |
| """With probability p_cond per sample, give the clean first frame of the block as conditioning. |
| Returns cond [B,16,F,H,W], cond_mask [B,1,F,1,1], loss_mask [B,1,F,1,1] (0 on the given frame).""" |
| B, C, Fr, H, W = x.shape |
| use = (torch.rand(B, device=x.device, generator=generator) < p_cond).to(x.dtype) |
| cond_mask = torch.zeros(B, 1, Fr, 1, 1, device=x.device, dtype=x.dtype) |
| cond_mask[:, :, 0] = use.view(B, 1, 1, 1) |
| cond = torch.zeros_like(x) |
| cond[:, :, 0] = x[:, :, 0] |
| return cond, cond_mask, 1 - cond_mask |
|
|
|
|
| def flow_loss(model, x, text, text_mask, p_cond=0.5, shift=3.0, t=None, noise=None): |
| """Rectified-flow MSE. x: clean block [B,16,F,H,W]. Returns (loss, dict of parts).""" |
| B = x.shape[0] |
| t = sample_timesteps(B, shift, x.device) if t is None else t |
| noise = torch.randn_like(x) if noise is None else noise |
| tt = t.view(B, 1, 1, 1, 1).to(x.dtype) |
| x_t = (1 - tt) * x + tt * noise |
| cond, cond_mask, loss_mask = make_frame_cond(x, p_cond) |
| v = model(x_t, t, text, text_mask, cond, cond_mask) |
| target = noise - x |
| err = (v.float() - target.float()) ** 2 * loss_mask |
| loss = err.sum() / (loss_mask.sum() * x.shape[1] * x.shape[3] * x.shape[4]) |
| return loss, {"t_mean": float(t.mean()), "cond_frac": float(cond_mask[:, :, 0].mean())} |
|
|
|
|
| @torch.no_grad() |
| def sample_block(model, text, text_mask, null_text, null_mask, steps=8, cfg_scale=4.0, shift=3.0, |
| first_frame=None, generator=None, dtype=torch.bfloat16): |
| """Euler sampler for one block. first_frame [B,16,H,W] (clean) continues a previous block; None starts fresh. |
| Returns clean latents [B,16,F,H,W] (float32).""" |
| c = model.cfg |
| B = text.shape[0] |
| dev = text.device |
| x = torch.randn(B, c.in_channels, c.frames, c.height, c.width, device=dev, generator=generator) |
| cond = cond_mask = None |
| eps0 = None |
| if first_frame is not None: |
| cond = torch.zeros_like(x) |
| cond[:, :, 0] = first_frame |
| cond_mask = torch.zeros(B, 1, c.frames, 1, 1, device=dev) |
| cond_mask[:, :, 0] = 1 |
| eps0 = torch.randn(B, c.in_channels, c.height, c.width, device=dev, generator=generator) |
| ts = shift_t(torch.linspace(1, 0, steps + 1, device=dev), shift) |
| for i in range(steps): |
| t, t_next = ts[i], ts[i + 1] |
| if first_frame is not None: |
| x[:, :, 0] = (1 - t) * first_frame + t * eps0 |
| tb = t.expand(B) |
| with torch.autocast(device_type=dev.type, dtype=dtype, enabled=dev.type == "cuda"): |
| if cfg_scale != 1.0: |
| xx = torch.cat([x, x]); tt = torch.cat([tb, tb]) |
| te = torch.cat([text, null_text]); tm = torch.cat([text_mask, null_mask]) |
| cc = None if cond is None else torch.cat([cond, cond]) |
| cm = None if cond_mask is None else torch.cat([cond_mask, cond_mask]) |
| v_c, v_u = model(xx.to(dtype), tt, te, tm, cc if cc is None else cc.to(dtype), cm if cm is None else cm.to(dtype)).float().chunk(2) |
| v = v_u + cfg_scale * (v_c - v_u) |
| else: |
| v = model(x.to(dtype), tb, text, text_mask, None if cond is None else cond.to(dtype), None if cond_mask is None else cond_mask.to(dtype)).float() |
| x = x + (t_next - t) * v |
| if first_frame is not None: |
| x[:, :, 0] = first_frame |
| return x |
|
|
|
|
| @torch.no_grad() |
| def sample_video(model, text, text_mask, null_text, null_mask, n_blocks=3, overlap=1, **kw): |
| """Chain blocks with `overlap` shared frames: 4 + 3 + 3 = 10 latent frames for n_blocks=3 (one full 37-frame clip).""" |
| blocks = [] |
| first = None |
| for _ in range(n_blocks): |
| blk = sample_block(model, text, text_mask, null_text, null_mask, first_frame=first, **kw) |
| blocks.append(blk if first is None else blk[:, :, overlap:]) |
| first = blk[:, :, -1] |
| return torch.cat(blocks, dim=2) |
|
|
|
|
| |
|
|
| def build(preset="v1", **overrides): |
| cfg = DiTConfig(**{**PRESETS[preset], **overrides}) |
| return BoraDiT(cfg), cfg |
|
|
|
|
| def count_params(model): |
| return sum(p.numel() for p in model.parameters()) |
|
|
|
|
| def config_dict(cfg): |
| d = asdict(cfg) |
| d["patch"] = list(d["patch"]) |
| return d |
|
|
