"""Experimental native-MLX inference for the pinned MiniMax-Music3 repack.

This module mirrors ComfyUI commit efd4e951's Music3 AR, flow-DiT, and DAV
implementations while reading PocketAI's existing verified safetensors files.
MLX is imported lazily so catalog and protocol tests remain platform-neutral.
"""

from __future__ import annotations

import gc
import hashlib
import math
import re
from functools import lru_cache
from pathlib import Path
from typing import Any, Callable


SAMPLE_RATE = 44_100
AUDIO_FRAMES_PER_SECOND = 25
MAX_AUDIO_FRAMES = 9_000
DAV_UPSAMPLE_FACTOR = 512
# MLX's Metal Conv1d path produces incorrect values for the multi-million-sample
# intermediate tensors created by long DAV decodes. Keep each fully-convolutional
# decode below that size, with substantially more context than DAV's receptive
# field, then retain only the context-independent center of each chunk.
DAV_MAX_CHUNK_FRAMES = 1_024
DAV_OVERLAP_FRAMES = 64
C0_VOCAB_SIZE = 16_384
AUDIO_VOCAB_SIZE = 1_024
NUM_CODEBOOKS = 8
AUDIO_CODE_OFFSET = 151_675
SPECIAL_TOKEN_IDS = {
    "<|im_start|>": 151_644,
    "<|im_end|>": 151_645,
    "<|audio_cfg|>": 151_654,
    "<|audio_start|>": 151_669,
    "<|audio_end|>": 151_670,
    "<|caption_start|>": 151_671,
    "<|caption_end|>": 151_672,
    "<|lyrics_start|>": 151_673,
    "<|lyrics_end|>": 151_674,
}
_SPECIAL_TAG_RE = re.compile(r"<\|([^|]*)\|>")
_LYRIC_TAG_RE = re.compile(r"\s*(\[[^\]]+\])\s*")
Progress = Callable[[int, str], None]


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

    return mx


def _numpy() -> Any:
    import numpy as np

    return np


def _remove_markdown_format(text: str) -> str:
    lines: list[str] = []
    for raw_line in text.splitlines():
        line = re.sub(r"^\s{0,3}#{1,6}\s+", "", raw_line)
        line = re.sub(r"^\s*[*+-]\s+", "", line)
        while "**" in line:
            updated = re.sub(r"\*\*([^*]+)\*\*", r"\1", line)
            if updated == line:
                break
            line = updated
        line = re.sub(r"(?<!\*)\*([^*\n]+)\*(?!\*)", r"\1", line)
        lines.append(line.rstrip())
    text = "\n".join(lines)
    text = re.sub(r"^\s*[-*_]{3,}\s*$", "", text, flags=re.MULTILINE)
    return text.replace("• ", "").replace("    ", "")


def clean_caption(caption: str) -> str:
    def replace_special(match: re.Match[str]) -> str:
        parts = match.group(1).strip().split(None, 1)
        return f"{parts[0]} is {parts[1]}" if len(parts) == 2 else parts[0]

    text = _SPECIAL_TAG_RE.sub(replace_special, caption)
    return re.sub(r"\n{2,}", "\n", _remove_markdown_format(text))


def normalize_lyrics(lyrics: str) -> str:
    parts = _LYRIC_TAG_RE.split(lyrics)
    text = "\n".join(part.lower() if part.startswith("[") else part for part in parts if part)
    text = text.replace(" ^ ", "\n")
    return f"[start]\n{text}"


def build_prompt(caption: str, lyrics: str) -> str:
    return (
        "<|im_start|><|caption_start|>"
        f"{clean_caption(caption)}"
        "<|caption_end|><|lyrics_start|>"
        f"{normalize_lyrics(lyrics)}"
        "<|lyrics_end|><|im_end|><|audio_start|>"
    )


def derive_seed(seed: int, *parts: object) -> int:
    digest = hashlib.blake2b(digest_size=8, person=b"minimax-ttm")
    digest.update(int(seed).to_bytes(8, "little", signed=False))
    for part in parts:
        value = str(part).encode("utf-8")
        digest.update(len(value).to_bytes(4, "little"))
        digest.update(value)
    return int.from_bytes(digest.digest(), "little") & ((1 << 63) - 1)


def latent_length(audio_frames: int) -> int:
    return max(1, int(audio_frames * 44_100 / 24_000 * 960 / 512))


def simple_flow_schedule(steps: int) -> list[float]:
    if not 1 <= steps <= 30:
        raise ValueError("steps must be between 1 and 30")
    schedule = [(1_000 - int(index * 1_000 / steps)) / 1_000 for index in range(steps)]
    return [*schedule, 0.0]


def dav_decode_slices(
    frames: int,
    max_chunk_frames: int = DAV_MAX_CHUNK_FRAMES,
    overlap_frames: int = DAV_OVERLAP_FRAMES,
) -> list[tuple[int, int, int, int]]:
    """Plan (context start/end, retained start/end) slices in latent frames."""
    if frames < 1:
        raise ValueError("DAV decode requires at least one latent frame")
    if overlap_frames < 0 or max_chunk_frames <= overlap_frames * 2:
        raise ValueError("DAV chunk size must be greater than twice its overlap")
    if frames <= max_chunk_frames:
        return [(0, frames, 0, frames)]

    retained_frames = max_chunk_frames - overlap_frames * 2
    slices: list[tuple[int, int, int, int]] = []
    retained_start = 0
    while retained_start < frames:
        retained_end = min(frames, retained_start + retained_frames)
        context_start = max(0, retained_start - overlap_frames)
        context_end = min(frames, retained_end + overlap_frames)
        local_start = retained_start - context_start
        local_end = local_start + retained_end - retained_start
        slices.append((context_start, context_end, local_start, local_end))
        retained_start = retained_end
    return slices


