mlx_bridge.py

"""Pont MLX pour le language_model de MiniMax-Music3.

Le pipeline diffusers n'utilise du Qwen3ForCausalLM que trois points de contact :

    language_model.model.embed_tokens(ids)
    language_model.model(inputs_embeds=..., past_key_values=..., use_cache=True)
    language_model.lm_head(hidden)

Cette classe les reimplemente sur un modele MLX quantifie. Le reste du pipeline
(transformer flow-matching, RVQ depth decoder, vocoder) reste en torch/MPS.
Les tenseurs echanges par pas de decodage sont minuscules, la conversion via numpy
ne pese rien face au forward du modele.
"""

import time
from collections import defaultdict
from pathlib import Path
from types import SimpleNamespace

import mlx.core as mx
import numpy as np
import torch
from mlx_lm.models.cache import make_prompt_cache
from mlx_lm.utils import load_model

STATS = defaultdict(float)
COUNTS = defaultdict(int)


class timed:
    def __init__(self, key):
        self.key = key

    def __enter__(self):
        self.t = time.perf_counter()

    def __exit__(self, *exc):
        STATS[self.key] += time.perf_counter() - self.t
        COUNTS[self.key] += 1


def report():
    if not STATS:
        return
    width = max(len(k) for k in STATS)
    print("\nprofil MLX bridge:")
    for key in sorted(STATS, key=STATS.get, reverse=True):
        print(f"  {key:<{width}}  {STATS[key]:7.1f}s  {COUNTS[key]:>6} appels")


def _to_mlx(tensor, dtype=mx.bfloat16):
    return mx.array(tensor.detach().float().cpu().numpy()).astype(dtype)


def _to_torch(array, device, dtype):
    return torch.from_numpy(np.array(array.astype(mx.float32))).to(device=device, dtype=dtype)


class _Embedding:
    def __init__(self, bridge):
        self._bridge = bridge

    def __call__(self, ids):
        b = self._bridge
        with timed("embed"):
            flat = mx.array(ids.detach().cpu().numpy().reshape(-1).astype(np.int64))
            out = b.mlx_model.model.embed_tokens(flat)
            out = _to_torch(out, b.device, b.dtype)
        return out.reshape(*ids.shape, -1)


class _Inner:
    """Expose `.embed_tokens` et l'appel forward attendus par le pipeline."""

    def __init__(self, bridge):
        self._bridge = bridge
        self.embed_tokens = _Embedding(bridge)

    def __call__(self, inputs_embeds=None, past_key_values=None, use_cache=True, **kwargs):
        b = self._bridge
        cache = past_key_values
        if cache is None:
            cache = make_prompt_cache(b.mlx_model)
        with timed("lm.in"):
            h = _to_mlx(inputs_embeds)
        with timed("lm.forward"):
            out = b.mlx_model.model(None, cache=cache, input_embeddings=h)
            mx.eval(out)
        with timed("lm.out"):
            hidden = _to_torch(out, b.device, b.dtype)
        return SimpleNamespace(last_hidden_state=hidden, past_key_values=cache)


class MlxLanguageModel:
    def __init__(self, mlx_path, device, dtype=torch.bfloat16):
        self.mlx_model, self.mlx_config = load_model(Path(mlx_path), lazy=False)
        self.device = torch.device(device)
        self.dtype = dtype
        self.config = SimpleNamespace(**self.mlx_config)
        self.model = _Inner(self)
        self._last_hidden_mlx = None

    def lm_head(self, hidden):
        with timed("lm_head"):
            logits = self.mlx_model.lm_head(_to_mlx(hidden))
            mx.eval(logits)
        with timed("lm_head.out"):
            return _to_torch(logits, self.device, torch.float32)

    def to(self, *args, **kwargs):
        return self

    def eval(self):
        return self