"""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()