anima_mlx/models/vae.py

"""Wan VAE decoder for Anima MLX."""

from __future__ import annotations

import math
from pathlib import Path
from typing import Any, Mapping


class CausalConv3d:
    def __init__(
        self,
        weights: Mapping[str, Any],
        prefix: str,
        *,
        kernel_size: int | tuple[int, int, int],
        padding: int | tuple[int, int, int] = 0,
        stride: int | tuple[int, int, int] = 1,
    ) -> None:
        self.weights = weights
        self.prefix = prefix
        self.kernel_size = _triple(kernel_size)
        self.padding = _triple(padding)
        self.stride = _triple(stride)

    def __call__(self, x: Any) -> Any:
        return causal_conv3d(
            x,
            self.weights[f"{self.prefix}.weight"],
            self.weights.get(f"{self.prefix}.bias"),
            padding=self.padding,
            stride=self.stride,
        )


class ResidualBlock:
    def __init__(self, weights: Mapping[str, Any], prefix: str) -> None:
        self.weights = weights
        self.prefix = prefix
        self.conv1 = CausalConv3d(weights, f"{prefix}.residual.2", kernel_size=3, padding=1)
        self.conv2 = CausalConv3d(weights, f"{prefix}.residual.6", kernel_size=3, padding=1)
        self.shortcut = (
            CausalConv3d(weights, f"{prefix}.shortcut", kernel_size=1)
            if f"{prefix}.shortcut.weight" in weights
            else None
        )

    def __call__(self, x: Any) -> Any:
        residual = vae_rms_norm(x, self.weights[f"{self.prefix}.residual.0.gamma"])
        residual = silu(residual)
        residual = self.conv1(residual)
        residual = vae_rms_norm(residual, self.weights[f"{self.prefix}.residual.3.gamma"])
        residual = silu(residual)
        residual = self.conv2(residual)
        shortcut = x if self.shortcut is None else self.shortcut(x)
        return residual + shortcut


