code/noncausal_vit_decoder.py

"""General (non-causal) fast H3 VAE decoder: the released H3 VAE ViT decoder pruned to 24 layers, decoding whole
native 5-latent blocks with window attention and no tiling / temporal-overlap recomputation.

* time: native 5-latent H3 blocks (17 video frames). Attention is bidirectional inside a temporal window of 5
  latent frames. Odd layers shift the temporal windows by 2 latent frames (Video-Swin style), so neighbouring
  blocks exchange information inside the network; nothing is recomputed and nothing is cross-faded.
* space: 16x16-token windows (256 px, the released tile size); odd layers shift them by half a window.
* global tokens: every native 5-latent block owns 5 global tokens (4 register + 1 class, from the released decoder).
  They attend to every token of their block; every window attends to the global tokens of the blocks it overlaps.
* rotary positions: spatial ids are window-local (as in the released tile decode); time ids are window-local too,
  so nothing depends on the absolute position in the video and any length decodes the same way.
* output: per token a 3 x 4 x 16 x 16 pixel patch; the first latent of every block keeps 1 of its 4 frames,
  exactly as the released decoder, so 5 latent frames give 17 video frames.

Input z: normalized latents (B, T, 24, H, W), t = 0 is the start of a native block. Output: decode_base-space
pixels (B, 3, F, 16H, 16W); apply vae.processor.revert_tensor for RGB.
"""
from __future__ import annotations

import copy

import torch
from torch import nn
from torch.nn import functional as F

from spark_compat import _apply_qk_norm, _apply_rotary_pos_emb_impl

PATCH, PATCH_T, OUT_CH = 16, 4, 3
KEEP_24 = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 19, 20, 21, 31, 32, 35]
# greedy single-layer removal on the released 36-layer decoder (vae_prune/layer_scan.py)
KEEP_16 = [0, 1, 2, 3, 5, 6, 8, 9, 10, 11, 14, 15, 20, 31, 32, 35]
KEEP_12 = [0, 1, 2, 3, 9, 10, 14, 15, 20, 31, 32, 35]
KEEP_8 = [0, 1, 2, 15, 20, 31, 32, 35]
KEEP_6 = [0, 1, 2, 15, 20, 35]          # KEEP_8 minus the two layers with the smallest single-removal PSNR drop
KEEP = {24: KEEP_24, 16: KEEP_16, 12: KEEP_12, 8: KEEP_8, 6: KEEP_6}


def shrink_ffn(ff, idx):
    """Keep the SwiGLU hidden channels idx: gate rows idx, up rows I+idx of w1 and columns idx of w2."""
    I = ff.w2.weight.shape[1]
    rows = torch.cat([idx, idx + I])
    w1 = nn.Linear(ff.w1.in_features, 2 * len(idx), bias=ff.w1.bias is not None).to(ff.w1.weight)
    w2 = nn.Linear(len(idx), ff.w2.out_features, bias=ff.w2.bias is not None).to(ff.w2.weight)
    with torch.no_grad():
        w1.weight.copy_(ff.w1.weight[rows]); w2.weight.copy_(ff.w2.weight[:, idx])
        if ff.w1.bias is not None: w1.bias.copy_(ff.w1.bias[rows])
        if ff.w2.bias is not None: w2.bias.copy_(ff.w2.bias)
    ff.w1, ff.w2 = w1, w2


