lance_mlx/vae_wan22.py

"""MLX port of Wan 2.2 VAE (used as Lance's video VAE).

Source: bytedance/Lance modeling/vae/wan/vae2_2.py
       (which mirrors Alibaba Wan 2.2's open-source VAE)

Layout convention: MLX uses (B, T, H, W, C) for Conv3d. PT uses (B, C, T, H, W).
All internal tensors here use the MLX NTHWC layout. Conv weights from the PT
checkpoint must be reshaped (out_C, in_C, kT, kH, kW) -> (out_C, kT, kH, kW, in_C)
at load time — see tools/convert_weights.py.

Caching: the PT model uses a feat_cache list to stream long videos in chunks.
This first MLX port runs in single-pass mode (whole tensor at once), which is
correct for images (T=1) and short videos. Streaming-cache mode is a follow-up.
"""
from __future__ import annotations

from typing import List, Optional

import mlx.core as mx
import mlx.nn as nn
from einops import rearrange


CACHE_T = 2  # matches PT constant — only used when streaming is enabled


# ---------------------------------------------------------------------------
# CausalConv3d — temporally-causal 3D convolution
# ---------------------------------------------------------------------------

# Number of trailing time frames to retain for streaming-cache continuity.
CACHE_T = 2


def _cache_pre(x: mx.array, feat_cache, feat_idx):
    """Streaming-cache helper for CausalConv3d calls.

    Returns (prev_cache, new_cache, idx). The caller passes prev_cache to the
    conv (as cache_x), then stores new_cache at feat_cache[idx] and bumps
    feat_idx. Inlining this avoids repeating the boilerplate at every conv.
    """
    if feat_cache is None:
        return None, None, None
    idx = feat_idx[0]
    # last CACHE_T frames on time axis (NTHWC -> axis 1)
    cache_x = x[:, -CACHE_T:, ...]
    prev = feat_cache[idx]
    is_rep = isinstance(prev, str) and prev == "Rep"
    has_prev = (prev is not None) and (not is_rep)
    if cache_x.shape[1] < 2 and has_prev:
        # prepend the last frame of the previous cache so the new cache has T>=2
        cache_x = mx.concatenate([prev[:, -1:, ...], cache_x], axis=1)
    feat_idx[0] = idx + 1
    return (prev if has_prev else None), cache_x, idx


def _set_cache(feat_cache, idx, value):
    if feat_cache is not None and idx is not None:
        feat_cache[idx] = value