class AttentionBlock:
    def __init__(self, weights: Mapping[str, Any], prefix: str) -> None:
        self.weights = weights
        self.prefix = prefix

    def __call__(self, x: Any) -> Any:
        import mlx.core as mx

        identity = x
        b, c, t, h, w = x.shape
        frames = x.transpose(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
        frames = vae_rms_norm(frames, self.weights[f"{self.prefix}.norm.gamma"])
        qkv = conv2d(
            frames,
            self.weights[f"{self.prefix}.to_qkv.weight"],
            self.weights.get(f"{self.prefix}.to_qkv.bias"),
        )
        q, k, v = mx.split(qkv, 3, axis=1)
        attended = vae_attention(q, k, v)
        projected = conv2d(
            attended,
            self.weights[f"{self.prefix}.proj.weight"],
            self.weights.get(f"{self.prefix}.proj.bias"),
        )
        projected = projected.reshape(b, t, c, h, w).transpose(0, 2, 1, 3, 4)
        return projected + identity


class Resample:
    def __init__(self, weights: Mapping[str, Any], prefix: str, *, mode: str) -> None:
        self.weights = weights
        self.prefix = prefix
        self.mode = mode

    def __call__(self, x: Any) -> Any:
        import mlx.core as mx

        if self.mode not in {"upsample2d", "upsample3d"}:
            raise NotImplementedError(f"unsupported decoder resample mode: {self.mode}")
        b, c, t, h, w = x.shape
        frames = x.transpose(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
        frames = mx.repeat(frames, 2, axis=2)
        frames = mx.repeat(frames, 2, axis=3)
        frames = conv2d(
            frames,
            self.weights[f"{self.prefix}.resample.1.weight"],
            self.weights.get(f"{self.prefix}.resample.1.bias"),
            padding=1,
        )
        out_channels = frames.shape[1]
        return frames.reshape(b, t, out_channels, h * 2, w * 2).transpose(0, 2, 1, 3, 4)


class Decoder3d:
    def __init__(self, weights: Mapping[str, Any]) -> None:
        self.weights = weights
        self.conv1 = CausalConv3d(weights, "decoder.conv1", kernel_size=3, padding=1)
        self.middle = [
            ResidualBlock(weights, "decoder.middle.0"),
            AttentionBlock(weights, "decoder.middle.1"),
            ResidualBlock(weights, "decoder.middle.2"),
        ]
        self.upsamples = self._build_upsamples()
        self.head = CausalConv3d(weights, "decoder.head.2", kernel_size=3, padding=1)

    def _build_upsamples(self) -> list[Any]:
        layers: list[Any] = []
        modes = {3: "upsample3d", 7: "upsample3d", 11: "upsample2d"}
        for index in range(15):
            prefix = f"decoder.upsamples.{index}"
            if f"{prefix}.resample.1.weight" in self.weights:
                layers.append(Resample(self.weights, prefix, mode=modes[index]))
            else:
                layers.append(ResidualBlock(self.weights, prefix))
        return layers

    def __call__(self, x: Any) -> Any:
        x = self.conv1(x)
        for layer in self.middle:
            x = layer(x)
        for layer in self.upsamples:
            x = layer(x)
        x = vae_rms_norm(x, self.weights["decoder.head.0.gamma"])
        x = silu(x)
        return self.head(x)


class WanVAEDecoder:
    """MLX implementation of the Anima/Wan VAE decoder path."""

    def __init__(self, weights: Mapping[str, Any]) -> None:
        self.weights = dict(weights)
        self.conv2 = CausalConv3d(self.weights, "conv2", kernel_size=1)
        self.decoder = Decoder3d(self.weights)

    @classmethod
    def from_safetensors(cls, path: str | Path, *, dtype: str = "float32") -> "WanVAEDecoder":
        from anima_mlx.utils.weights import load_mlx_safetensors_subset

        path = _resolve_vae_path(path)
        weights = load_mlx_safetensors_subset(
            path,
            key_filter=lambda key: key.startswith(("conv2.", "decoder.")),
            dtype=dtype,
        )
        return cls(weights)

    def decode(self, latent: Any) -> Any:
        return self.decoder(self.conv2(latent))

    def decode_tiled(self, latent: Any, *, tile_size: int = 64, overlap: int = 16) -> Any:
        """Decode latent spatial tiles.

        This is a memory fallback, not a quality-equivalent path: the Wan decoder
        has a middle spatial attention block, so each tile attends over only its
        local crop instead of the full latent plane.
        """

        import mlx.core as mx

        if tile_size <= 0:
            raise ValueError("tile_size must be positive")
        if overlap < 0:
            raise ValueError("overlap must be non-negative")
        if overlap >= tile_size:
            raise ValueError("overlap must be smaller than tile_size")

        batch, _, frames, latent_h, latent_w = latent.shape
        if latent_h <= tile_size and latent_w <= tile_size:
            return self.decode(latent)

        output_h = latent_h * 8
        output_w = latent_w * 8
        output = mx.zeros((batch, 3, frames, output_h, output_w), dtype=latent.dtype)
        output_div = mx.zeros((batch, 1, frames, output_h, output_w), dtype=latent.dtype)
        y_positions = _tile_positions(latent_h, tile_size, overlap)
        x_positions = _tile_positions(latent_w, tile_size, overlap)

        for y in y_positions:
            tile_h = min(tile_size, latent_h - y)
            for x in x_positions:
                tile_w = min(tile_size, latent_w - x)
                tile = latent[:, :, :, y : y + tile_h, x : x + tile_w]
                decoded = self.decode(tile)
                mx.eval(decoded)
                out_y = y * 8
                out_x = x * 8
                mask = _tile_mask(decoded.shape, overlap_h=min(overlap, tile_h), overlap_w=min(overlap, tile_w), dtype=decoded.dtype)
                output = _add_tile(output, decoded * mask, out_y=out_y, out_x=out_x)
                output_div = _add_tile(output_div, mask, out_y=out_y, out_x=out_x)
                mx.eval(output, output_div)
        return output / output_div

    def __call__(self, latent: Any) -> Any:
        return self.decode(latent)


def _tile_positions(length: int, tile_size: int, overlap: int) -> list[int]:
    if length <= tile_size:
        return [0]
    stride = tile_size - overlap
    positions: list[int] = []
    current = 0
    while current < length:
        pos = max(0, min(length - overlap, current))
        if positions and pos == positions[-1]:
            break
        positions.append(pos)
        if pos + tile_size >= length:
            break
        current += stride
    return positions


def _tile_mask(shape: tuple[int, ...], *, overlap_h: int, overlap_w: int, dtype: Any) -> Any:
    import mlx.core as mx

    _, _, frames, height, width = shape
    mask = mx.ones((1, 1, frames, height, width), dtype=dtype)
    feather_h = min(overlap_h * 8, height)
    feather_w = min(overlap_w * 8, width)
    if feather_h < height:
        for index in range(feather_h):
            value = mx.array((index + 1) / feather_h, dtype=dtype)
            mask = mask.at[:, :, :, index : index + 1, :].multiply(value)
            mask = mask.at[:, :, :, height - 1 - index : height - index, :].multiply(value)
    if feather_w < width:
        for index in range(feather_w):
            value = mx.array((index + 1) / feather_w, dtype=dtype)
            mask = mask.at[:, :, :, :, index : index + 1].multiply(value)
            mask = mask.at[:, :, :, :, width - 1 - index : width - index].multiply(value)
    return mask


def _add_tile(output: Any, tile: Any, *, out_y: int, out_x: int) -> Any:
    height = tile.shape[-2]
    width = tile.shape[-1]
    return output.at[:, :, :, out_y : out_y + height, out_x : out_x + width].add(tile)


def _triple(value: int | tuple[int, int, int]) -> tuple[int, int, int]:
    if isinstance(value, int):
        return (value, value, value)
    return value


def _pair(value: int | tuple[int, int]) -> tuple[int, int]:
    if isinstance(value, int):
        return (value, value)
    return value


def _conv2d_weight(weight: Any) -> Any:
    return weight.transpose(0, 2, 3, 1)


def _conv3d_weight(weight: Any) -> Any:
    return weight.transpose(0, 2, 3, 4, 1)


def conv2d(
    x: Any,
    weight: Any,
    bias: Any | None = None,
    *,
    padding: int | tuple[int, int] = 0,
    stride: int | tuple[int, int] = 1,
) -> Any:
    import mlx.core as mx

    x_cl = x.transpose(0, 2, 3, 1)
    y = mx.conv2d(x_cl, _conv2d_weight(weight), stride=_pair(stride), padding=_pair(padding))
    if bias is not None:
        y = y + bias.reshape(1, 1, 1, -1)
    return y.transpose(0, 3, 1, 2)


def causal_conv3d(
    x: Any,
    weight: Any,
    bias: Any | None = None,
    *,
    padding: int | tuple[int, int, int] = 0,
    stride: int | tuple[int, int, int] = 1,
) -> Any:
    import mlx.core as mx

    pad_t, pad_h, pad_w = _triple(padding)
    stride_t, stride_h, stride_w = _triple(stride)
    if x.shape[2] == 1:
        frame = x[:, :, 0, :, :]
        kernel = weight[:, :, -1, :, :]
        y = conv2d(frame, kernel, bias, padding=(pad_h, pad_w), stride=(stride_h, stride_w))
        return y[:, :, None, :, :]

    x_cl = x.transpose(0, 2, 3, 4, 1)
    if pad_t:
        x_cl = mx.pad(x_cl, [(0, 0), (2 * pad_t, 0), (0, 0), (0, 0), (0, 0)])
    y = mx.conv3d(
        x_cl,
        _conv3d_weight(weight),
        stride=(stride_t, stride_h, stride_w),
        padding=(0, pad_h, pad_w),
    )
    if bias is not None:
        y = y + bias.reshape(1, 1, 1, 1, -1)
    return y.transpose(0, 4, 1, 2, 3)


def vae_rms_norm(x: Any, gamma: Any) -> Any:
    import mlx.core as mx

    x32 = x.astype(mx.float32)
    denom = mx.sqrt(mx.sum(mx.square(x32), axis=1, keepdims=True))
    denom = mx.maximum(denom, mx.array(1e-12, dtype=mx.float32))
    scale = math.sqrt(x.shape[1])
    return x32 / denom * scale * gamma.astype(mx.float32)


def silu(x: Any) -> Any:
    import mlx.core as mx

    return x * mx.sigmoid(x)


def vae_attention(q: Any, k: Any, v: Any) -> Any:
    import mlx.core as mx

    b, c, h, w = q.shape
    tokens = h * w
    q_tokens = q.reshape(b, c, tokens).transpose(0, 2, 1)
    k_tokens = k.reshape(b, c, tokens)
    v_tokens = v.reshape(b, c, tokens).transpose(0, 2, 1)
    scores = (q_tokens @ k_tokens) * (c ** -0.5)
    probs = mx.softmax(scores, axis=-1)
    out = probs @ v_tokens
    return out.transpose(0, 2, 1).reshape(b, c, h, w)


def _resolve_vae_path(path: str | Path) -> Path:
    resolved = Path(path)
    if resolved.is_dir():
        return resolved / "vae.safetensors"
    return resolved