class NonCausalViTDecoder(nn.Module):
    def __init__(self, vae, keep_layers=KEEP_24, *, window: int = 16, block: int = 5, t_shift: int = 2,
                 frames_to_trim: int = 3, time_norm: float | None = None, ffn_keep: int | None = None):
        super().__init__()
        dec = vae.decoder
        self.keep_layers = list(keep_layers)
        self.window, self.block, self.t_shift, self.frames_to_trim = window, block, t_shift, frames_to_trim
        # time ids: (window-local index + 0.5) / time_norm * 2 - 1. The released decoder decodes 7-latent chunks
        # (5 + 2 overlap) with length-normalized ids, i.e. time_norm = 7; default = block keeps the v1 behaviour.
        self.time_norm = float(time_norm if time_norm is not None else block)
        self.post_quant_conv = copy.deepcopy(vae.post_quant_conv)
        self.x_embedder = copy.deepcopy(dec.x_embedder)
        self.register_tokens = nn.Parameter(dec.register_tokens.detach().clone())
        self.blocks = nn.ModuleList([copy.deepcopy(dec.transformer_blocks[i]) for i in self.keep_layers])
        self.norm_out = copy.deepcopy(dec.norm_out)
        self.proj_out = copy.deepcopy(dec.proj_out)
        self.pos_embed = copy.deepcopy(dec.pos_embed)
        self.heads = self.blocks[0].attn.heads
        self.dh = self.blocks[0].attn.dim_head
        self.register_buffer("latents_mean", vae.latents_mean.detach().clone().float(), persistent=False)
        self.register_buffer("latents_std", vae.latents_std.detach().clone().float(), persistent=False)
        for p in self.pos_embed.parameters():
            p.requires_grad_(False)
        self.grad_checkpointing = False      # per-layer activation checkpointing (training only)
        if ffn_keep is not None:             # placeholder channels; prune_ffn picks the real ones before training
            for blk in self.blocks:
                shrink_ffn(blk.ff, torch.arange(ffn_keep, device=blk.ff.w2.weight.device))

    def prune_ffn(self, vae, scores, keep: int, src_model=None):
        """Rebuild every FFN with the `keep` highest-scoring hidden channels of the source FFN.
        Source: the full-width layers of src_model (e.g. a trained decoder), else the released VAE layers.
        scores[li]: (I,) importance of the source layer's hidden channels."""
        for li, blk in enumerate(self.blocks):
            ff = src_model.blocks[li].ff if src_model is not None else vae.decoder.transformer_blocks[self.keep_layers[li]].ff
            src = copy.deepcopy(ff).to(blk.ff.w2.weight.device)
            idx = torch.topk(scores[li].to(src.w2.weight.device), keep).indices.sort().values
            shrink_ffn(src, idx)
            blk.ff = src

    # ------------------------------------------------------------------ geometry
    def _shifts(self, li):
        odd = li % 2 == 1
        return (self.window // 2 if odd else 0), (self.t_shift if odd else 0)

    def _rope(self, T, H, W, s, st, device):
        """cos/sin for tokens in (T, H, W) order with window-local time and space ids for shifts (st, s)."""
        lt = ((torch.arange(T, device=device) + st) % self.block).float()
        ly = ((torch.arange(H, device=device) + s) % self.window).float()
        lx = ((torch.arange(W, device=device) + s) % self.window).float()
        t = (lt + 0.5) / self.time_norm * 2 - 1
        hy = (ly + 0.5) / self.window * 2 - 1
        wx = (lx + 0.5) / self.window * 2 - 1
        grid = torch.stack(torch.meshgrid(t, hy, wx, indexing="ij"), -1).reshape(1, -1, 3)
        return self.pos_embed(grid.float())

    def _partition(self, x, s, st):
        """x: (B, T, H, W, ...) -> windows (B*nW, Tw*ws*ws, ...), valid (nW, Tw*ws*ws), info.
        Window order is (time window, row, col)."""
        B, T, H, W = x.shape[:4]
        ws, tb = self.window, self.block
        rest = x.shape[4:]
        Tp = ((T + st + tb - 1) // tb) * tb
        Hp = ((H + s + ws - 1) // ws) * ws
        Wp = ((W + s + ws - 1) // ws) * ws
        xp = x.new_zeros(B, Tp, Hp, Wp, *rest)
        xp[:, st:st + T, s:s + H, s:s + W] = x
        valid = torch.zeros(Tp, Hp, Wp, dtype=torch.bool, device=x.device)
        valid[st:st + T, s:s + H, s:s + W] = True
        nt, nh, nw = Tp // tb, Hp // ws, Wp // ws
        r = len(rest)
        xw = xp.view(B, nt, tb, nh, ws, nw, ws, *rest).permute(0, 1, 3, 5, 2, 4, 6, *range(7, 7 + r))
        xw = xw.reshape(B * nt * nh * nw, tb * ws * ws, *rest)
        vw = valid.view(nt, tb, nh, ws, nw, ws).permute(0, 2, 4, 1, 3, 5).reshape(nt * nh * nw, tb * ws * ws)
        return xw, vw, (B, T, H, W, Tp, Hp, Wp, nt, nh, nw)

    def _unpartition(self, xw, info, s, st):
        B, T, H, W, Tp, Hp, Wp, nt, nh, nw = info
        ws, tb = self.window, self.block
        rest = xw.shape[2:]
        r = len(rest)
        x = xw.view(B, nt, nh, nw, tb, ws, ws, *rest).permute(0, 1, 4, 2, 5, 3, 6, *range(7, 7 + r))
        x = x.reshape(B, Tp, Hp, Wp, *rest)
        return x[:, st:st + T, s:s + H, s:s + W]

    def _layer(self, li, blk, x, g):
        """One transformer layer on the tokens x (B,T,H,W,D) and the per-block global tokens g (B,nb,S,D)."""
        B, T, H, W, nb, S = self._geom
        dev = x.device
        tb = self.block
        s, st = self._shifts(li)
        h = blk.norm1(x.float()).to(x.dtype)
        hg = blk.norm1(g.float()).to(g.dtype)
        q, k, v = blk.attn.to_qkv(h).view(B, T, H, W, self.heads, 3 * self.dh).chunk(3, -1)
        qg, kg, vg = blk.attn.to_qkv(hg).view(B, nb, S, self.heads, 3 * self.dh).chunk(3, -1)
        if blk.attn.norm_q is not None:
            q, qg = _apply_qk_norm(blk.attn.norm_q, q), _apply_qk_norm(blk.attn.norm_q, qg)
            k, kg = _apply_qk_norm(blk.attn.norm_k, k), _apply_qk_norm(blk.attn.norm_k, kg)
        cos, sin = self._rope(T, H, W, s, st, dev)
        q = _apply_rotary_pos_emb_impl(q.reshape(B, T * H * W, self.heads, self.dh), (cos, sin)).view(B, T, H, W, self.heads, self.dh)
        k = _apply_rotary_pos_emb_impl(k.reshape(B, T * H * W, self.heads, self.dh), (cos, sin)).view(B, T, H, W, self.heads, self.dh)
        # ---- window attention: tokens of a window + global tokens of the native blocks the window overlaps
        qw, vmask, info = self._partition(q, s, st)
        kw, _, _ = self._partition(k, s, st)
        vw, _, _ = self._partition(v, s, st)
        nt, nh, nw = info[7:]
        nW = nt * nh * nw
        # native blocks overlapped by temporal window j: latent range [j*tb - st, (j+1)*tb - st)
        tw_lo = (torch.arange(nt, device=dev) * tb - st).clamp(min=0)
        tw_hi = ((torch.arange(nt, device=dev) + 1) * tb - st - 1).clamp(max=T - 1)
        b0, b1 = tw_lo // tb, tw_hi // tb                         # at most 2 blocks per window
        gk_idx = torch.stack([b0, b1], 1)                          # nt, 2
        gk_valid = torch.stack([torch.ones_like(b0, dtype=torch.bool), b1 != b0], 1)   # nt, 2
        kgw = kg[:, gk_idx].reshape(B, nt, 2 * S, self.heads, self.dh)
        vgw = vg[:, gk_idx].reshape(B, nt, 2 * S, self.heads, self.dh)
        kgw = kgw[:, :, None, None].expand(B, nt, nh, nw, 2 * S, self.heads, self.dh).reshape(B * nW, 2 * S, self.heads, self.dh)
        vgw = vgw[:, :, None, None].expand(B, nt, nh, nw, 2 * S, self.heads, self.dh).reshape(B * nW, 2 * S, self.heads, self.dh)
        gmask = gk_valid[:, :, None].expand(nt, 2, S).reshape(nt, 2 * S)
        gmask = gmask[:, None, None].expand(nt, nh, nw, 2 * S).reshape(nW, 2 * S)
        Kc = torch.cat([kw, kgw], 1).transpose(1, 2)
        Vc = torch.cat([vw, vgw], 1).transpose(1, 2)
        kmask = torch.cat([vmask, gmask], 1).repeat(B, 1)[:, None, None, :]
        ow = F.scaled_dot_product_attention(qw.transpose(1, 2), Kc, Vc, attn_mask=kmask).transpose(1, 2)
        o = self._unpartition(ow, info, s, st).reshape(B, T, H, W, self.heads * self.dh)
        # ---- global tokens: attend to every token of their own native block + the block's global tokens
        Tp = nb * tb
        kb = k.new_zeros(B, Tp, H, W, self.heads, self.dh); kb[:, :T] = k
        vb = v.new_zeros(B, Tp, H, W, self.heads, self.dh); vb[:, :T] = v
        kb = kb.view(B, nb, tb * H * W, self.heads, self.dh)
        vb = vb.view(B, nb, tb * H * W, self.heads, self.dh)
        tvalid = torch.arange(Tp, device=dev) < T
        bmask = tvalid.view(nb, tb, 1).expand(nb, tb, H * W).reshape(nb, tb * H * W)
        Kg = torch.cat([kb, kg], 2).reshape(B * nb, tb * H * W + S, self.heads, self.dh).transpose(1, 2)
        Vg = torch.cat([vb, vg], 2).reshape(B * nb, tb * H * W + S, self.heads, self.dh).transpose(1, 2)
        gm = torch.cat([bmask, torch.ones(nb, S, dtype=torch.bool, device=dev)], 1).repeat(B, 1)[:, None, None, :]
        og = F.scaled_dot_product_attention(qg.reshape(B * nb, S, self.heads, self.dh).transpose(1, 2), Kg, Vg,
                                            attn_mask=gm).transpose(1, 2).reshape(B, nb, S, -1)
        x = x + blk.attn.to_out(o) * blk.scale1
        g = g + blk.attn.to_out(og) * blk.scale1
        x = x + blk.ff(blk.norm2(x.float()).to(x.dtype)) * blk.scale2
        g = g + blk.ff(blk.norm2(g.float()).to(g.dtype)) * blk.scale2
        return x, g

    # ------------------------------------------------------------------ forward
    def forward(self, z):
        B, T, _, H, W = z.shape
        dev = z.device
        tb, ws = self.block, self.window
        nb = (T + tb - 1) // tb                                     # native blocks
        compute_dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled("cuda") else z.dtype
        raw = z.float() * self.latents_std.view(1, 1, -1, 1, 1) + self.latents_mean.view(1, 1, -1, 1, 1)
        raw = self.post_quant_conv(raw.transpose(1, 2))
        with torch.autocast("cuda", enabled=False):
            x = self.x_embedder(raw.permute(0, 2, 3, 4, 1).float())   # B,T,H,W,D
        x = x.to(compute_dtype)
        D = x.shape[-1]
        S = self.register_tokens.shape[1] + 1
        g = torch.cat([self.register_tokens.expand(B, -1, -1), x.new_zeros(B, 1, D)], 1).to(compute_dtype)
        g = g.unsqueeze(1).expand(B, nb, S, D).contiguous()          # B, nb, S, D  (one set per native block)
        self._geom = (B, T, H, W, nb, S)
        use_ckpt = self.grad_checkpointing and torch.is_grad_enabled() and self.training
        for li, blk in enumerate(self.blocks):
            if use_ckpt:
                x, g = torch.utils.checkpoint.checkpoint(self._layer, li, blk, x, g, use_reentrant=False)
            else:
                x, g = self._layer(li, blk, x, g)
        x = self.norm_out(x)
        with torch.autocast("cuda", enabled=False):
            out = self.proj_out(x.float())                                  # B,T,H,W,3*4*16*16
        out = out.view(B, T, H, W, OUT_CH, PATCH_T, PATCH, PATCH).permute(0, 4, 1, 5, 2, 6, 3, 7)
        out = out.reshape(B, OUT_CH, T * PATCH_T, H * PATCH, W * PATCH)
        keep = torch.ones(T * PATCH_T, dtype=torch.bool, device=dev)
        for t in range(0, T, tb):
            keep[t * PATCH_T: t * PATCH_T + self.frames_to_trim] = False
        return out[:, :, keep]