scripts/moss_audio_mlx_bridge_v3.py

"""Pure-MLX MOSS-Audio hybrid bridge (v3). No PyTorch at runtime.

Uses:
  - MLX INT4 Qwen3 LLM (existing moss4b_mlx_int4/)
  - MLX BF16 MossAudioEncoder (ported to mlx.nn)
  - MLX BF16 GatedMLP for audio_adapter + deepstack_audio_merger_list
  - MossAudioProcessor from upstream repo (CPU/numpy only, used to compute mel spectrogram)

All compute on MLX; only mel-spectrogram construction uses CPU.
"""
from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path.home() / "benchmark" / "moss-audio"))
sys.path.insert(0, str(Path.home() / "benchmark" / "scripts"))

import librosa
import mlx.core as mx
import mlx.nn as mnn
import numpy as np
from mlx_lm import load as mlx_load
from mlx_lm.generate import generate_step
from mlx_lm.sample_utils import make_sampler, make_logits_processors

from moss_audio_encoder_mlx import MossAudioEncoderMLX, EncoderConfig, GatedMLP


# ---- Load MLX encoder + adapter + mergers ----


def _infer_llm_hidden(ad_w: dict) -> int:
    """Sniff the LLM hidden dim from the adapter's down_proj.

    Supports both BF16 weights (`down_proj.weight` shape = (llm_hidden, 8192))
    and INT4-quantized weights (`down_proj.scales` shape = (llm_hidden, n_groups)).
    Returns 2560 for 4B, 4096 for 8B.
    """
    if "down_proj.scales" in ad_w:
        return ad_w["down_proj.scales"].shape[0]
    return ad_w["down_proj.weight"].shape[0]


def load_mlx_audio_path(weights_dir: Path, *, int4: bool = False):
    """Load encoder + adapter + mergers into MLX.

    Adapter/merger output dim (LLM hidden) is inferred from the saved weights,
    so the same loader works for both 4B (2560) and 8B (4096).

    When int4=True, weights_dir is expected to contain already-quantized
    INT4 safetensors (see scripts/save_moss_audio_int4.py). This avoids
    the double-buffer memory cost of live quantization from BF16.
    """
    ad_w = mx.load(str(weights_dir / "audio_adapter.safetensors"))
    llm_hidden = _infer_llm_hidden(ad_w)

    cfg = EncoderConfig()
    enc = MossAudioEncoderMLX(cfg)
    adapter = GatedMLP(1280, 8192, llm_hidden)
    mergers = [GatedMLP(1280, 8192, llm_hidden) for _ in range(3)]

    if int4:
        # Build quantized module structure FIRST, then load INT4 weights directly.
        # This avoids holding BF16 + INT4 copies in memory simultaneously.
        mnn.quantize(enc, group_size=64, bits=4)
        mnn.quantize(adapter, group_size=64, bits=4)
        for m in mergers:
            mnn.quantize(m, group_size=64, bits=4)

    enc_w = mx.load(str(weights_dir / "audio_encoder.safetensors"))
    enc.load_weights(list(enc_w.items()), strict=True)

    adapter.load_weights(list(ad_w.items()), strict=True)

    dm_w = mx.load(str(weights_dir / "deepstack_mergers.safetensors"))
    for i, merger in enumerate(mergers):
        subw = {k[len(f"{i}."):]: v for k, v in dm_w.items() if k.startswith(f"{i}.")}
        merger.load_weights(list(subw.items()), strict=True)

    mx.eval(enc.parameters(), adapter.parameters(), *[m.parameters() for m in mergers])
    return enc, adapter, mergers