def stereo_collapse_fraction(samples: Any, sample_rate: int = SAMPLE_RATE) -> float:
    """Return the fraction of whole-second blocks with a near-missing channel."""
    np = _numpy()
    if samples.ndim != 2 or samples.shape[0] != 2:
        raise ValueError("MiniMax-Music3 audio must have shape (2, samples)")
    blocks = samples.shape[1] // sample_rate
    if blocks < 1:
        return 0.0
    framed = samples[:, : blocks * sample_rate].reshape(2, blocks, sample_rate)
    block_rms = np.sqrt(np.mean(framed.astype(np.float64) ** 2, axis=-1))
    lower = np.min(block_rms, axis=0)
    upper = np.max(block_rms, axis=0)
    return float(np.mean((upper > 1e-7) & (lower < upper * 0.1)))


@lru_cache(maxsize=1)
def _normalized_hadamard() -> Any:
    np = _numpy()
    h4 = np.array(
        [[1, 1, 1, -1], [1, 1, -1, 1], [1, -1, 1, 1], [-1, 1, 1, 1]],
        dtype=np.float32,
    )
    matrix = h4
    while matrix.shape[0] < 256:
        matrix = np.kron(matrix, h4)
    matrix /= math.sqrt(256)
    matrix.setflags(write=False)
    return matrix


def _pack_bytes_as_uint32(values: Any, rows: int) -> Any:
    mx = _mlx()
    if values.size % (rows * 4):
        raise ValueError("packed ConvRot weights must contain whole uint32 words")
    return values.reshape(-1).view(mx.uint32).reshape(rows, -1)


