anima_mlx/models/dit.py

"""DiT blocks for Anima MLX."""

from __future__ import annotations

from pathlib import Path
from typing import Any, Mapping

from .attention import scaled_dot_product_attention
from .primitives import gelu, layer_norm, linear, rms_norm, silu


def apply_cosmos_rotary_pos_emb(x: Any, freqs: Any) -> Any:
    import mlx.core as mx

    x_pairs = x.reshape(*x.shape[:-1], 2, -1)
    x_pairs = mx.swapaxes(x_pairs, -2, -1)
    x_pairs = mx.expand_dims(x_pairs.astype(mx.float32), axis=-2)
    out = freqs[..., 0] * x_pairs[..., 0] + freqs[..., 1] * x_pairs[..., 1]
    return mx.swapaxes(out, -1, -2).reshape(*x.shape).astype(x.dtype)


def video_rope3d(shape: tuple[int, int, int, int, int], head_dim: int = 128) -> Any:
    import mlx.core as mx

    _, t_len, h_len, w_len, _ = shape
    dim_h = head_dim // 6 * 2
    dim_w = dim_h
    dim_t = head_dim - 2 * dim_h
    dim_spatial_range = mx.arange(0, dim_h, 2).astype(mx.float32)[: dim_h // 2] / dim_h
    dim_temporal_range = mx.arange(0, dim_t, 2).astype(mx.float32)[: dim_t // 2] / dim_t
    h_ntk_factor = 4.0 ** (dim_h / (dim_h - 2))
    w_ntk_factor = 4.0 ** (dim_w / (dim_w - 2))
    t_ntk_factor = 1.0 ** (dim_t / (dim_t - 2))
    h_freqs = 1.0 / ((10000.0 * h_ntk_factor) ** dim_spatial_range)
    w_freqs = 1.0 / ((10000.0 * w_ntk_factor) ** dim_spatial_range)
    t_freqs = 1.0 / ((10000.0 * t_ntk_factor) ** dim_temporal_range)
    seq = mx.arange(max(h_len, w_len, t_len)).astype(mx.float32)
    half_h = mx.outer(seq[:h_len], h_freqs)
    half_w = mx.outer(seq[:w_len], w_freqs)
    half_t = mx.outer(seq[:t_len], t_freqs)

    def stack_freqs(x: Any) -> Any:
        return mx.stack([mx.cos(x), -mx.sin(x), mx.sin(x), mx.cos(x)], axis=-1)

    half_h = stack_freqs(half_h)
    half_w = stack_freqs(half_w)
    half_t = stack_freqs(half_t)
    emb_t = mx.broadcast_to(half_t[:, None, None, :, :], (t_len, h_len, w_len, dim_t // 2, 4))
    emb_h = mx.broadcast_to(half_h[None, :, None, :, :], (t_len, h_len, w_len, dim_h // 2, 4))
    emb_w = mx.broadcast_to(half_w[None, None, :, :, :], (t_len, h_len, w_len, dim_w // 2, 4))
    emb = mx.concatenate([emb_t, emb_h, emb_w], axis=-2)
    return emb.reshape(t_len * h_len * w_len, head_dim // 2, 2, 2).astype(mx.float32)


class DiTAttention:
    def __init__(self, weights: Mapping[str, Any], prefix: str, *, n_heads: int = 16, head_dim: int = 128) -> None:
        self.weights = weights
        self.prefix = prefix
        self.n_heads = n_heads
        self.head_dim = head_dim

    def _weight(self, name: str) -> Any:
        return self.weights[f"{self.prefix}.{name}"]

    def _linear(self, name: str, x: Any) -> Any:
        return linear(x, self._weight(f"{name}.weight"))

    def __call__(
        self,
        x: Any,
        *,
        context: Any | None = None,
        rope_emb: Any | None = None,
        trace: dict[str, Any] | None = None,
        trace_prefix: str = "",
    ) -> Any:
        import mlx.core as mx

        is_selfattn = context is None
        context = x if context is None else context
        input_shape = x.shape[:-1]
        context_shape = context.shape[:-1]

        q = self._linear("q_proj", x).reshape(*input_shape, self.n_heads, self.head_dim)
        k = self._linear("k_proj", context).reshape(*context_shape, self.n_heads, self.head_dim)
        v = self._linear("v_proj", context).reshape(*context_shape, self.n_heads, self.head_dim)
        if trace is not None:
            trace[f"{trace_prefix}q_proj"] = q
            trace[f"{trace_prefix}k_proj"] = k
            trace[f"{trace_prefix}v_proj"] = v
        q = rms_norm(q, self._weight("q_norm.weight"))
        k = rms_norm(k, self._weight("k_norm.weight"))
        if trace is not None:
            trace[f"{trace_prefix}q_norm"] = q
            trace[f"{trace_prefix}k_norm"] = k

        if is_selfattn and rope_emb is not None:
            q = apply_cosmos_rotary_pos_emb(q, rope_emb)
            k = apply_cosmos_rotary_pos_emb(k, rope_emb)
            if trace is not None:
                trace[f"{trace_prefix}q_rope"] = q
                trace[f"{trace_prefix}k_rope"] = k

        q = mx.swapaxes(q, 1, 2)
        k = mx.swapaxes(k, 1, 2)
        v = mx.swapaxes(v, 1, 2)
        out = scaled_dot_product_attention(q, k, v)
        out = mx.swapaxes(out, 1, 2).reshape(*input_shape, self.n_heads * self.head_dim)
        if trace is not None:
            trace[f"{trace_prefix}attention_kernel_output"] = out
        output = self._linear("output_proj", out)
        if trace is not None:
            trace[f"{trace_prefix}output_proj"] = output
            trace[f"{trace_prefix}output"] = output
        return output


class DiTBlock:
    def __init__(self, weights: Mapping[str, Any], prefix: str = "blocks.0") -> None:
        self.weights = self._normalize_keys(weights)
        self.prefix = prefix
        self.self_attn = DiTAttention(self.weights, f"{prefix}.self_attn")
        self.cross_attn = DiTAttention(self.weights, f"{prefix}.cross_attn")

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

        path = _resolve_diffusion_block_path(path, block_index)
        prefix = f"net.blocks.{block_index}."
        weights = load_mlx_safetensors_subset(path, prefix=prefix, strip_prefix="net.", dtype=dtype)
        return cls(weights, prefix=f"blocks.{block_index}")

    @staticmethod
    def _normalize_keys(weights: Mapping[str, Any]) -> dict[str, Any]:
        normalized: dict[str, Any] = {}
        for key, value in weights.items():
            if key.startswith("net."):
                key = key.removeprefix("net.")
            normalized[key] = value
        return normalized

    def _weight(self, name: str) -> Any:
        return self.weights[f"{self.prefix}.{name}"]

    def _adaln(self, name: str, emb: Any, adaln_lora: Any) -> tuple[Any, Any, Any]:
        hidden = silu(emb)
        hidden = linear(hidden, self._weight(f"adaln_modulation_{name}.1.weight"))
        hidden = linear(hidden, self._weight(f"adaln_modulation_{name}.2.weight"))
        hidden = hidden + adaln_lora
        size = hidden.shape[-1] // 3
        return hidden[..., :size], hidden[..., size : 2 * size], hidden[..., 2 * size :]

    @staticmethod
    def _modulation_shape(x: Any, value: Any) -> Any:
        return value.reshape(value.shape[0], value.shape[1], *([1] * (len(x.shape) - 3)), value.shape[-1])

    def __call__(
        self,
        x: Any,
        emb: Any,
        context: Any,
        rope_emb: Any,
        adaln_lora: Any,
        *,
        trace: dict[str, Any] | None = None,
        trace_prefix: str = "",
    ) -> Any:
        residual_dtype = x.dtype
        b, t_len, h_len, w_len, dim = x.shape

        shift, scale, gate = self._adaln("self_attn", emb, adaln_lora)
        if trace is not None:
            trace[f"{trace_prefix}self_attn_shift"] = shift
            trace[f"{trace_prefix}self_attn_scale"] = scale
            trace[f"{trace_prefix}self_attn_gate"] = gate
        shift = self._modulation_shape(x, shift)
        scale = self._modulation_shape(x, scale)
        gate = self._modulation_shape(x, gate)
        normed = layer_norm(x) * (1 + scale) + shift
        if trace is not None:
            trace[f"{trace_prefix}self_attn_normed"] = normed
        self_attn_input = normed.reshape(b, t_len * h_len * w_len, dim)
        if trace is not None:
            trace[f"{trace_prefix}self_attn_input"] = self_attn_input
        result = self.self_attn(
            self_attn_input,
            rope_emb=rope_emb,
            trace=trace,
            trace_prefix=f"{trace_prefix}self_attn_",
        )
        if trace is not None:
            trace[f"{trace_prefix}self_attn_output_flat"] = result
        result = result.reshape(b, t_len, h_len, w_len, dim)
        if trace is not None:
            trace[f"{trace_prefix}self_attn_output_reshaped"] = result
        x = x + gate.astype(residual_dtype) * result.astype(residual_dtype)
        if trace is not None:
            trace[f"{trace_prefix}self_attn_residual_output"] = x

        shift, scale, gate = self._adaln("cross_attn", emb, adaln_lora)
        if trace is not None:
            trace[f"{trace_prefix}cross_attn_shift"] = shift
            trace[f"{trace_prefix}cross_attn_scale"] = scale
            trace[f"{trace_prefix}cross_attn_gate"] = gate
        shift = self._modulation_shape(x, shift)
        scale = self._modulation_shape(x, scale)
        gate = self._modulation_shape(x, gate)
        normed = layer_norm(x) * (1 + scale) + shift
        if trace is not None:
            trace[f"{trace_prefix}cross_attn_normed"] = normed
        cross_attn_input = normed.reshape(b, t_len * h_len * w_len, dim)
        if trace is not None:
            trace[f"{trace_prefix}cross_attn_input"] = cross_attn_input
        result = self.cross_attn(
            cross_attn_input,
            context=context,
            trace=trace,
            trace_prefix=f"{trace_prefix}cross_attn_",
        )
        if trace is not None:
            trace[f"{trace_prefix}cross_attn_output_flat"] = result
        result = result.reshape(b, t_len, h_len, w_len, dim)
        if trace is not None:
            trace[f"{trace_prefix}cross_attn_output_reshaped"] = result
        x = x + gate.astype(residual_dtype) * result.astype(residual_dtype)
        if trace is not None:
            trace[f"{trace_prefix}cross_attn_residual_output"] = x

        shift, scale, gate = self._adaln("mlp", emb, adaln_lora)
        if trace is not None:
            trace[f"{trace_prefix}mlp_shift"] = shift
            trace[f"{trace_prefix}mlp_scale"] = scale
            trace[f"{trace_prefix}mlp_gate"] = gate
        shift = self._modulation_shape(x, shift)
        scale = self._modulation_shape(x, scale)
        gate = self._modulation_shape(x, gate)
        normed = layer_norm(x) * (1 + scale) + shift
        if trace is not None:
            trace[f"{trace_prefix}mlp_normed"] = normed
        hidden = linear(normed, self._weight("mlp.layer1.weight"))
        if trace is not None:
            trace[f"{trace_prefix}mlp_layer1_output"] = hidden
        hidden = gelu(hidden)
        if trace is not None:
            trace[f"{trace_prefix}mlp_gelu_output"] = hidden
        result = linear(hidden, self._weight("mlp.layer2.weight"))
        if trace is not None:
            trace[f"{trace_prefix}mlp_layer2_output"] = result
        residual_product = gate.astype(residual_dtype) * result.astype(residual_dtype)
        output = x + residual_product
        if trace is not None:
            trace[f"{trace_prefix}mlp_residual_product"] = residual_product
            trace[f"{trace_prefix}mlp_residual_output"] = output
        return output


def dit_timesteps(timesteps_b_t: Any, num_channels: int = 2048) -> Any:
    import math
    import mlx.core as mx

    timesteps = timesteps_b_t.reshape(-1).astype(mx.float32)
    half_dim = num_channels // 2
    exponent = -math.log(10000) * mx.arange(half_dim).astype(mx.float32)
    exponent = exponent / (half_dim - 0.0)
    emb = mx.exp(exponent)
    emb = timesteps[:, None] * emb[None, :]
    emb = mx.concatenate([mx.cos(emb), mx.sin(emb)], axis=-1)
    return emb.reshape(timesteps_b_t.shape[0], timesteps_b_t.shape[1], num_channels)


def patchify_latent(latent: Any) -> Any:
    import mlx.core as mx

    b, _, t_len, h_len, w_len = latent.shape
    pad_h = h_len % 2
    pad_w = w_len % 2
    if pad_h or pad_w:
        latent = mx.pad(latent, [(0, 0), (0, 0), (0, 0), (0, pad_h), (0, pad_w)])
        h_len = latent.shape[-2]
        w_len = latent.shape[-1]
    padding_mask = mx.zeros((b, 1, h_len, w_len), dtype=latent.dtype)
    if pad_h or pad_w:
        valid_h = h_len - pad_h
        valid_w = w_len - pad_w
        valid = mx.zeros((b, 1, valid_h, valid_w), dtype=latent.dtype)
        if pad_w:
            valid = mx.concatenate([valid, mx.ones((b, 1, valid_h, pad_w), dtype=latent.dtype)], axis=-1)
        if pad_h:
            valid = mx.concatenate([valid, mx.ones((b, 1, pad_h, w_len), dtype=latent.dtype)], axis=-2)
        padding_mask = valid
    patch_input = mx.concatenate([latent, mx.repeat(padding_mask[:, :, None, :, :], t_len, axis=2)], axis=1)
    b, channels, t_len, h_len, w_len = patch_input.shape
    return (
        patch_input.reshape(b, channels, t_len, h_len // 2, 2, w_len // 2, 2)
        .transpose(0, 2, 3, 5, 1, 4, 6)
        .reshape(b, t_len, h_len // 2, w_len // 2, 68)
    )


def unpatchify(x: Any) -> Any:
    b, t_len, h_len, w_len, _ = x.shape
    return (
        x.reshape(b, t_len, h_len, w_len, 2, 2, 1, 16)
        .transpose(0, 7, 1, 6, 2, 4, 3, 5)
        .reshape(b, 16, t_len, h_len * 2, w_len * 2)
    )


class DiT:
    def __init__(
        self,
        weights: Mapping[str, Any],
        *,
        block_count: int = 28,
        weights_path: str | Path | None = None,
        dtype: str = "float32",
        eval_each_block: bool = True,
        eval_interval: int | None = None,
    ) -> None:
        self.weights = DiTBlock._normalize_keys(weights)
        self.block_count = block_count
        self.weights_path = Path(weights_path) if weights_path is not None else None
        self.dtype = dtype
        self.eval_interval = 1 if eval_interval is None and eval_each_block else eval_interval or 0
        self.blocks = None if self.weights_path is not None else [
            DiTBlock(self.weights, prefix=f"blocks.{index}") for index in range(block_count)
        ]

    @classmethod
    def from_safetensors(
        cls,
        path: str | Path,
        *,
        block_count: int = 28,
        dtype: str = "float32",
        lazy_blocks: bool = True,
        eval_interval: int = 1,
    ) -> "DiT":
        from anima_mlx.utils.weights import load_mlx_safetensors_subset

        if eval_interval < 0:
            raise ValueError("eval_interval must be non-negative")
        path = _resolve_diffusion_path(path)
        core_prefixes = (
            "net.x_embedder.",
            "net.t_embedder.",
            "net.t_embedding_norm.",
            "net.final_layer.",
        )
        weights = load_mlx_safetensors_subset(
            path,
            strip_prefix="net.",
            key_filter=lambda key: key.startswith(("net.blocks.", *core_prefixes) if not lazy_blocks else core_prefixes),
            dtype=dtype,
        )
        if not lazy_blocks and not any(key.startswith("blocks.") for key in weights):
            for block_index in range(block_count):
                weights.update(_load_diffusion_block_weights(path, block_index, dtype=dtype))
        return cls(
            weights,
            block_count=block_count,
            weights_path=path if lazy_blocks else None,
            dtype=dtype,
            eval_interval=eval_interval,
        )

    def _weight(self, name: str) -> Any:
        return self.weights[name]

    def prepare_inputs_trace(self, latent: Any, timestep: Any) -> tuple[Any, Any, Any, dict[str, Any]]:
        patch_input = patchify_latent(latent)
        x = linear(patch_input, self._weight("x_embedder.proj.1.weight"))
        t_emb = dit_timesteps(timestep, 2048)
        t_emb_for_embedder = t_emb.astype(x.dtype)
        hidden = linear(t_emb_for_embedder, self._weight("t_embedder.1.linear_1.weight"))
        hidden_activated = silu(hidden)
        adaln_lora = linear(hidden_activated, self._weight("t_embedder.1.linear_2.weight"))
        t_embedding = rms_norm(t_emb_for_embedder, self._weight("t_embedding_norm.weight"))
        return x, t_embedding, adaln_lora, {
            "patch_input": patch_input,
            "x_embedder_output": x,
            "timestep_embedding_raw": t_emb,
            "t_embedder_linear_1_output": hidden,
            "t_embedder_silu_output": hidden_activated,
            "t_embedding": t_embedding,
            "adaln_lora": adaln_lora,
        }

    def prepare_inputs(self, latent: Any, timestep: Any) -> tuple[Any, Any, Any]:
        x, t_embedding, adaln_lora, _ = self.prepare_inputs_trace(latent, timestep)
        return x, t_embedding, adaln_lora

    def final_layer(self, x: Any, emb: Any, adaln_lora: Any, *, trace: dict[str, Any] | None = None) -> Any:
        hidden = silu(emb)
        if trace is not None:
            trace["final_adaln_silu_output"] = hidden
        hidden = linear(hidden, self._weight("final_layer.adaln_modulation.1.weight"))
        if trace is not None:
            trace["final_adaln_linear_1_output"] = hidden
        hidden = linear(hidden, self._weight("final_layer.adaln_modulation.2.weight"))
        if trace is not None:
            trace["final_adaln_linear_2_output"] = hidden
        hidden = hidden + adaln_lora[:, :, :4096]
        shift, scale = hidden[..., :2048], hidden[..., 2048:]
        if trace is not None:
            trace["final_shift"] = shift
            trace["final_scale"] = scale
        shift = DiTBlock._modulation_shape(x, shift)
        scale = DiTBlock._modulation_shape(x, scale)
        normed = layer_norm(x) * (1 + scale) + shift
        if trace is not None:
            trace["final_normed"] = normed
        output = linear(normed, self._weight("final_layer.linear.weight"))
        if trace is not None:
            trace["final_linear_output"] = output
        return output

    def _block(self, index: int) -> DiTBlock:
        if self.blocks is not None:
            return self.blocks[index]
        if self.weights_path is None:
            raise RuntimeError("DiT block weights are not available")
        return DiTBlock.from_safetensors(self.weights_path, block_index=index, dtype=self.dtype)

    def __call__(
        self,
        latent: Any,
        timestep: Any,
        context: Any,
        *,
        checkpoint_indices: set[int] | None = None,
        trace: bool = False,
    ) -> Any | tuple[Any, dict[str, Any], Any] | tuple[Any, dict[str, Any]]:
        checkpoint_indices = checkpoint_indices or set()
        x, t_embedding, adaln_lora, trace_data = self.prepare_inputs_trace(latent, timestep)
        rope = video_rope3d(tuple(x.shape), head_dim=128)
        import mlx.core as mx

        if trace:
            trace_data["rope"] = rope
        rope = mx.expand_dims(mx.expand_dims(rope, axis=1), axis=0)
        if trace:
            trace_data["rope_expanded"] = rope
        checkpoints: dict[str, Any] = {}
        for index in range(self.block_count):
            block = self._block(index)
            if trace and index in {0, 1, 13, 27}:
                trace_data[f"block_{index}_input"] = x
            block_trace = trace_data if trace and index in {0, 1} else None
            block_x = x.astype(mx.float32) if x.dtype == mx.float16 else x
            block_t_embedding = t_embedding.astype(mx.float32) if t_embedding.dtype == mx.float16 else t_embedding
            block_adaln_lora = adaln_lora.astype(mx.float32) if adaln_lora.dtype == mx.float16 else adaln_lora
            x = block(
                block_x,
                block_t_embedding,
                context,
                rope,
                block_adaln_lora,
                trace=block_trace,
                trace_prefix=f"block_{index}_",
            )
            if self.eval_interval and ((index + 1) % self.eval_interval == 0 or index == self.block_count - 1):
                mx.eval(x)
            if trace and index in {0, 1, 13, 27}:
                trace_data[f"block_{index}_output"] = x
            if index in checkpoint_indices:
                checkpoints[f"block_{index}_output"] = x
        if trace:
            trace_data["final_layer_input"] = x
        final_x = x.astype(mx.float32) if x.dtype == mx.float16 else x
        final_t_embedding = t_embedding.astype(mx.float32) if t_embedding.dtype == mx.float16 else t_embedding
        final_adaln_lora = adaln_lora.astype(mx.float32) if adaln_lora.dtype == mx.float16 else adaln_lora
        patch_output = self.final_layer(final_x, final_t_embedding, final_adaln_lora, trace=trace_data if trace else None)
        unpatchified = unpatchify(patch_output)
        denoised = unpatchified[:, :, : latent.shape[-3], : latent.shape[-2], : latent.shape[-1]]
        if trace:
            trace_data["patch_output"] = patch_output
            trace_data["unpatchified"] = unpatchified
            trace_data["denoised"] = denoised
            return denoised, trace_data
        if checkpoint_indices:
            return denoised, checkpoints, patch_output
        return denoised


def _resolve_diffusion_path(path: str | Path) -> Path:
    resolved = Path(path)
    if resolved.is_dir():
        core_path = resolved / "diffusion_core.safetensors"
        if core_path.exists():
            return core_path
        return resolved / "diffusion.safetensors"
    return resolved


def _resolve_diffusion_block_path(path: str | Path, block_index: int) -> Path:
    resolved = Path(path)
    if resolved.is_dir():
        block_path = resolved / "diffusion_blocks" / f"block_{block_index:02d}.safetensors"
        if block_path.exists():
            return block_path
        return resolved / "diffusion.safetensors"
    if resolved.name == "diffusion_core.safetensors":
        block_path = resolved.parent / "diffusion_blocks" / f"block_{block_index:02d}.safetensors"
        if block_path.exists():
            return block_path
        fallback = resolved.parent / "diffusion.safetensors"
        if fallback.exists():
            return fallback
    return resolved


def _load_diffusion_block_weights(path: str | Path, block_index: int, *, dtype: str) -> dict[str, Any]:
    from anima_mlx.utils.weights import load_mlx_safetensors_subset

    block_path = _resolve_diffusion_block_path(path, block_index)
    prefix = f"net.blocks.{block_index}."
    return load_mlx_safetensors_subset(block_path, prefix=prefix, strip_prefix="net.", dtype=dtype)