def build_mel_spectrogram(audio_np: np.ndarray, processor_or_tokenizer, *, fp32_compute: bool = True) -> tuple[mx.array, mx.array]:
    """Build mel spectrogram + audio-expanded input_ids.

    Two supported inputs for `processor_or_tokenizer`:
      - A HuggingFace tokenizer (pure-Python, no torch) — pure-MLX path. Default.
      - A MossAudioProcessor instance — delegates to upstream (kept for the
        parity regression tests and anyone pinning to the old behavior).

    `fp32_compute=True` (default) keeps the mel input as fp32, which causes
    MLX to auto-promote all encoder activations to fp32 during forward
    (weights stay BF16 on disk). This cuts relative error in the adapter
    output from 6.5% to 0.98% vs PyTorch reference, at the cost of ~1 GB
    extra peak memory during the one-shot encoder pass.
    """
    prompt = ("Describe this audio in detail. Include speech content, speaker "
              "characteristics, background sounds, music, and any notable temporal events.")

    # Pure-MLX path: the object has `.encode()` but not MossAudioProcessor's `audio_token_id`
    # set alongside a `.tokenizer` attribute. Detect by method presence.
    is_tokenizer = hasattr(processor_or_tokenizer, "encode") and not hasattr(processor_or_tokenizer, "audio_token_id")

    if is_tokenizer:
        from moss_audio_mel_mlx import build_mel_and_input_ids
        mel_mx, lens_mx, input_ids_mx, audio_token_id = build_mel_and_input_ids(
            audio_np, processor_or_tokenizer, prompt=prompt, enable_time_marker=True,
        )
        if not fp32_compute:
            mel_mx = mel_mx.astype(mx.bfloat16)
        return mel_mx, lens_mx, input_ids_mx, audio_token_id

    # Legacy torch processor path (unchanged)
    processor = processor_or_tokenizer
    inputs = processor(text=prompt, audios=[audio_np], return_tensors="pt")
    import torch
    mel = inputs["audio_data"]
    mel_np = mel.to(torch.float32).numpy()
    mel_mx = mx.array(mel_np) if fp32_compute else mx.array(mel_np).astype(mx.bfloat16)
    lens = inputs["audio_data_seqlens"].cpu().numpy().astype(np.int32) if inputs.get("audio_data_seqlens") is not None else None
    lens_mx = mx.array(lens) if lens is not None else None
    input_ids_np = inputs["input_ids"].cpu().numpy()
    input_ids_mx = mx.array(input_ids_np)
    audio_token_id = processor.audio_token_id
    return mel_mx, lens_mx, input_ids_mx, audio_token_id


def run_mlx_audio_pipeline(encoder, adapter, mergers,
                           mel: mx.array, lens: mx.array):
    """Returns (primary_embeds, deepstack_embeds) all on MLX."""
    last, deepstack = encoder(mel, feature_lens=lens, return_deepstack=True)
    primary = adapter(last)  # (B, N_audio, llm_hidden)
    ds_embeds = [mergers[i](ds) for i, ds in enumerate(deepstack)]
    return primary, ds_embeds


# ---- DeepStack injection on MLX decoder ----


def install_deepstack_hooks(mlx_model, deepstack_embeds: list[mx.array], audio_positions: np.ndarray):
    """Same class-level Qwen3Model.__call__ override as v2, but with MLX-native inputs."""
    from mlx_lm.models.base import create_attention_mask

    audio_positions_mx = mx.array(audio_positions.astype(np.int32))
    num_inject = len(deepstack_embeds)

    ModelCls = type(mlx_model.model)
    orig_call = ModelCls.__call__

    def new_call(self, inputs, cache=None, input_embeddings=None):
        if self is not mlx_model.model:
            return orig_call(self, inputs, cache, input_embeddings)

        if input_embeddings is not None:
            h = input_embeddings
        else:
            h = self.embed_tokens(inputs)

        if cache is None:
            cache = [None] * len(self.layers)
        mask = create_attention_mask(h, cache[0])
        is_prefill = h.shape[1] > 1

        for layer_idx, (layer, c) in enumerate(zip(self.layers, cache)):
            h = layer(h, mask, c)
            if is_prefill and layer_idx < num_inject:
                ds = deepstack_embeds[layer_idx]
                if ds.dtype != h.dtype:
                    ds = ds.astype(h.dtype)
                if ds.ndim == 3:
                    ds = ds[0]  # flatten batch
                h_batch0 = h[0]
                h_batch0 = h_batch0.at[audio_positions_mx].add(ds)
                h = h_batch0[None]

        return self.norm(h)

    ModelCls.__call__ = new_call