class CausalConv3d(nn.Conv3d):
    """3D conv with causal padding on the temporal dimension.

    Inherits from nn.Conv3d so parameter names (weight, bias) match the PT
    checkpoint directly — PT's CausalConv3d also subclasses nn.Conv3d.

    PT pads time as (2*pt, 0) so each output frame only depends on current +
    previous frames. Spatial padding is symmetric.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size,
        stride=1,
        padding=0,
        bias: bool = True,
    ):
        if isinstance(kernel_size, int):
            kernel_size = (kernel_size, kernel_size, kernel_size)
        if isinstance(stride, int):
            stride = (stride, stride, stride)
        if isinstance(padding, int):
            padding = (padding, padding, padding)

        # Pass padding=0 to base — we do explicit pad before forward.
        super().__init__(
            in_channels, out_channels,
            kernel_size=kernel_size, stride=stride, padding=0, bias=bias,
        )
        self._pt, self._ph, self._pw = padding

    def __call__(self, x: mx.array, cache_x: Optional[mx.array] = None) -> mx.array:
        pt = 2 * self._pt
        if cache_x is not None and pt > 0:
            x = mx.concatenate([cache_x, x], axis=1)
            pt = max(0, pt - cache_x.shape[1])
        if pt > 0 or self._ph > 0 or self._pw > 0:
            x = mx.pad(
                x,
                [(0, 0), (pt, 0), (self._ph, self._ph), (self._pw, self._pw), (0, 0)],
            )
        return super().__call__(x)


# ---------------------------------------------------------------------------
# RMS_norm — channel-axis RMS normalization with learnable scale
# ---------------------------------------------------------------------------

class RMS_norm(nn.Module):
    """Equivalent to PT version: F.normalize(x, dim=channel) * sqrt(C) * gamma + bias.

    In NTHWC layout, channel is always axis -1. The PT version has a
    `channel_first` flag and a `images` flag controlling broadcast shape;
    we collapse these — gamma is shape (C,) and broadcasts trivially on -1.
    """

    def __init__(self, dim: int, bias: bool = False):
        super().__init__()
        self.scale = dim ** 0.5
        self.gamma = mx.ones((dim,))
        self.bias = mx.zeros((dim,)) if bias else None

    def __call__(self, x: mx.array) -> mx.array:
        # L2 normalize along channel (-1)
        norm = mx.sqrt(mx.sum(x * x, axis=-1, keepdims=True) + 1e-12)
        x = x / norm
        out = x * self.scale * self.gamma
        if self.bias is not None:
            out = out + self.bias
        return out


# ---------------------------------------------------------------------------
# Resample — 2D/3D up/down sample blocks
# ---------------------------------------------------------------------------

class Resample(nn.Module):
    """Spatial (and optional temporal) resampling.

    Modes:
      none         — passthrough
      upsample2d   — nearest 2x spatial + 3x3 conv2d refine
      upsample3d   — upsample2d + temporal expand (2x via channel-doubling conv)
      downsample2d — zero-pad + stride-2 conv2d
      downsample3d — downsample2d + stride-2 temporal conv

    First implementation: single-pass only (no feat_cache streaming).
    """

    def __init__(self, dim: int, mode: str):
        super().__init__()
        assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d")
        self.dim = dim
        self.mode = mode

        if mode == "upsample2d":
            self.up = nn.Upsample(scale_factor=(2.0, 2.0), mode="nearest")
            self.spatial_conv = nn.Conv2d(dim, dim, 3, padding=1)
        elif mode == "upsample3d":
            self.up = nn.Upsample(scale_factor=(2.0, 2.0), mode="nearest")
            self.spatial_conv = nn.Conv2d(dim, dim, 3, padding=1)
            # doubles channels — second half becomes the new temporal frames
            self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
        elif mode == "downsample2d":
            self.spatial_conv = nn.Conv2d(dim, dim, 3, stride=(2, 2), padding=0)
        elif mode == "downsample3d":
            self.spatial_conv = nn.Conv2d(dim, dim, 3, stride=(2, 2), padding=0)
            self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))

    def _spatial_pad_zeropad2d(self, x: mx.array) -> mx.array:
        # PT ZeroPad2d((0,1,0,1)) -> pad right=1, bottom=1 in (H,W).
        # x is (BT, H, W, C). Pad axes 1 (H bottom) and 2 (W right).
        return mx.pad(x, [(0, 0), (0, 1), (0, 1), (0, 0)])

    def __call__(self, x: mx.array, feat_cache=None, feat_idx=None,
                 first_chunk: bool = True) -> mx.array:
        """
        x: (B, T, H, W, C)

        Single-pass mode (feat_cache is None): controlled by first_chunk —
        time_conv is skipped on the first chunk (no prior cache frame).

        Streaming mode (feat_cache is a list): each upsample3d / downsample3d
        time_conv consumes feat_cache[idx] and updates it.
        """
        b, t, h, w, c = x.shape

        if self.mode == "upsample3d":
            if feat_cache is None:
                if not first_chunk:
                    x = self.time_conv(x)
                    x = x.reshape(b, t, h, w, 2, c)
                    x = mx.stack([x[..., 0, :], x[..., 1, :]], axis=2)
                    x = x.reshape(b, t * 2, h, w, c)
                    t = t * 2
            else:
                idx = feat_idx[0]
                if feat_cache[idx] is None:
                    # First chunk: mark sentinel and bump
                    feat_cache[idx] = "Rep"
                    feat_idx[0] = idx + 1
                else:
                    prev, new_cache, _ = _cache_pre(x, feat_cache, feat_idx)
                    # Same gotcha as PT: if prev is "Rep" sentinel, pass None to time_conv
                    is_rep = isinstance(prev, str) and prev == "Rep"
                    x = self.time_conv(x, cache_x=None if is_rep else prev)
                    feat_cache[idx] = new_cache
                    x = x.reshape(b, t, h, w, 2, c)
                    x = mx.stack([x[..., 0, :], x[..., 1, :]], axis=2)
                    x = x.reshape(b, t * 2, h, w, c)
                    t = t * 2

        if self.mode in ("upsample2d", "upsample3d"):
            x = rearrange(x, "b t h w c -> (b t) h w c")
            x = self.up(x)
            x = self.spatial_conv(x)
            x = rearrange(x, "(b t) h w c -> b t h w c", b=b, t=t)
        elif self.mode in ("downsample2d", "downsample3d"):
            x = rearrange(x, "b t h w c -> (b t) h w c")
            x = self._spatial_pad_zeropad2d(x)
            x = self.spatial_conv(x)
            x = rearrange(x, "(b t) h w c -> b t h w c", b=b, t=t)

        if self.mode == "downsample3d":
            if feat_cache is None:
                if not first_chunk:
                    x = self.time_conv(x)
            else:
                idx = feat_idx[0]
                if feat_cache[idx] is None:
                    # PT: store current x as the cache, no time_conv yet
                    feat_cache[idx] = x
                    feat_idx[0] = idx + 1
                else:
                    cache_x = x[:, -1:, ...]
                    # PT: prepend the last frame of the previous cache to x, then time_conv
                    x = self.time_conv(mx.concatenate([feat_cache[idx][:, -1:, ...], x], axis=1))
                    feat_cache[idx] = cache_x
                    feat_idx[0] = idx + 1

        return x


# ---------------------------------------------------------------------------
# ResidualBlock — 3D ResBlock with RMS + SiLU + CausalConv3d
# ---------------------------------------------------------------------------

class ResidualBlock(nn.Module):

    def __init__(self, in_dim: int, out_dim: int, dropout: float = 0.0):
        super().__init__()
        self.in_dim = in_dim
        self.out_dim = out_dim

        self.norm1 = RMS_norm(in_dim)
        self.conv1 = CausalConv3d(in_dim, out_dim, 3, padding=1)
        self.norm2 = RMS_norm(out_dim)
        self.conv2 = CausalConv3d(out_dim, out_dim, 3, padding=1)

        if in_dim != out_dim:
            self.shortcut = CausalConv3d(in_dim, out_dim, 1)
        else:
            self.shortcut = None

    def __call__(self, x: mx.array, feat_cache=None, feat_idx=None) -> mx.array:
        h = x if self.shortcut is None else self.shortcut(x)
        x = self.norm1(x)
        x = nn.silu(x)
        prev, new_cache, idx = _cache_pre(x, feat_cache, feat_idx)
        x = self.conv1(x, cache_x=prev)
        _set_cache(feat_cache, idx, new_cache)
        x = self.norm2(x)
        x = nn.silu(x)
        prev, new_cache, idx = _cache_pre(x, feat_cache, feat_idx)
        x = self.conv2(x, cache_x=prev)
        _set_cache(feat_cache, idx, new_cache)
        return x + h


# ---------------------------------------------------------------------------
# AttentionBlock — single-head spatial self-attention applied per frame
# ---------------------------------------------------------------------------

class AttentionBlock(nn.Module):

    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim
        self.norm = RMS_norm(dim)
        # PT uses Conv2d 1x1 — same as Linear over channels. Use Conv2d for
        # parameter-name compatibility with the checkpoint.
        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
        self.proj = nn.Conv2d(dim, dim, 1)

    def __call__(self, x: mx.array) -> mx.array:
        # x: (B, T, H, W, C)
        b, t, h, w, c = x.shape
        identity = x
        x = rearrange(x, "b t h w c -> (b t) h w c")
        x = self.norm(x)
        qkv = self.to_qkv(x)                                       # (BT, H, W, 3C)
        qkv = qkv.reshape(b * t, h * w, 3, c)
        q, k, v = qkv[:, :, 0, :], qkv[:, :, 1, :], qkv[:, :, 2, :]   # each (BT, HW, C)
        # Add head dim of size 1 for SDPA
        q = q[:, None, :, :]                                       # (BT, 1, HW, C)
        k = k[:, None, :, :]
        v = v[:, None, :, :]
        out = mx.fast.scaled_dot_product_attention(q, k, v, scale=c ** -0.5)
        out = out[:, 0, :, :]                                      # (BT, HW, C)
        out = out.reshape(b * t, h, w, c)
        out = self.proj(out)
        out = rearrange(out, "(b t) h w c -> b t h w c", b=b, t=t)
        return out + identity


# ---------------------------------------------------------------------------
# patchify / unpatchify — 2x spatial pixel-shuffle in/out
# ---------------------------------------------------------------------------

def patchify(x: mx.array, patch_size: int) -> mx.array:
    if patch_size == 1:
        return x
    # NTHWC: (B, T, H, W, C) -> (B, T, H//p, W//p, C*p*p)
    if x.ndim == 4:                                            # (B, H, W, C)
        return rearrange(x, "b (h q) (w r) c -> b h w (c r q)", q=patch_size, r=patch_size)
    elif x.ndim == 5:                                          # (B, T, H, W, C)
        return rearrange(x, "b t (h q) (w r) c -> b t h w (c r q)", q=patch_size, r=patch_size)
    raise ValueError(f"patchify: invalid ndim {x.ndim}")


def unpatchify(x: mx.array, patch_size: int) -> mx.array:
    if patch_size == 1:
        return x
    if x.ndim == 4:
        return rearrange(x, "b h w (c r q) -> b (h q) (w r) c", q=patch_size, r=patch_size)
    elif x.ndim == 5:
        return rearrange(x, "b t h w (c r q) -> b t (h q) (w r) c", q=patch_size, r=patch_size)
    raise ValueError(f"unpatchify: invalid ndim {x.ndim}")


# ---------------------------------------------------------------------------
# AvgDown3D / DupUp3D — average-pool downsample / repeat-then-shuffle upsample
# (used as the shortcut path in Down_/Up_ResidualBlock)
# ---------------------------------------------------------------------------

class AvgDown3D(nn.Module):

    def __init__(self, in_channels: int, out_channels: int, factor_t: int, factor_s: int = 1):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.factor_t = factor_t
        self.factor_s = factor_s
        self.factor = factor_t * factor_s * factor_s
        assert in_channels * self.factor % out_channels == 0
        self.group_size = in_channels * self.factor // out_channels

    def __call__(self, x: mx.array) -> mx.array:
        # x: (B, T, H, W, C)
        B, T, H, W, C = x.shape
        pad_t = (self.factor_t - T % self.factor_t) % self.factor_t
        if pad_t:
            x = mx.pad(x, [(0, 0), (pad_t, 0), (0, 0), (0, 0), (0, 0)])
            T = T + pad_t
        ft, fs = self.factor_t, self.factor_s
        x = x.reshape(B, T // ft, ft, H // fs, fs, W // fs, fs, C)
        # PT permutes (B, C, ft, fs, fs, T/ft, H/fs, W/fs) so the merged channel
        # dim has order [C(slow), ft, fs, fs(fast)]. In NTHWC the slow factor is
        # batch, then spatial outer, then we want C slowest among the merged tail.
        # Permute (B, T/ft, H/fs, W/fs, C, ft, fs(H), fs(W)).
        x = mx.transpose(x, (0, 1, 3, 5, 7, 2, 4, 6))
        x = x.reshape(B, T // ft, H // fs, W // fs, C * self.factor)
        x = x.reshape(B, T // ft, H // fs, W // fs, self.out_channels, self.group_size)
        return mx.mean(x, axis=-1)


class DupUp3D(nn.Module):

    def __init__(self, in_channels: int, out_channels: int, factor_t: int, factor_s: int = 1):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.factor_t = factor_t
        self.factor_s = factor_s
        self.factor = factor_t * factor_s * factor_s
        assert out_channels * self.factor % in_channels == 0
        self.repeats = out_channels * self.factor // in_channels

    def __call__(self, x: mx.array, first_chunk: bool = False) -> mx.array:
        # x: (B, T, H, W, C). Repeat channel dim self.repeats times -> (B, T, H, W, C * repeats)
        # then split into (out_channels, factor_t, factor_s, factor_s) and pixel-shuffle.
        x = mx.repeat(x, self.repeats, axis=-1)
        B, T, H, W, _ = x.shape
        x = x.reshape(B, T, H, W, self.out_channels, self.factor_t, self.factor_s, self.factor_s)
        # Interleave factor dims into spatial/temporal: target
        # (B, T, ft, H, fs, W, fs, Cout) -> (B, T*ft, H*fs, W*fs, Cout)
        x = mx.transpose(x, (0, 1, 5, 2, 6, 3, 7, 4))
        x = x.reshape(B, T * self.factor_t, H * self.factor_s, W * self.factor_s, self.out_channels)
        if first_chunk and self.factor_t > 1:
            # PT drops the first (factor_t - 1) "anticipated" frames of the first chunk.
            x = x[:, self.factor_t - 1 :, :, :, :]
        return x


# ---------------------------------------------------------------------------
# Down/Up_ResidualBlock — repeated ResBlocks + optional sample, with avg shortcut
# ---------------------------------------------------------------------------

class Down_ResidualBlock(nn.Module):

    def __init__(self, in_dim: int, out_dim: int, dropout: float, mult: int,
                 temperal_downsample: bool = False, down_flag: bool = False):
        super().__init__()
        self.avg_shortcut = AvgDown3D(
            in_dim, out_dim,
            factor_t=2 if temperal_downsample else 1,
            factor_s=2 if down_flag else 1,
        )

        layers: List[nn.Module] = []
        cur_in = in_dim
        for _ in range(mult):
            layers.append(ResidualBlock(cur_in, out_dim, dropout))
            cur_in = out_dim
        if down_flag:
            mode = "downsample3d" if temperal_downsample else "downsample2d"
            layers.append(Resample(out_dim, mode=mode))
        self.downsamples = layers

    def __call__(self, x: mx.array, feat_cache=None, feat_idx=None) -> mx.array:
        x_copy = x
        for module in self.downsamples:
            if isinstance(module, (ResidualBlock, Resample)):
                x = module(x, feat_cache=feat_cache, feat_idx=feat_idx)
            else:
                x = module(x)
        return x + self.avg_shortcut(x_copy)


class Up_ResidualBlock(nn.Module):

    def __init__(self, in_dim: int, out_dim: int, dropout: float, mult: int,
                 temperal_upsample: bool = False, up_flag: bool = False):
        super().__init__()
        if up_flag:
            self.avg_shortcut = DupUp3D(
                in_dim, out_dim,
                factor_t=2 if temperal_upsample else 1,
                factor_s=2 if up_flag else 1,
            )
        else:
            self.avg_shortcut = None

        layers: List[nn.Module] = []
        cur_in = in_dim
        for _ in range(mult):
            layers.append(ResidualBlock(cur_in, out_dim, dropout))
            cur_in = out_dim
        if up_flag:
            mode = "upsample3d" if temperal_upsample else "upsample2d"
            layers.append(Resample(out_dim, mode=mode))
        self.upsamples = layers

    def __call__(self, x: mx.array, feat_cache=None, feat_idx=None,
                 first_chunk: bool = False) -> mx.array:
        x_main = x
        for module in self.upsamples:
            if isinstance(module, (ResidualBlock, Resample)):
                x_main = module(x_main, feat_cache=feat_cache, feat_idx=feat_idx)
            else:
                x_main = module(x_main)
        if self.avg_shortcut is not None:
            return x_main + self.avg_shortcut(x, first_chunk=first_chunk)
        return x_main


# ---------------------------------------------------------------------------
# Encoder3d / Decoder3d
# ---------------------------------------------------------------------------

class Encoder3d(nn.Module):

    def __init__(self, dim: int = 128, z_dim: int = 4, dim_mult=(1, 2, 4, 4),
                 num_res_blocks: int = 2, attn_scales=(),
                 temperal_downsample=(True, True, False), dropout: float = 0.0):
        super().__init__()
        self.dim = dim
        self.z_dim = z_dim
        self.dim_mult = list(dim_mult)
        self.num_res_blocks = num_res_blocks

        dims = [dim * u for u in [1] + list(dim_mult)]

        # input: 3-channel RGB patchified 2x -> 12 channels
        self.conv1 = CausalConv3d(12, dims[0], 3, padding=1)

        downsamples: List[nn.Module] = []
        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
            t_down_flag = temperal_downsample[i] if i < len(temperal_downsample) else False
            downsamples.append(
                Down_ResidualBlock(
                    in_dim=in_dim, out_dim=out_dim, dropout=dropout, mult=num_res_blocks,
                    temperal_downsample=t_down_flag, down_flag=i != len(dim_mult) - 1,
                )
            )
        self.downsamples = downsamples

        bottleneck_dim = dims[-1]
        self.middle = [
            ResidualBlock(bottleneck_dim, bottleneck_dim, dropout),
            AttentionBlock(bottleneck_dim),
            ResidualBlock(bottleneck_dim, bottleneck_dim, dropout),
        ]

        self.head_norm = RMS_norm(bottleneck_dim)
        self.head_conv = CausalConv3d(bottleneck_dim, z_dim, 3, padding=1)

    def __call__(self, x: mx.array, feat_cache=None, feat_idx=None) -> mx.array:
        prev, new_cache, idx = _cache_pre(x, feat_cache, feat_idx)
        x = self.conv1(x, cache_x=prev)
        _set_cache(feat_cache, idx, new_cache)
        for blk in self.downsamples:
            x = blk(x, feat_cache=feat_cache, feat_idx=feat_idx)
        for blk in self.middle:
            if isinstance(blk, ResidualBlock):
                x = blk(x, feat_cache=feat_cache, feat_idx=feat_idx)
            else:
                x = blk(x)
        x = self.head_norm(x)
        x = nn.silu(x)
        prev, new_cache, idx = _cache_pre(x, feat_cache, feat_idx)
        x = self.head_conv(x, cache_x=prev)
        _set_cache(feat_cache, idx, new_cache)
        return x


class Decoder3d(nn.Module):

    def __init__(self, dim: int = 128, z_dim: int = 4, dim_mult=(1, 2, 4, 4),
                 num_res_blocks: int = 2, attn_scales=(),
                 temperal_upsample=(False, True, True), dropout: float = 0.0):
        super().__init__()
        self.dim = dim
        self.z_dim = z_dim
        self.dim_mult = list(dim_mult)
        self.num_res_blocks = num_res_blocks

        # dim list runs in reverse for the decoder
        rev_mult = list(dim_mult[::-1])
        dims = [dim * u for u in [dim_mult[-1]] + rev_mult]

        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)

        bottleneck_dim = dims[0]
        self.middle = [
            ResidualBlock(bottleneck_dim, bottleneck_dim, dropout),
            AttentionBlock(bottleneck_dim),
            ResidualBlock(bottleneck_dim, bottleneck_dim, dropout),
        ]

        upsamples: List[nn.Module] = []
        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
            t_up_flag = temperal_upsample[i] if i < len(temperal_upsample) else False
            upsamples.append(
                Up_ResidualBlock(
                    in_dim=in_dim, out_dim=out_dim, dropout=dropout, mult=num_res_blocks + 1,
                    temperal_upsample=t_up_flag, up_flag=i != len(dim_mult) - 1,
                )
            )
        self.upsamples = upsamples

        out_final = dims[-1]
        self.head_norm = RMS_norm(out_final)
        self.head_conv = CausalConv3d(out_final, 12, 3, padding=1)

    def __call__(self, x: mx.array, feat_cache=None, feat_idx=None,
                 first_chunk: bool = False) -> mx.array:
        prev, new_cache, idx = _cache_pre(x, feat_cache, feat_idx)
        x = self.conv1(x, cache_x=prev)
        _set_cache(feat_cache, idx, new_cache)
        for blk in self.middle:
            if isinstance(blk, ResidualBlock):
                x = blk(x, feat_cache=feat_cache, feat_idx=feat_idx)
            else:
                x = blk(x)
        for blk in self.upsamples:
            x = blk(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
        x = self.head_norm(x)
        x = nn.silu(x)
        prev, new_cache, idx = _cache_pre(x, feat_cache, feat_idx)
        x = self.head_conv(x, cache_x=prev)
        _set_cache(feat_cache, idx, new_cache)
        return x


# ---------------------------------------------------------------------------
# WanVAE_ — inner model: encoder + conv1 (mu/log_var) + conv2 + decoder
# ---------------------------------------------------------------------------

class WanVAE_(nn.Module):

    def __init__(self, dim: int = 160, dec_dim: int = 256, z_dim: int = 16,
                 dim_mult=(1, 2, 4, 4), num_res_blocks: int = 2, attn_scales=(),
                 temperal_downsample=(True, True, False), dropout: float = 0.0):
        super().__init__()
        self.z_dim = z_dim
        self.temperal_downsample = list(temperal_downsample)
        self.temperal_upsample = list(temperal_downsample[::-1])

        self.encoder = Encoder3d(
            dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales,
            self.temperal_downsample, dropout,
        )
        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
        self.decoder = Decoder3d(
            dec_dim, z_dim, dim_mult, num_res_blocks, attn_scales,
            self.temperal_upsample, dropout,
        )

    # -----------------------------------------------------------------------
    # Conv counting (for cache slot allocation)
    # -----------------------------------------------------------------------
    @staticmethod
    def _count_conv3d(module: nn.Module) -> int:
        """Count CausalConv3d instances that participate in streaming-cache.

        A CausalConv3d takes cache iff its kernel has a temporal extent > 1
        (i.e. ._pt > 0). 1×1 shortcuts inside ResidualBlock have no temporal
        kernel and are called without cache in the forward path, so they
        must NOT consume a cache slot.

        Resample with mode upsample3d/downsample3d contributes one slot for
        its time_conv (also a CausalConv3d with temporal kernel).
        """
        n = 0
        visited = set()

        def _walk(obj):
            nonlocal n
            obj_id = id(obj)
            if obj_id in visited:
                return
            visited.add(obj_id)
            if isinstance(obj, Resample):
                # upsample3d / downsample3d each have a time_conv that
                # consumes exactly one cache slot. downsample3d's time_conv
                # has padding=0 (_pt=0) but still uses the cache to prepend
                # the last frame of the previous chunk before convolution.
                # Count unconditionally for these modes.
                if obj.mode in ("upsample3d", "downsample3d"):
                    n += 1
                return
            if isinstance(obj, CausalConv3d):
                if obj._pt > 0:
                    n += 1
                return
            if isinstance(obj, nn.Module):
                # MLX nn.Module is dict-like: children are stored via .items()
                # (and __setattr__ routes through __setitem__).
                for _, v in obj.items():
                    _walk(v)
            elif isinstance(obj, (list, tuple)):
                for v in obj:
                    _walk(v)

        _walk(module)
        return n

    def encode(self, x: mx.array, scale) -> tuple:
        """Encode (B, T, H, W, 3) in [-1, 1] -> (mu, log_var) each (B, T', H', W', z_dim).

        For T=1 (single image): single-pass.
        For T>1: streaming-cache chunked encode matching the PT reference. The
        encoder consumes frames in chunks of (1, 4, 4, 4, ...) and shares
        intermediate temporal-conv state across chunks via feat_cache.
        """
        x = patchify(x, patch_size=2)                  # (B, T, H/2, W/2, 12)
        T = x.shape[1]
        if T == 1:
            out = self.encoder(x)
        else:
            # chunked: first chunk is frame [0:1], then chunks of 4
            n_cache = self._count_conv3d(self.encoder)
            feat_cache = [None] * n_cache
            feat_idx = [0]
            iter_n = 1 + (T - 1) // 4
            outs = []
            for i in range(iter_n):
                feat_idx[0] = 0
                if i == 0:
                    xi = x[:, :1, ...]
                else:
                    start = 1 + 4 * (i - 1)
                    end = min(1 + 4 * i, T)
                    xi = x[:, start:end, ...]
                outs.append(self.encoder(xi, feat_cache=feat_cache, feat_idx=feat_idx))
            out = mx.concatenate(outs, axis=1)
        h = self.conv1(out)
        mu, log_var = mx.split(h, 2, axis=-1)
        mean, inv_std = scale
        mu = (mu - mean) * inv_std
        return mu, log_var

    def decode(self, z: mx.array, scale) -> mx.array:
        """Decode (B, T', H', W', z_dim) -> (B, T, H, W, 3).

        For T'=1 (single image): single-pass.
        For T'>1: streaming-cache chunked decode (one latent frame at a time).
        """
        mean, inv_std = scale
        z = z / inv_std + mean
        x = self.conv2(z)
        T_lat = x.shape[1]
        if T_lat == 1:
            out = self.decoder(x, first_chunk=True)
        else:
            n_cache = self._count_conv3d(self.decoder)
            feat_cache = [None] * n_cache
            feat_idx = [0]
            outs = []
            for i in range(T_lat):
                feat_idx[0] = 0
                xi = x[:, i:i+1, ...]
                first_chunk = (i == 0)
                outs.append(self.decoder(
                    xi, feat_cache=feat_cache, feat_idx=feat_idx,
                    first_chunk=first_chunk,
                ))
            out = mx.concatenate(outs, axis=1)
        return unpatchify(out, patch_size=2)


# ---------------------------------------------------------------------------
# Wan2_2_VAE — public wrapper with checkpoint-derived mean/std normalization
# ---------------------------------------------------------------------------

# Lance/Wan 2.2 normalization constants (z_dim=48). Copied verbatim from PT
# vae2_2.py so behavior matches.
_WAN22_MEAN = [
    -0.2289, -0.0052, -0.1323, -0.2339, -0.2799,  0.0174,  0.1838,  0.1557,
    -0.1382,  0.0542,  0.2813,  0.0891,  0.1570, -0.0098,  0.0375, -0.1825,
    -0.2246, -0.1207, -0.0698,  0.5109,  0.2665, -0.2108, -0.2158,  0.2502,
    -0.2055, -0.0322,  0.1109,  0.1567, -0.0729,  0.0899, -0.2799, -0.1230,
    -0.0313, -0.1649,  0.0117,  0.0723, -0.2839, -0.2083, -0.0520,  0.3748,
     0.0152,  0.1957,  0.1433, -0.2944,  0.3573, -0.0548, -0.1681, -0.0667,
]
_WAN22_STD = [
    0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013,
    0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978,
    0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659,
    0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093,
    0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887,
    0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744,
]


class Wan2_2_VAE:
    """Outer wrapper: applies channel-wise mean/std normalization around the inner VAE."""

    def __init__(
        self,
        z_dim: int = 48,
        c_dim: int = 160,
        dim_mult=(1, 2, 4, 4),
        temperal_downsample=(False, True, True),
        dtype=mx.float32,
    ):
        self.dtype = dtype
        mean = mx.array(_WAN22_MEAN, dtype=dtype)             # (48,)
        inv_std = 1.0 / mx.array(_WAN22_STD, dtype=dtype)     # (48,)
        # channel-last broadcasting: mean/inv_std are (C,) and broadcast over (B,T,H,W,C)
        self.scale = (mean, inv_std)

        self.model = WanVAE_(
            dim=c_dim,
            z_dim=z_dim,
            dim_mult=dim_mult,
            num_res_blocks=2,
            attn_scales=(),
            temperal_downsample=temperal_downsample,
            dropout=0.0,
        )

    def encode(self, video: mx.array):
        """video: (B, T, H, W, 3) float in [-1, 1]"""
        mu, log_var = self.model.encode(video, self.scale)
        return mu, log_var

    def decode(self, u: mx.array) -> mx.array:
        """u: (B, T', H', W', 48). Returns (B, T, H, W, 3) clamped to [-1, 1]."""
        x = self.model.decode(u, self.scale)
        return mx.clip(x, -1.0, 1.0)