anima_mlx/runtime/latent.py

"""Latent shape, noise, and Wan 2.1 latent-format helpers."""

from __future__ import annotations

from typing import Any

WAN21_LATENTS_MEAN = (
    -0.7571,
    -0.7089,
    -0.9113,
    0.1075,
    -0.1745,
    0.9653,
    -0.1517,
    1.5508,
    0.4134,
    -0.0715,
    0.5517,
    -0.3632,
    -0.1922,
    -0.9497,
    0.2503,
    -0.2921,
)

WAN21_LATENTS_STD = (
    2.8184,
    1.4541,
    2.3275,
    2.6558,
    1.2196,
    1.7708,
    2.6052,
    2.0743,
    3.2687,
    2.1526,
    2.8652,
    1.5579,
    1.6382,
    1.1253,
    2.8251,
    1.9160,
)


def latent_shape(batch_size: int, height: int, width: int, frames: int = 1) -> tuple[int, int, int, int, int]:
    if batch_size <= 0:
        raise ValueError("batch_size must be positive")
    if frames <= 0:
        raise ValueError("frames must be positive")
    if height <= 0 or width <= 0:
        raise ValueError("height and width must be positive")
    if height % 8 != 0 or width % 8 != 0:
        raise ValueError("height and width must be divisible by 8 for Anima VAE latents")
    return (batch_size, 16, frames, height // 8, width // 8)


def empty_anima_latent(
    *,
    batch_size: int,
    height: int,
    width: int,
    frames: int = 1,
    dtype: str = "float32",
) -> Any:
    import mlx.core as mx

    return mx.zeros(latent_shape(batch_size, height, width, frames), dtype=_mlx_dtype(dtype))


def random_anima_noise(shape: tuple[int, ...], seed: int, *, dtype: str = "float32") -> Any:
    import numpy as np
    import mlx.core as mx

    rng = np.random.default_rng(int(seed))
    noise = rng.standard_normal(shape).astype(np.float32)
    return mx.array(noise, dtype=_mlx_dtype(dtype))


def apply_const_noise_scaling(noise: Any, latent: Any, sigma: Any) -> Any:
    return _reshape_sigma(sigma, noise) * noise + (1.0 - _reshape_sigma(sigma, latent)) * latent


def wan21_process_in(latent: Any) -> Any:
    mean, std = _wan21_stats(latent)
    return (latent - mean) / std


def wan21_process_out(latent: Any) -> Any:
    mean, std = _wan21_stats(latent)
    return latent * std + mean


def _wan21_stats(reference: Any) -> tuple[Any, Any]:
    import mlx.core as mx

    mean = mx.array(WAN21_LATENTS_MEAN, dtype=reference.dtype).reshape(1, 16, 1, 1, 1)
    std = mx.array(WAN21_LATENTS_STD, dtype=reference.dtype).reshape(1, 16, 1, 1, 1)
    return mean, std


def _reshape_sigma(sigma: Any, reference: Any) -> Any:
    if hasattr(sigma, "shape") and len(sigma.shape) > 0:
        return sigma.reshape(sigma.shape[:1] + (1,) * (len(reference.shape) - 1))
    return sigma


def _mlx_dtype(dtype: str) -> Any:
    import mlx.core as mx

    dtype_map = {
        "float32": mx.float32,
        "float16": mx.float16,
        "bfloat16": mx.bfloat16,
    }
    if dtype not in dtype_map:
        raise ValueError(f"unsupported MLX dtype: {dtype}")
    return dtype_map[dtype]