# ---- Main ----


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mlx-llm", default=str(Path.home() / "benchmark" / "moss4b_mlx_int4"))
    parser.add_argument("--mlx-audio", default=str(Path.home() / "benchmark" / "moss4b_audio_mlx_int4"))
    parser.add_argument("--moss-source", default="OpenMOSS-Team/MOSS-Audio-4B-Thinking",
                        help="HF repo for the processor only (mel spectrogram computation)")
    parser.add_argument("--audio", required=True)
    parser.add_argument("--max-tokens", type=int, default=2048)
    parser.add_argument("--temp", type=float, default=1.0)
    parser.add_argument("--int4-audio", action="store_true", default=True,
                        help="Load pre-quantized INT4 audio encoder + adapter + mergers from --mlx-audio")
    parser.add_argument("--no-int4-audio", dest="int4_audio", action="store_false",
                        help="Use BF16 audio weights instead of INT4")
    parser.add_argument("--repetition-penalty", type=float, default=1.02,
                        help="mlx_lm repetition_penalty. 1.02 is the shipped EN-scope default "
                             "(kills loop-decode on non-speech clips without starving genre descriptions).")
    parser.add_argument("--repetition-context-size", type=int, default=20)
    parser.add_argument("--use-torch-processor", action="store_true",
                        help="Use upstream MossAudioProcessor (torch+torchaudio) for mel "
                             "computation. Default off: pure-MLX mel path has <0.2%% "
                             "rel-err parity and no torch dep.")
    args = parser.parse_args()

    print(f"[mlx] loading LLM from {args.mlx_llm}", flush=True)
    t0 = time.perf_counter()
    mlx_model, mlx_tokenizer = mlx_load(args.mlx_llm)
    print(f"[mlx] LLM loaded in {time.perf_counter()-t0:.1f}s  peak={mx.get_peak_memory()/1e9:.2f}GB", flush=True)

    print(f"[mlx] loading audio path from {args.mlx_audio} (int4={args.int4_audio})", flush=True)
    t0 = time.perf_counter()
    encoder, adapter, mergers = load_mlx_audio_path(Path(args.mlx_audio), int4=args.int4_audio)
    print(f"[mlx] audio path loaded in {time.perf_counter()-t0:.1f}s  peak={mx.get_peak_memory()/1e9:.2f}GB", flush=True)

    # Pure-MLX path: use the MLX-LLM tokenizer directly for text encoding.
    # Mel spectrogram + input_ids expansion are computed in pure MLX (no torch).
    # The upstream MossAudioProcessor is only needed if you opt into the legacy
    # torch mel path via `--use-torch-processor` (not on by default).
    processor = mlx_tokenizer
    if args.use_torch_processor:
        from src.processing_moss_audio import MossAudioProcessor
        processor = MossAudioProcessor.from_pretrained(
            args.moss_source, trust_remote_code=True, enable_time_marker=True,
        )

    y, _ = librosa.load(args.audio, sr=16000, mono=True)
    y = y.astype(np.float32)
    print(f"[audio] {args.audio}: {len(y)/16000:.1f}s", flush=True)

    t0 = time.perf_counter()
    mel, lens, input_ids_mx, audio_token_id = build_mel_spectrogram(y, processor)
    primary, ds_embeds = run_mlx_audio_pipeline(encoder, adapter, mergers, mel, lens)
    # Cast to bf16 and force materialization NOW, so the fp32 activations can be freed
    # before we start the decode phase (the encoder's fp32 forward is our peak-memory hotspot).
    primary = primary.astype(mx.bfloat16)
    ds_embeds = [d.astype(mx.bfloat16) for d in ds_embeds]
    mx.eval(primary, *ds_embeds)
    print(f"[mlx] audio encoded in {time.perf_counter()-t0:.2f}s  primary={primary.shape}", flush=True)
    print(f"[mem] after encode (pre-cleanup): peak={mx.get_peak_memory()/1e9:.2f}GB active={mx.get_active_memory()/1e9:.2f}GB", flush=True)

    # FREE ENCODER + ADAPTER + MERGERS now that we have the embeddings.
    # These modules are only needed at prefill; all downstream work (merge,
    # deepstack injection, decode) uses only primary + ds_embeds tensors.
    # Dropping them reclaims ~1.3 GB for clip_01 (up to ~2.5 GB if fp32
    # activations were still held). Captured closures for install_deepstack_hooks
    # keep only the embedding mx.array references, not the modules.
    del encoder, adapter, mergers, mel, lens
    import gc; gc.collect()
    mx.clear_cache()
    try: mx.reset_peak_memory()
    except Exception: pass
    print(f"[mem] after free encoder: peak={mx.get_peak_memory()/1e9:.2f}GB active={mx.get_active_memory()/1e9:.2f}GB", flush=True)

    # Build merged text+audio embeddings
    audio_mask = input_ids_mx == audio_token_id
    audio_positions = np.where(np.array(audio_mask[0]))[0]
    assert len(audio_positions) == primary.shape[1], \
        f"mask positions {len(audio_positions)} != primary len {primary.shape[1]}"

    text_embeds = mlx_model.model.embed_tokens(input_ids_mx)  # (1, seq, hidden)
    text_np = np.array(text_embeds.astype(mx.float32))
    primary_np = np.array(primary.astype(mx.float32))
    text_np[0, audio_positions, :] = primary_np[0, :, :]
    merged = mx.array(text_np).astype(mx.bfloat16)
    print(f"[bridge] merged embeds {merged.shape}", flush=True)

    # Install DeepStack
    ds_flat = [d[0] for d in ds_embeds]  # drop batch dim
    install_deepstack_hooks(mlx_model, ds_flat, audio_positions)
    print(f"[deepstack] installed {len(ds_flat)} layer injections", flush=True)

    # Generate
    print(f"[gen] decoding max_tokens={args.max_tokens}...", flush=True)
    sampler = make_sampler(temp=args.temp, top_p=1.0, top_k=50)
    gen_kwargs = dict(
        prompt=input_ids_mx[0],
        model=mlx_model,
        input_embeddings=merged[0],
        max_tokens=args.max_tokens,
        sampler=sampler,
    )
    if args.repetition_penalty:
        gen_kwargs["logits_processors"] = make_logits_processors(
            repetition_penalty=args.repetition_penalty,
            repetition_context_size=args.repetition_context_size,
        )
        print(f"[gen] repetition_penalty={args.repetition_penalty} ctx={args.repetition_context_size}", flush=True)
    t0 = time.perf_counter()
    generated = []
    ttft = None
    for tok, _ in generate_step(**gen_kwargs):
        if ttft is None:
            ttft = time.perf_counter() - t0
        generated.append(int(tok))
        if tok == mlx_tokenizer.eos_token_id:
            break
    elapsed = time.perf_counter() - t0
    n_tok = len(generated)
    decode_s = max(elapsed - (ttft or 0), 1e-6)
    print(f"[gen] elapsed={elapsed:.2f}s  ttft={ttft:.3f}s  out={n_tok}tok  decode={max(n_tok-1,0)/decode_s:.1f}t/s", flush=True)
    text = mlx_tokenizer.decode(generated)
    print(f"\n=== OUTPUT ===\n{text[:1500]}\n=== END ===\n")

    print(f"[mem] mlx peak={mx.get_peak_memory()/1e9:.2f}GB", flush=True)


if __name__ == "__main__":
    main()