class ConvRotInt8Linear:
    """Comfy tensorwise INT8 plus 256-wide ConvRot, evaluated by MLX QMM."""

    def __init__(self, weight: Any, scale: Any, bias: Any | None = None) -> None:
        mx = _mlx()
        if weight.ndim != 2 or weight.shape[1] % 256:
            raise ValueError(f"invalid ConvRot weight shape: {weight.shape}")
        self.out_features, self.in_features = weight.shape
        row_scales = scale.reshape(self.out_features, -1)
        if row_scales.shape[1] != 1:
            raise ValueError("ConvRot requires one scale per output row")
        scales = mx.repeat(row_scales, self.in_features // 128, axis=1)
        unsigned = (weight.astype(mx.int16) + 128).astype(mx.uint8)
        self.weight = _pack_bytes_as_uint32(unsigned, self.out_features)
        self.scales = scales
        self.quant_biases = -128 * scales
        self.rotation = mx.array(_normalized_hadamard()).astype(mx.bfloat16)
        self.bias = bias
        # Materialize one repack at a time. Leaving every conversion lazy would
        # retain the signed source and INT16 intermediate for the full model.
        mx.eval(self.weight, self.scales, self.quant_biases, self.rotation)

    def __call__(self, inputs: Any) -> Any:
        mx = _mlx()
        shape = inputs.shape
        rotated = (
            inputs.reshape(-1, self.in_features // 256, 256)
            @ self.rotation.astype(inputs.dtype)
        ).reshape(shape)
        output = mx.quantized_matmul(
            rotated,
            self.weight,
            self.scales,
            self.quant_biases,
            group_size=128,
            bits=8,
            mode="affine",
        ).astype(inputs.dtype)
        return output if self.bias is None else output + self.bias.astype(output.dtype)


class DenseLinear:
    def __init__(self, weight: Any, bias: Any | None = None) -> None:
        self.weight = weight
        self.bias = bias

    def __call__(self, inputs: Any) -> Any:
        output = inputs @ self.weight.astype(inputs.dtype).T
        return output if self.bias is None else output + self.bias.astype(output.dtype)


def _linear(weights: dict[str, Any], prefix: str) -> Any:
    weight = weights.pop(f"{prefix}.weight")
    bias = weights.pop(f"{prefix}.bias", None)
    scale = weights.pop(f"{prefix}.weight_scale", None)
    weights.pop(f"{prefix}.comfy_quant", None)
    if scale is not None:
        return ConvRotInt8Linear(weight, scale, bias)
    return DenseLinear(weight, bias)


def _rms_norm(inputs: Any, weight: Any, eps: float = 1e-6) -> Any:
    mx = _mlx()
    variance = mx.mean(mx.square(inputs.astype(mx.float32)), axis=-1, keepdims=True)
    return inputs * mx.rsqrt(variance + eps).astype(inputs.dtype) * weight.astype(inputs.dtype)


def _layer_norm(inputs: Any, gamma: Any, beta: Any, eps: float = 1e-5) -> Any:
    mx = _mlx()
    values = inputs.astype(mx.float32)
    mean = mx.mean(values, axis=-1, keepdims=True)
    variance = mx.mean(mx.square(values - mean), axis=-1, keepdims=True)
    normalized = (values - mean) * mx.rsqrt(variance + eps)
    return (normalized * gamma + beta).astype(inputs.dtype)


def _silu(inputs: Any) -> Any:
    mx = _mlx()
    return inputs * mx.sigmoid(inputs)


def _split_half_rope(values: Any, positions: Any, theta: float, dims: int | None = None) -> Any:
    mx = _mlx()
    rotary_dims = dims or values.shape[-1]
    rotary = values[..., :rotary_dims]
    tail = values[..., rotary_dims:]
    inv_freq = 1.0 / mx.power(
        mx.array(theta, dtype=mx.float32),
        mx.arange(0, rotary_dims, 2, dtype=mx.float32) / rotary_dims,
    )
    angles = positions.astype(mx.float32)[:, None] * inv_freq[None, :]
    cos = mx.concatenate((mx.cos(angles), mx.cos(angles)), axis=-1)[None, None]
    sin = mx.concatenate((mx.sin(angles), mx.sin(angles)), axis=-1)[None, None]
    half = rotary_dims // 2
    rotated = mx.concatenate((-rotary[..., half:], rotary[..., :half]), axis=-1)
    output = (rotary * cos + rotated * sin).astype(values.dtype)
    return mx.concatenate((output, tail), axis=-1) if tail.shape[-1] else output


class KVCache:
    def __init__(self, batch: int, heads: int, capacity: int, head_dim: int, dtype: Any) -> None:
        mx = _mlx()
        self.keys = mx.zeros((batch, heads, capacity, head_dim), dtype=dtype)
        self.values = mx.zeros((batch, heads, capacity, head_dim), dtype=dtype)
        self.offset = 0

    def update(self, keys: Any, values: Any) -> tuple[Any, Any]:
        mx = _mlx()
        length = keys.shape[2]
        end = self.offset + length
        if end > self.keys.shape[2]:
            raise ValueError("MiniMax-Music3 KV cache capacity exceeded")
        start = mx.array([self.offset], dtype=mx.int32)
        self.keys = mx.slice_update(self.keys, keys, start, axes=(2,))
        self.values = mx.slice_update(self.values, values, start, axes=(2,))
        self.offset = end
        return self.keys[:, :, :end], self.values[:, :, :end]


class QwenLayer:
    def __init__(self, weights: dict[str, Any], index: int) -> None:
        prefix = f"model.layers.{index}"
        self.input_norm = weights.pop(f"{prefix}.input_layernorm.weight")
        self.post_norm = weights.pop(f"{prefix}.post_attention_layernorm.weight")
        self.q_norm = weights.pop(f"{prefix}.self_attn.q_norm.weight")
        self.k_norm = weights.pop(f"{prefix}.self_attn.k_norm.weight")
        self.qkv = _linear(weights, f"{prefix}.self_attn.qkv_proj")
        self.out = _linear(weights, f"{prefix}.self_attn.o_proj")
        self.gate_up = _linear(weights, f"{prefix}.mlp.gate_up_proj")
        self.down = _linear(weights, f"{prefix}.mlp.down_proj")

    def __call__(self, inputs: Any, cache: KVCache) -> Any:
        mx = _mlx()
        batch, length, _ = inputs.shape
        residual = inputs
        qkv = self.qkv(_rms_norm(inputs, self.input_norm))
        q, k, v = mx.split(qkv, (4_096, 5_120), axis=-1)
        q = q.reshape(batch, length, 32, 128).transpose(0, 2, 1, 3)
        k = k.reshape(batch, length, 8, 128).transpose(0, 2, 1, 3)
        v = v.reshape(batch, length, 8, 128).transpose(0, 2, 1, 3)
        q = _rms_norm(q, self.q_norm)
        k = _rms_norm(k, self.k_norm)
        positions = mx.arange(cache.offset, cache.offset + length)
        q = _split_half_rope(q, positions, 1_000_000.0)
        k = _split_half_rope(k, positions, 1_000_000.0)
        keys, values = cache.update(k, v)
        attention = mx.fast.scaled_dot_product_attention(
            q,
            keys,
            values,
            scale=128**-0.5,
            mask="causal" if length > 1 else None,
        )
        attention = attention.transpose(0, 2, 1, 3).reshape(batch, length, 4_096)
        hidden = residual + self.out(attention)
        residual = hidden
        gate, up = mx.split(self.gate_up(_rms_norm(hidden, self.post_norm)), 2, axis=-1)
        return residual + self.down(_silu(gate) * up)


class RVQLayer:
    def __init__(self, weights: dict[str, Any], index: int) -> None:
        prefix = f"model.audio_decoder.layers.{index}"
        self.input_norm = weights.pop(f"{prefix}.input_layernorm.weight")
        self.post_norm = weights.pop(f"{prefix}.post_attention_layernorm.weight")
        self.qkv = _linear(weights, f"{prefix}.self_attn.qkv_proj")
        self.out = _linear(weights, f"{prefix}.self_attn.o_proj")
        self.gate_up = _linear(weights, f"{prefix}.mlp.gate_up_proj")
        self.down = _linear(weights, f"{prefix}.mlp.down_proj")

    def __call__(self, inputs: Any) -> Any:
        mx = _mlx()
        batch, length, _ = inputs.shape
        residual = inputs
        q, k, v = mx.split(self.qkv(_rms_norm(inputs, self.input_norm)), 3, axis=-1)
        q = q.reshape(batch, length, 16, 256).transpose(0, 2, 1, 3)
        k = k.reshape(batch, length, 16, 256).transpose(0, 2, 1, 3)
        v = v.reshape(batch, length, 16, 256).transpose(0, 2, 1, 3)
        attention = mx.fast.scaled_dot_product_attention(
            q, k, v, scale=256**-0.5, mask="causal"
        )
        attention = attention.transpose(0, 2, 1, 3).reshape(batch, length, 4_096)
        hidden = residual + self.out(attention)
        residual = hidden
        gate, up = mx.split(self.gate_up(_rms_norm(hidden, self.post_norm)), 2, axis=-1)
        return residual + self.down(_silu(gate) * up)


def _sample_top_k(logits: Any, top_k: int, rng: Any) -> int:
    mx = _mlx()
    np = _numpy()
    mx.eval(logits)
    values = np.nan_to_num(np.asarray(logits, dtype=np.float32), nan=-1e9, posinf=1e9, neginf=-1e9)
    values = values.reshape(-1)
    count = min(top_k, values.size)
    indices = np.argpartition(values, -count)[-count:]
    selected = values[indices]
    probabilities = np.exp(selected - selected.max())
    probabilities /= probabilities.sum()
    return int(rng.choice(indices, p=probabilities))


class Music3AR:
    def __init__(self, checkpoint: Path) -> None:
        mx = _mlx()
        np = _numpy()
        weights = dict(mx.load(str(checkpoint)))
        tokenizer_bytes = np.asarray(weights.pop("tokenizer_json"), dtype=np.uint8).tobytes()
        from tokenizers import Tokenizer

        self.tokenizer = Tokenizer.from_str(tokenizer_bytes.decode("utf-8"))
        for token, expected in SPECIAL_TOKEN_IDS.items():
            if self.tokenizer.token_to_id(token) != expected:
                raise ValueError(f"MiniMax-Music3 tokenizer mismatch for {token}")
        self.prefill_embedding = weights.pop("model.embed_tokens_prefill.weight")
        self.audio_embedding = weights.pop("model.embed_tokens_audio.weight")
        self.extra_embedding = weights.pop("model.audio_extra_embedding.weight")
        self.norm = weights.pop("model.norm.weight")
        self.lm_head = DenseLinear(weights.pop("model.lm_head_pruned.weight"))
        self.layers = [QwenLayer(weights, index) for index in range(36)]
        self.depth_projection = DenseLinear(weights.pop("model.audio_decoder.projection.weight"))
        self.depth_positions = weights.pop("model.audio_decoder.pos_embedding.weight")
        self.depth_layers = [RVQLayer(weights, index) for index in range(4)]
        self.depth_norm = weights.pop("model.audio_decoder.norm.weight")
        self.audio_heads = [
            DenseLinear(weights.pop(f"model.audio_decoder.audio_heads.{index}.weight"))
            for index in range(7)
        ]
        mx.eval(
            self.prefill_embedding,
            self.audio_embedding,
            self.extra_embedding,
            self.norm,
            self.depth_positions,
            self.depth_norm,
        )

    def _forward(self, embeds: Any, caches: list[KVCache]) -> Any:
        hidden = embeds
        for layer, cache in zip(self.layers, caches, strict=True):
            hidden = layer(hidden, cache)
        return _rms_norm(hidden, self.norm)

    def _depth_forward(self, sequence: Any) -> Any:
        hidden = sequence + self.depth_positions[: sequence.shape[1]].astype(sequence.dtype)[None]
        for layer in self.depth_layers:
            hidden = layer(hidden)
        return _rms_norm(hidden, self.depth_norm)

    def generate(
        self,
        caption: str,
        lyrics: str,
        seed: int,
        max_audio_frames: int,
        *,
        cfg_scale: float = 1.5,
        top_k: int = 50,
    ) -> Any:
        mx = _mlx()
        np = _numpy()
        prompt = build_prompt(caption, lyrics)
        token_ids = self.tokenizer.encode(prompt, add_special_tokens=False).ids
        if len(token_ids) > 5_000:
            raise ValueError(f"MiniMax-Music3 prompt has {len(token_ids)} tokens; maximum is 5000")
        conditioned = np.asarray(token_ids, dtype=np.int32)
        unconditioned = conditioned.copy()
        unconditioned[1:-2] = SPECIAL_TOKEN_IDS["<|audio_cfg|>"]
        ids = mx.array(np.stack((conditioned, unconditioned)))
        embeds = self.prefill_embedding[ids].astype(mx.bfloat16)
        limit = min(int(max_audio_frames), MAX_AUDIO_FRAMES)
        capacity = len(token_ids) + limit + 1
        caches = [KVCache(2, 8, capacity, 128, mx.bfloat16) for _ in range(36)]
        hidden = self._forward(embeds, caches)[:, -1]
        mx.eval(hidden, *[value for cache in caches for value in (cache.keys, cache.values)])
        rng = np.random.default_rng(derive_seed(seed, "ar"))
        hidden_frames: list[Any] = []
        pending_code: int | None = None
        pending_hidden: Any | None = None

        for frame_index in range(limit + 1):
            if pending_code == 0:
                pending_hidden = None
                break
            if pending_hidden is not None:
                hidden_frames.append(pending_hidden)
                if len(hidden_frames) >= limit:
                    break

            logits = self.lm_head(hidden).astype(mx.float32)
            conditioned_logits = logits[0:1]
            guided = logits[1:2] + (conditioned_logits - logits[1:2]) * cfg_scale
            threshold = mx.sort(conditioned_logits, axis=-1)[..., -min(top_k, C0_VOCAB_SIZE + 1)]
            guided = mx.where(conditioned_logits < threshold[..., None], -mx.inf, guided)
            token = _sample_top_k(guided, top_k, rng)
            pending_code = token
            c0_value = 0 if token == 0 else token - 1
            c0 = mx.array([c0_value, c0_value])
            c0_embed = self.audio_embedding[c0].astype(mx.bfloat16)
            sequence = [
                self.depth_projection(hidden)[:, None],
                self.depth_projection(c0_embed)[:, None],
            ]
            codes = [c0]
            depth_hidden: list[Any] = []
            for index in range(1, NUM_CODEBOOKS):
                depth = self._depth_forward(mx.concatenate(sequence, axis=1))[:, -1]
                depth_hidden.append(depth[:1])
                code_logits = self.audio_heads[index - 1](depth).astype(mx.float32)
                guided_depth = code_logits[1:2] + (code_logits[0:1] - code_logits[1:2]) * cfg_scale
                code_value = _sample_top_k(guided_depth, top_k, rng)
                code = mx.array([code_value, code_value])
                codes.append(code)
                if index < NUM_CODEBOOKS - 1:
                    embedding = self.extra_embedding[code + (index - 1) * AUDIO_VOCAB_SIZE]
                    sequence.append(self.depth_projection(embedding.astype(mx.bfloat16))[:, None])
            frame_hidden = mx.concatenate(
                (hidden[:1], mx.concatenate(depth_hidden, axis=-1)), axis=-1
            )[0]
            if frame_index > 0:
                pending_hidden = frame_hidden
            feedback_codes = mx.stack(codes, axis=1)
            c0_feedback = self.audio_embedding[feedback_codes[:, 0]]
            offsets = mx.arange(NUM_CODEBOOKS - 1) * AUDIO_VOCAB_SIZE
            extra = mx.sum(self.extra_embedding[feedback_codes[:, 1:] + offsets[None]], axis=1)
            feedback = ((c0_feedback + extra) * (NUM_CODEBOOKS**-0.5)).astype(mx.bfloat16)[:, None]
            hidden = self._forward(feedback, caches)[:, -1]
            mx.eval(
                hidden,
                frame_hidden,
                *[value for cache in caches for value in (cache.keys, cache.values)],
            )

        if pending_hidden is not None and pending_code != 0 and len(hidden_frames) < limit:
            hidden_frames.append(pending_hidden)
        if not hidden_frames:
            raise RuntimeError("MiniMax-Music3 generated zero audio frames")
        context = mx.stack(hidden_frames)[None]
        mx.eval(context)
        return context


def _conv1d(
    inputs: Any,
    weight: Any,
    bias: Any | None = None,
    *,
    stride: int = 1,
    padding: int = 0,
    dilation: int = 1,
) -> Any:
    mx = _mlx()
    output = mx.conv1d(
        inputs.transpose(0, 2, 1),
        weight.transpose(0, 2, 1).astype(inputs.dtype),
        stride=stride,
        padding=padding,
        dilation=dilation,
    ).transpose(0, 2, 1)
    return output if bias is None else output + bias.astype(output.dtype).reshape(1, -1, 1)


def _conv_transpose1d(
    inputs: Any, weight: Any, bias: Any | None, *, stride: int, padding: int
) -> Any:
    mx = _mlx()
    output = mx.conv_transpose1d(
        inputs.transpose(0, 2, 1),
        weight.transpose(1, 2, 0).astype(inputs.dtype),
        stride=stride,
        padding=padding,
    ).transpose(0, 2, 1)
    return output if bias is None else output + bias.astype(output.dtype).reshape(1, -1, 1)


class DiTLayer:
    def __init__(self, weights: dict[str, Any], index: int) -> None:
        prefix = f"diffusion_transformer.transformer.layers.{index}"
        self.pre_gamma = weights.pop(f"{prefix}.pre_norm.gamma")
        self.pre_beta = weights.pop(f"{prefix}.pre_norm.beta")
        self.ff_gamma = weights.pop(f"{prefix}.ff_norm.gamma")
        self.ff_beta = weights.pop(f"{prefix}.ff_norm.beta")
        self.qkv = _linear(weights, f"{prefix}.self_attn.to_qkv")
        self.out = _linear(weights, f"{prefix}.self_attn.to_out")
        self.ff_in = _linear(weights, f"{prefix}.ff.ff.0.proj")
        self.ff_out = _linear(weights, f"{prefix}.ff.ff.2")

    def __call__(self, inputs: Any, inv_freq: Any) -> Any:
        mx = _mlx()
        batch, length, _ = inputs.shape
        residual = inputs
        q, k, v = mx.split(
            self.qkv(_layer_norm(inputs, self.pre_gamma, self.pre_beta)), 3, axis=-1
        )
        q = q.reshape(batch, length, 32, 64).transpose(0, 2, 1, 3)
        k = k.reshape(batch, length, 32, 64).transpose(0, 2, 1, 3)
        v = v.reshape(batch, length, 32, 64).transpose(0, 2, 1, 3)
        positions = mx.arange(length, dtype=mx.float32)
        angles = positions[:, None] * inv_freq.astype(mx.float32)[None]
        cos = mx.concatenate((mx.cos(angles), mx.cos(angles)), axis=-1)[None, None]
        sin = mx.concatenate((mx.sin(angles), mx.sin(angles)), axis=-1)[None, None]
        for name, values in (("q", q), ("k", k)):
            rotary = values[..., :32]
            half = rotary.shape[-1] // 2
            rotated = mx.concatenate((-rotary[..., half:], rotary[..., :half]), axis=-1)
            updated = mx.concatenate(
                ((rotary * cos + rotated * sin).astype(values.dtype), values[..., 32:]),
                axis=-1,
            )
            if name == "q":
                q = updated
            else:
                k = updated
        attention = mx.fast.scaled_dot_product_attention(q, k, v, scale=64**-0.5)
        attention = attention.transpose(0, 2, 1, 3).reshape(batch, length, 2_048)
        hidden = residual + self.out(attention)
        residual = hidden
        value, gate = mx.split(
            self.ff_in(_layer_norm(hidden, self.ff_gamma, self.ff_beta)), 2, axis=-1
        )
        return residual + self.ff_out(value * _silu(gate))


class Music3DiT:
    def __init__(self, checkpoint: Path) -> None:
        mx = _mlx()
        weights = dict(mx.load(str(checkpoint)))
        self.cond_logits = weights.pop("cond_layer_logits")
        self.cond_scale = weights.pop("cond_layer_scale")
        self.condition_weight = weights.pop("latent_conditioners.0.weight")
        self.condition_bias = weights.pop("latent_conditioners.0.bias")
        self.preprocess_weight = weights.pop("diffusion_transformer.preprocess_conv.weight")
        self.postprocess_weight = weights.pop("diffusion_transformer.postprocess_conv.weight")
        self.fourier_weight = weights.pop("diffusion_transformer.timestep_features.weight")
        self.time_in = _linear(weights, "diffusion_transformer.to_timestep_embed.0")
        self.time_out = _linear(weights, "diffusion_transformer.to_timestep_embed.2")
        self.project_in = _linear(weights, "diffusion_transformer.transformer.project_in")
        self.project_out = _linear(weights, "diffusion_transformer.transformer.project_out")
        self.inv_freq = weights.pop("diffusion_transformer.transformer.rotary_pos_emb.inv_freq")
        self.layers = [DiTLayer(weights, index) for index in range(36)]
        mx.eval(
            self.cond_logits,
            self.cond_scale,
            self.condition_weight,
            self.condition_bias,
            self.preprocess_weight,
            self.postprocess_weight,
            self.fourier_weight,
            self.inv_freq,
        )

    def aligned_condition(self, context: Any) -> Any:
        mx = _mlx()
        batch, frames, _ = context.shape
        hidden = context.transpose(0, 2, 1).reshape(batch, 8, 4_096, frames)
        weights = mx.softmax(self.cond_logits.astype(hidden.dtype), axis=0)
        hidden = mx.sum(hidden * weights[None, :, None, None], axis=1)
        hidden = hidden * self.cond_scale.astype(hidden.dtype)
        condition = _conv1d(hidden, self.condition_weight, self.condition_bias, padding=1)
        length = latent_length(frames)
        indices = mx.floor(mx.arange(length) * frames / length).astype(mx.int32)
        return condition[..., indices]

    def _transform(self, latent: Any, timestep: Any, condition: Any) -> Any:
        mx = _mlx()
        full = mx.concatenate((latent, mx.zeros_like(latent), condition), axis=1)
        full = _conv1d(full, self.preprocess_weight) + full
        features = 2 * math.pi * timestep[:, None, None] @ self.fourier_weight.astype(mx.float32).T
        features = mx.concatenate((mx.cos(features[:, 0]), mx.sin(features[:, 0])), axis=-1)
        timestep_embedding = self.time_out(_silu(self.time_in(features.astype(latent.dtype))))
        hidden = self.project_in(full.transpose(0, 2, 1))
        hidden = mx.concatenate((timestep_embedding[:, None], hidden), axis=1)
        for layer in self.layers:
            hidden = layer(hidden, self.inv_freq)
        output = self.project_out(hidden[:, 1:]).transpose(0, 2, 1)
        return _conv1d(output, self.postprocess_weight) + output

    def velocity(self, latent: Any, sigma: float, condition: Any) -> Any:
        mx = _mlx()
        timestep = mx.full((latent.shape[0],), 1.0 - sigma, dtype=latent.dtype)
        window = latent_length(200)
        if latent.shape[-1] <= window:
            return -self._transform(latent, timestep, condition)
        output = mx.zeros_like(latent)
        count = mx.zeros((1, 1, latent.shape[-1]), dtype=latent.dtype)
        hop = latent_length(100)
        start = 0
        while start < latent.shape[-1]:
            end = min(start + window, latent.shape[-1])
            update = -self._transform(latent[..., start:end], timestep, condition[..., start:end])
            start_index = mx.array([start], dtype=mx.int32)
            output = mx.slice_update(
                output, output[..., start:end] + update, start_index, axes=(2,)
            )
            count = mx.slice_update(
                count, count[..., start:end] + 1, start_index, axes=(2,)
            )
            if end == latent.shape[-1]:
                break
            start += hop
        return output / count

    def sample(self, context: Any, steps: int, seed: int, cfg: float = 1.7) -> Any:
        mx = _mlx()
        condition = self.aligned_condition(context)
        uncondition = self.aligned_condition(mx.zeros_like(context))
        length = condition.shape[-1]
        latent = mx.random.normal(
            (1, 128, length), dtype=mx.float16, key=mx.random.key(seed)
        )
        mx.eval(latent, condition, uncondition)
        schedule = simple_flow_schedule(steps)
        for sigma, next_sigma in zip(schedule[:-1], schedule[1:], strict=True):
            batch_latent = mx.concatenate((latent, latent), axis=0)
            batch_condition = mx.concatenate((condition, uncondition), axis=0)
            velocity = self.velocity(batch_latent, sigma, batch_condition)
            guided = velocity[1:2] + (velocity[0:1] - velocity[1:2]) * cfg
            latent = latent + guided * (next_sigma - sigma)
            mx.eval(latent)
        return latent


def _weight_norm(weight_v: Any, weight_g: Any) -> Any:
    mx = _mlx()
    norm = mx.sqrt(mx.sum(mx.square(weight_v.astype(mx.float32)), axis=(1, 2), keepdims=True))
    weight = weight_v * (weight_g / mx.maximum(norm, 1e-12)).astype(weight_v.dtype)
    mx.eval(weight)
    return weight


class ResidualUnit:
    def __init__(self, weights: dict[str, Any], prefix: str, dilation: int) -> None:
        self.alpha1 = weights.pop(f"{prefix}.block.0.alpha")
        self.weight1 = _weight_norm(
            weights.pop(f"{prefix}.block.1.weight_v"), weights.pop(f"{prefix}.block.1.weight_g")
        )
        self.bias1 = weights.pop(f"{prefix}.block.1.bias")
        self.alpha2 = weights.pop(f"{prefix}.block.2.alpha")
        self.weight2 = _weight_norm(
            weights.pop(f"{prefix}.block.3.weight_v"), weights.pop(f"{prefix}.block.3.weight_g")
        )
        self.bias2 = weights.pop(f"{prefix}.block.3.bias")
        self.dilation = dilation

    @staticmethod
    def snake(inputs: Any, alpha: Any) -> Any:
        mx = _mlx()
        alpha = alpha.astype(inputs.dtype)
        return inputs + mx.square(mx.sin(alpha * inputs)) / (alpha + 1e-9)

    def __call__(self, inputs: Any) -> Any:
        hidden = self.snake(inputs, self.alpha1)
        hidden = _conv1d(
            hidden,
            self.weight1,
            self.bias1,
            padding=3 * self.dilation,
            dilation=self.dilation,
        )
        hidden = self.snake(hidden, self.alpha2)
        hidden = _conv1d(hidden, self.weight2, self.bias2)
        if hidden.shape[-1] != inputs.shape[-1]:
            padding = (inputs.shape[-1] - hidden.shape[-1]) // 2
            inputs = inputs[..., padding : inputs.shape[-1] - padding]
        return inputs + hidden


class DecoderBlock:
    def __init__(self, weights: dict[str, Any], index: int, stride: int) -> None:
        prefix = f"decoder.model.{index}.block"
        self.alpha = weights.pop(f"{prefix}.0.alpha")
        weight_v = weights.pop(f"{prefix}.1.weight_v")
        weight_g = weights.pop(f"{prefix}.1.weight_g")
        # ConvTranspose1d weight norm is per input channel (axis zero).
        mx = _mlx()
        norm = mx.sqrt(mx.sum(mx.square(weight_v.astype(mx.float32)), axis=(1, 2), keepdims=True))
        self.weight = weight_v * (weight_g / mx.maximum(norm, 1e-12)).astype(weight_v.dtype)
        mx.eval(self.weight)
        self.bias = weights.pop(f"{prefix}.1.bias")
        self.residuals = [
            ResidualUnit(weights, f"{prefix}.{offset}", dilation)
            for offset, dilation in zip((2, 3, 4), (1, 3, 9), strict=True)
        ]
        self.stride = stride

    def __call__(self, inputs: Any) -> Any:
        hidden = ResidualUnit.snake(inputs, self.alpha)
        hidden = _conv_transpose1d(
            hidden,
            self.weight,
            self.bias,
            stride=self.stride,
            padding=math.ceil(self.stride / 2),
        )
        for residual in self.residuals:
            hidden = residual(hidden)
        return hidden


class Music3DAV:
    def __init__(self, checkpoint: Path) -> None:
        mx = _mlx()
        weights = dict(mx.load(str(checkpoint)))
        self.input_weight = weights.pop("dec_in_proj.weight")
        self.input_bias = weights.pop("dec_in_proj.bias")
        self.first_weight = _weight_norm(
            weights.pop("decoder.model.0.weight_v"), weights.pop("decoder.model.0.weight_g")
        )
        self.first_bias = weights.pop("decoder.model.0.bias")
        self.blocks = [
            DecoderBlock(weights, index, stride)
            for index, stride in zip(range(1, 5), (8, 8, 4, 2), strict=True)
        ]
        self.final_alpha = weights.pop("decoder.model.5.alpha")
        self.final_weight = _weight_norm(
            weights.pop("decoder.model.6.weight_v"), weights.pop("decoder.model.6.weight_g")
        )
        self.final_bias = weights.pop("decoder.model.6.bias")
        mx.eval(self.input_weight, self.first_weight, self.final_weight)

    def _decode_chunk(self, latent: Any) -> Any:
        mx = _mlx()
        batch, _, frames = latent.shape
        hidden = latent.reshape(batch * 2, 64, frames).astype(mx.float32)
        hidden = _conv1d(hidden, self.input_weight, self.input_bias)
        hidden = _conv1d(hidden, self.first_weight, self.first_bias, padding=3)
        for block in self.blocks:
            hidden = block(hidden)
        hidden = ResidualUnit.snake(hidden, self.final_alpha)
        hidden = mx.tanh(_conv1d(hidden, self.final_weight, self.final_bias, padding=3))
        return hidden.reshape(batch, 2, -1)

    def decode(self, latent: Any) -> Any:
        mx = _mlx()
        pieces = []
        for context_start, context_end, local_start, local_end in dav_decode_slices(
            latent.shape[-1]
        ):
            waveform = self._decode_chunk(latent[..., context_start:context_end])
            pieces.append(
                waveform[
                    ...,
                    local_start * DAV_UPSAMPLE_FACTOR : local_end * DAV_UPSAMPLE_FACTOR,
                ]
            )
        return pieces[0] if len(pieces) == 1 else mx.concatenate(pieces, axis=-1)


class MiniMaxMusic3MlxPipeline:
    def __init__(self, model_dir: Path, progress: Progress | None = None) -> None:
        self.model_dir = model_dir.resolve(strict=True)
        notify = progress or (lambda _stage, _message: None)
        notify(1, "Loading MiniMax-Music3 autoregressive model into MLX…")
        self.ar = Music3AR(
            self.model_dir
            / "text_encoders/minimax_music3_text_encoder_pruned_int8_convrot.safetensors"
        )
        notify(2, "Loading MiniMax-Music3 flow transformer into MLX…")
        self.dit = Music3DiT(
            self.model_dir / "diffusion_models/minimax_music3_dit_int8_convrot.safetensors"
        )
        notify(3, "Loading MiniMax-Music3 DAV decoder into MLX…")
        self.dav = Music3DAV(self.model_dir / "vae/minimax_music3_dav.safetensors")

    def generate(
        self,
        caption: str,
        lyrics: str,
        seconds: float,
        steps: int,
        seed: int,
        progress: Progress | None = None,
    ) -> Any:
        mx = _mlx()
        np = _numpy()
        notify = progress or (lambda _stage, _message: None)
        requested_frames = min(MAX_AUDIO_FRAMES, max(1, round(seconds * AUDIO_FRAMES_PER_SECOND)))
        notify(1, "Generating MiniMax-Music3 acoustic tokens with MLX…")
        context = self.ar.generate(caption, lyrics, seed, requested_frames)
        notify(2, f"Sampling MiniMax-Music3 flow transformer ({steps} steps)…")
        latent = self.dit.sample(context, steps, seed)
        mx.eval(latent)
        notify(4, "Decoding MiniMax-Music3 stereo audio with DAV…")
        waveform = self.dav.decode(latent)
        mx.eval(waveform)
        samples = np.asarray(waveform, dtype=np.float32)[0]
        maximum_samples = round(context.shape[1] / AUDIO_FRAMES_PER_SECOND * SAMPLE_RATE)
        samples = samples[:, :maximum_samples]
        if not np.isfinite(samples).all():
            raise RuntimeError("MiniMax-Music3 produced non-finite audio")
        peak = float(np.max(np.abs(samples)))
        rms = float(np.sqrt(np.mean(samples.astype(np.float64) ** 2)))
        if rms <= 1e-7:
            raise RuntimeError("MiniMax-Music3 produced silent audio")
        if peak > 8.0:
            raise RuntimeError(f"MiniMax-Music3 produced an implausible peak ({peak:.3f})")
        collapse_fraction = stereo_collapse_fraction(samples)
        if collapse_fraction >= 0.75:
            raise RuntimeError(
                "MiniMax-Music3 DAV decoder produced a collapsed stereo channel "
                f"in {collapse_fraction:.0%} of the audio"
            )
        if peak > 0.99:
            samples *= 0.99 / peak
        del context, latent, waveform
        gc.collect()
        clear_cache = getattr(mx, "clear_cache", None) or getattr(
            getattr(mx, "metal", None), "clear_cache", None
        )
        if clear_cache is not None:
            clear_cache()
        return samples