Add MLX inference engine with custom Metal kernels (0.3.0)

e8ff114 verified 1 day ago

12.3 kB

"""MayaTTS on MLX: the same engine (prompt, sampler, stop rule, guards, voices, text splitting, streaming Mimi), with
the model run by maya_mac.mlx_model instead of torch MPS.
What makes it faster than the torch engine, per 80 ms frame:
  * the cb0 sampler, the 31-level depth decoder and the feedback embedding are one compiled Metal graph
    (FrameStep.step): no per-op dispatch from Python, no host round trip between the backbone and depth;
  * the frame loop is pipelined: frame t+1 is queued on the GPU before the host reads frame t (stop flag + codes, one
    33-int read), so the host-side work (guards, Mimi hand-off) overlaps the GPU instead of stalling it. A
    speculative frame is thrown away when frame t stops, or recomputed when the loop guard changes the AEOS bias;
  * a voice's reference prefix (the ICL clip) is prefilled once and its KV reused by every later request;
  * optional 8 / 4-bit weights (EngineConfig.backbone_bits / depth_bits).
Sampling draws from the same distributions as the torch engine but with MLX's RNG, so a given seed yields different
(equally likely) audio on the two backends.
"""
from __future__ import annotations

import json
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import mlx.core as mx
import numpy as np
import torch

from maya_mac import guards as G
from maya_mac import reference as S
from maya_mac.engine import HERE, EngineConfig, MayaTTS
from maya_mac.mlx_model import DTYPES, STOP_REASONS, Backbone, Composer, FrameStep, load_weights
from maya_mac.prompt import DEFAULT_BUILDER, MODES, Voice, VoiceRegistry, load_named_voices
from maya_mac.tts.tokens import N_Q


class MayaMLX(MayaTTS):
    def __init__(self, model_dir, mimi_dir, ecapa_dir=None, device=None, config: EngineConfig = None, log=print):
        from transformers import AutoTokenizer, MimiModel
        from maya_mac.codec import StreamingMimi
        from maya_mac.tts.data.packing_v2 import PackerV2
        from maya_mac.tts.tokens import extend_tokenizer

        self.device = "mlx"
        self.cfg = cfg = config or EngineConfig()
        self.log = log
        t0 = time.perf_counter()
        dtype = DTYPES[cfg.backbone_dtype]
        self.dtype = dtype
        bb = Path(model_dir) / "backbone"
        # MLX streams are thread-local: every MLX call (load, compile, voices, frames) runs on this one thread, so the
        # engine can be built on one thread and driven from another (the server loads on one and serves on another)
        self._mlx = ThreadPoolExecutor(1, thread_name_prefix="maya-mlx")
        self._mlx_thread = None

        def build():
            self._mlx_thread = threading.current_thread()
            w = load_weights(model_dir)
            self.backbone = Backbone(w, json.loads((bb / "config.json").read_text()), dtype=dtype,
                                     bits=cfg.backbone_bits, capacity=cfg.capacity)
            self.step = FrameStep(w, dtype=DTYPES[cfg.depth_dtype], bits=cfg.depth_bits, depth_temp=cfg.depth_temp,
                                  depth_top_k=cfg.depth_top_k, rep_win=cfg.cb0_rep_win, fb_dtype=dtype)
            self.compose = Composer(w, dtype=dtype)
        self._on_mlx(build)
        tokenizer, token_ids = extend_tokenizer(AutoTokenizer.from_pretrained(str(bb)))
        self.packer = PackerV2(tokenizer, token_ids, max_tokens=8192)
        self.prior = S.load_prior(str(HERE / "configs/text_frames_prior.json"))
        self.dur = G.DurationPrior()
        self.mimi = MimiModel.from_pretrained(str(mimi_dir)).to(cfg.codec_device).eval()
        self.codec = StreamingMimi(self.mimi)
        self._codec_pool = ThreadPoolExecutor(1, thread_name_prefix="maya-mimi")
        self.ecapa_dir = ecapa_dir
        self._ecapa = None
        self._ecapa_device = "cpu"
        self.registry = VoiceRegistry()
        voices_json = Path(model_dir) / "voices" / "voices.json"
        if voices_json.is_file():
            load_named_voices(self.registry, str(voices_json))
        if cfg.warmup:      # compile the frame graph and the Metal kernels at load, not on the first request
            list(self.stream("Hello, this is a warmup.", "en", max_frames=6, seed=1))
        self.load_s = time.perf_counter() - t0
        self._seed_rng = np.random.default_rng(cfg.seed)
        self.last_info = None

    def _on_mlx(self, fn, *args):
        if threading.current_thread() is self._mlx_thread:
            return fn(*args)
        return self._mlx.submit(fn, *args).result()

    def close(self):
        super().close()
        self._mlx.shutdown(wait=True)

    # ------------------------------------------------------------------ voices
    def prepare_voice(self, voice: Voice):
        """Pack the voice's reference prefix and compute its depth cross-attention constants (once per voice)."""
        return self._on_mlx(self._prepare_voice, voice)

    def _prepare_voice(self, voice):
        with voice.lock:
            if voice.mem is not None:
                return voice
            _, use_mem = MODES[voice.mode]
            builder = voice.builder or DEFAULT_BUILDER
            voice.prefix = builder.prefix(self.packer, voice)
            voice.spk_refs = (np.asarray(voice.ecapa, np.float32).reshape(1, 192) if voice.ecapa is not None
                              else np.zeros((1, 192), np.float32))
            voice.mem = self.step.cross(voice.spk_refs if use_mem else None)
            mx.eval(voice.mem)
            voice.ts = int(voice.prefix["tokens"].shape[0]) if voice.prefix is not None else 0
            voice.kv = voice.emb = None
            return voice

    def _embed(self, w):
        if self._ecapa is None:
            if self.ecapa_dir is None:
                raise ValueError("ecapa_dir is required for reference-audio voices")
            from maya_mac.engine import Ecapa
            self._ecapa = Ecapa(self.ecapa_dir, self._ecapa_device)
        return self._ecapa(w)

    # ------------------------------------------------------------------ synthesis
    def _plan(self, text, lang, voice, max_frames):
        from maya_mac.tts.data.packing_v2 import collate_v2
        from maya_mac.tts.model.load_s3 import truncate_prompt_batch
        cfg, g = self.cfg, self.cfg.guards
        builder = voice.builder or DEFAULT_BUILDER
        ex = builder.example(self.packer, voice, text, lang)
        # reference.prompt_from_example, composed by MLX
        b = collate_v2([ex], q_choices=(N_Q,), q_rng=np.random.default_rng(0))
        kind = b["kind"][0]
        a = (kind == 1).nonzero().flatten()
        a = a[a >= int(ex.get("target_start", 0))]
        assert a.numel(), "prompt has no target audio position"
        prompt = self.compose(truncate_prompt_batch(b, int(a[0]), require_target_span=self.compose.anchor is not None))
        ef, er = S.expected_frames(self.prior, lang, len(text))
        loose = min(cfg.hard_cap, int(2.0 * (ef + er) + 40))
        cap = loose
        if g.enabled:
            cap = min(cap, max(int(ef * 0.5) + 8, self.dur.cap_frames(text, lang, g.dur_k, g.dur_margin_s)))
        if max_frames:
            cap = min(cap, max(1, int(max_frames)))
        minimum = max(4, int(0.5 * ef))
        suspend = G.suspend_windows_v2(text, lang, self.dur) if (g.enabled and g.loop and g.aware_v2) else ()
        return prompt, cap, loose, minimum, suspend

    def _prefill(self, prompt, voice):
        """Run the prompt through the backbone, reusing the voice's cached reference-prefix KV when the prompt starts
        with exactly that prefix."""
        bb, ts = self.backbone, voice.ts
        bb.reset()
        if ts and ts < prompt.shape[1]:
            head = prompt[:, :ts]
            if voice.kv is not None and mx.array_equal(head, voice.emb).item():
                bb.restore(voice.kv)
                return bb(prompt[:, ts:])
            bb(head)
            voice.kv, voice.emb = bb.snapshot(ts), head
            mx.eval(voice.kv, voice.emb)   # materialise now, so later in-place cache writes don't have to copy the cache
            return bb(prompt[:, ts:])
        return bb(prompt)

    def _frames(self, text, lang, voice, sampling, max_frames, info):
        """Yield the (32,) code tensor of each frame of one text piece; info receives the stop record. Each frame is
        produced on the MLX thread."""
        gen = self._frames_mlx(text, lang, voice, sampling, max_frames, info)
        try:
            while True:
                try:
                    codes = self._on_mlx(next, gen)
                except StopIteration:
                    return
                yield codes
        finally:
            self._on_mlx(gen.close)

    def _frames_mlx(self, text, lang, voice, sampling, max_frames, info):
        cfg, g = self.cfg, self.cfg.guards
        prompt, cap, loose, minimum, suspend = self._plan(text, lang, voice, max_frames)
        seed = sampling.get("seed")
        seed = int(seed) if seed is not None else int(self._seed_rng.integers(1, 2 ** 31 - 1))
        if prompt.shape[1] + cap + 1 > cfg.capacity:
            raise ValueError(f"prompt of {prompt.shape[1]} tokens + {cap} frames exceeds KV capacity {cfg.capacity}")
        rp = sampling.get("repetition_penalty")
        knobs = (mx.array(float(sampling.get("temperature") or cfg.cb0_temp)),
                 mx.array(min(int(sampling.get("top_k") or cfg.cb0_top_k), 2049)),
                 mx.array(float(sampling.get("top_p") or cfg.cb0_top_p)),
                 mx.array(float(rp if rp is not None else cfg.cb0_rep)))
        cross = voice.mem
        keys = iter(mx.random.split(mx.random.key(seed), cap + 1))
        backbone = self.backbone

        def issue(h, hazard, hist, frame, bias, key):
            """Queue one frame on the GPU; nothing is read back here."""
            out, fb, hazard, hist = self.step.step(h, cross, hist, hazard, mx.array(frame >= minimum),
                                                   mx.array(bias, mx.float32), *knobs, key)
            h = backbone(fb[None, None])
            mx.async_eval(out, h, hazard, hist)
            return out, h, hazard, hist, key

        h = self._prefill(prompt, voice)
        hist0 = mx.full((max(1, cfg.cb0_rep_win),), -1, mx.int32)
        cur = issue(h, mx.array(0.0), hist0, 0, 0.0, next(keys))
        cb0_hist = np.full(g.hist, -1, np.int64)
        frames, reason, boost_at, n_boost, bias = 0, "cap", None, 0, 0.0
        while True:
            nxt = issue(*cur[1:4], frames + 1, bias, next(keys)) if frames + 1 < cap else None
            out = np.array(cur[0])                      # waits for this frame only; the next one is already queued
            stop = int(out[0])
            if stop:
                reason = STOP_REASONS[stop]
                break
            codes = torch.from_numpy(out[1:].astype(np.int64))
            yield codes
            frames += 1
            cb0_hist[:-1] = cb0_hist[1:]
            cb0_hist[-1] = int(out[1])
            if frames >= cap:
                reason = "max_frames" if max_frames and frames >= max_frames else (
                    "cap_guard" if g.enabled and cap < loose else "cap")
                break
            new_bias = bias
            if (g.enabled and g.loop and frames >= g.start and frames % g.every == 0
                    and not any(a <= frames <= b for a, b in suspend)):
                hit = G.loop_detect_batch(cb0_hist[None], np.array([frames]), g)[0]
                if hit is not None and g.mode == "boost":
                    if boost_at is None:
                        boost_at, n_boost = frames, n_boost + 1
                        new_bias = g.boost
                    if frames - boost_at >= g.backstop:
                        reason = "loop_backstop"
                        break
                elif hit is not None:
                    reason = "loop_" + hit.split(":")[0]
                    break
                elif boost_at is not None:
                    new_bias, boost_at = 0.0, None
            if new_bias != bias:      # the queued frame used the old AEOS bias: redo it (it rewrites the same KV slot)
                bias = new_bias
                nxt = issue(*cur[1:4], frames, bias, nxt[4])
            cur = nxt
        info.update({"reason": reason, "frames": frames, "cap": cap, "n_boost": n_boost})

	"""MayaTTS on MLX: the same engine (prompt, sampler, stop rule, guards, voices, text splitting, streaming Mimi), with
	the model run by maya_mac.mlx_model instead of torch MPS.

	What makes it faster than the torch engine, per 80 ms frame:
	* the cb0 sampler, the 31-level depth decoder and the feedback embedding are one compiled Metal graph
	(FrameStep.step): no per-op dispatch from Python, no host round trip between the backbone and depth;
	* the frame loop is pipelined: frame t+1 is queued on the GPU before the host reads frame t (stop flag + codes, one
	33-int read), so the host-side work (guards, Mimi hand-off) overlaps the GPU instead of stalling it. A
	speculative frame is thrown away when frame t stops, or recomputed when the loop guard changes the AEOS bias;
	* a voice's reference prefix (the ICL clip) is prefilled once and its KV reused by every later request;
	* optional 8 / 4-bit weights (EngineConfig.backbone_bits / depth_bits).

	Sampling draws from the same distributions as the torch engine but with MLX's RNG, so a given seed yields different
	(equally likely) audio on the two backends.
	"""
	from __future__ import annotations

	import json
	import threading
	import time
	from concurrent.futures import ThreadPoolExecutor
	from pathlib import Path

	import mlx.core as mx
	import numpy as np
	import torch

	from maya_mac import guards as G
	from maya_mac import reference as S
	from maya_mac.engine import HERE, EngineConfig, MayaTTS
	from maya_mac.mlx_model import DTYPES, STOP_REASONS, Backbone, Composer, FrameStep, load_weights
	from maya_mac.prompt import DEFAULT_BUILDER, MODES, Voice, VoiceRegistry, load_named_voices
	from maya_mac.tts.tokens import N_Q


	class MayaMLX(MayaTTS):
	def __init__(self, model_dir, mimi_dir, ecapa_dir=None, device=None, config: EngineConfig = None, log=print):
	from transformers import AutoTokenizer, MimiModel
	from maya_mac.codec import StreamingMimi
	from maya_mac.tts.data.packing_v2 import PackerV2
	from maya_mac.tts.tokens import extend_tokenizer

	self.device = "mlx"
	self.cfg = cfg = config or EngineConfig()
	self.log = log
	t0 = time.perf_counter()
	dtype = DTYPES[cfg.backbone_dtype]
	self.dtype = dtype
	bb = Path(model_dir) / "backbone"
	# MLX streams are thread-local: every MLX call (load, compile, voices, frames) runs on this one thread, so the
	# engine can be built on one thread and driven from another (the server loads on one and serves on another)
	self._mlx = ThreadPoolExecutor(1, thread_name_prefix="maya-mlx")
	self._mlx_thread = None

	def build():
	self._mlx_thread = threading.current_thread()
	w = load_weights(model_dir)
	self.backbone = Backbone(w, json.loads((bb / "config.json").read_text()), dtype=dtype,
	bits=cfg.backbone_bits, capacity=cfg.capacity)
	self.step = FrameStep(w, dtype=DTYPES[cfg.depth_dtype], bits=cfg.depth_bits, depth_temp=cfg.depth_temp,
	depth_top_k=cfg.depth_top_k, rep_win=cfg.cb0_rep_win, fb_dtype=dtype)
	self.compose = Composer(w, dtype=dtype)
	self._on_mlx(build)
	tokenizer, token_ids = extend_tokenizer(AutoTokenizer.from_pretrained(str(bb)))
	self.packer = PackerV2(tokenizer, token_ids, max_tokens=8192)
	self.prior = S.load_prior(str(HERE / "configs/text_frames_prior.json"))
	self.dur = G.DurationPrior()
	self.mimi = MimiModel.from_pretrained(str(mimi_dir)).to(cfg.codec_device).eval()
	self.codec = StreamingMimi(self.mimi)
	self._codec_pool = ThreadPoolExecutor(1, thread_name_prefix="maya-mimi")
	self.ecapa_dir = ecapa_dir
	self._ecapa = None
	self._ecapa_device = "cpu"
	self.registry = VoiceRegistry()
	voices_json = Path(model_dir) / "voices" / "voices.json"
	if voices_json.is_file():
	load_named_voices(self.registry, str(voices_json))
	if cfg.warmup: # compile the frame graph and the Metal kernels at load, not on the first request
	list(self.stream("Hello, this is a warmup.", "en", max_frames=6, seed=1))
	self.load_s = time.perf_counter() - t0
	self._seed_rng = np.random.default_rng(cfg.seed)
	self.last_info = None

	def _on_mlx(self, fn, *args):
	if threading.current_thread() is self._mlx_thread:
	return fn(*args)
	return self._mlx.submit(fn, *args).result()

	def close(self):
	super().close()
	self._mlx.shutdown(wait=True)

	# ------------------------------------------------------------------ voices
	def prepare_voice(self, voice: Voice):
	"""Pack the voice's reference prefix and compute its depth cross-attention constants (once per voice)."""
	return self._on_mlx(self._prepare_voice, voice)

	def _prepare_voice(self, voice):
	with voice.lock:
	if voice.mem is not None:
	return voice
	_, use_mem = MODES[voice.mode]
	builder = voice.builder or DEFAULT_BUILDER
	voice.prefix = builder.prefix(self.packer, voice)
	voice.spk_refs = (np.asarray(voice.ecapa, np.float32).reshape(1, 192) if voice.ecapa is not None
	else np.zeros((1, 192), np.float32))
	voice.mem = self.step.cross(voice.spk_refs if use_mem else None)
	mx.eval(voice.mem)
	voice.ts = int(voice.prefix["tokens"].shape[0]) if voice.prefix is not None else 0
	voice.kv = voice.emb = None
	return voice

	def _embed(self, w):
	if self._ecapa is None:
	if self.ecapa_dir is None:
	raise ValueError("ecapa_dir is required for reference-audio voices")
	from maya_mac.engine import Ecapa
	self._ecapa = Ecapa(self.ecapa_dir, self._ecapa_device)
	return self._ecapa(w)

	# ------------------------------------------------------------------ synthesis
	def _plan(self, text, lang, voice, max_frames):
	from maya_mac.tts.data.packing_v2 import collate_v2
	from maya_mac.tts.model.load_s3 import truncate_prompt_batch
	cfg, g = self.cfg, self.cfg.guards
	builder = voice.builder or DEFAULT_BUILDER
	ex = builder.example(self.packer, voice, text, lang)
	# reference.prompt_from_example, composed by MLX
	b = collate_v2([ex], q_choices=(N_Q,), q_rng=np.random.default_rng(0))
	kind = b["kind"][0]
	a = (kind == 1).nonzero().flatten()
	a = a[a >= int(ex.get("target_start", 0))]
	assert a.numel(), "prompt has no target audio position"
	prompt = self.compose(truncate_prompt_batch(b, int(a[0]), require_target_span=self.compose.anchor is not None))
	ef, er = S.expected_frames(self.prior, lang, len(text))
	loose = min(cfg.hard_cap, int(2.0 * (ef + er) + 40))
	cap = loose
	if g.enabled:
	cap = min(cap, max(int(ef * 0.5) + 8, self.dur.cap_frames(text, lang, g.dur_k, g.dur_margin_s)))
	if max_frames:
	cap = min(cap, max(1, int(max_frames)))
	minimum = max(4, int(0.5 * ef))
	suspend = G.suspend_windows_v2(text, lang, self.dur) if (g.enabled and g.loop and g.aware_v2) else ()
	return prompt, cap, loose, minimum, suspend

	def _prefill(self, prompt, voice):
	"""Run the prompt through the backbone, reusing the voice's cached reference-prefix KV when the prompt starts
	with exactly that prefix."""
	bb, ts = self.backbone, voice.ts
	bb.reset()
	if ts and ts < prompt.shape[1]:
	head = prompt[:, :ts]
	if voice.kv is not None and mx.array_equal(head, voice.emb).item():
	bb.restore(voice.kv)
	return bb(prompt[:, ts:])
	bb(head)
	voice.kv, voice.emb = bb.snapshot(ts), head
	mx.eval(voice.kv, voice.emb) # materialise now, so later in-place cache writes don't have to copy the cache
	return bb(prompt[:, ts:])
	return bb(prompt)

	def _frames(self, text, lang, voice, sampling, max_frames, info):
	"""Yield the (32,) code tensor of each frame of one text piece; info receives the stop record. Each frame is
	produced on the MLX thread."""
	gen = self._frames_mlx(text, lang, voice, sampling, max_frames, info)
	try:
	while True:
	try:
	codes = self._on_mlx(next, gen)
	except StopIteration:
	return
	yield codes
	finally:
	self._on_mlx(gen.close)

	def _frames_mlx(self, text, lang, voice, sampling, max_frames, info):
	cfg, g = self.cfg, self.cfg.guards
	prompt, cap, loose, minimum, suspend = self._plan(text, lang, voice, max_frames)
	seed = sampling.get("seed")
	seed = int(seed) if seed is not None else int(self._seed_rng.integers(1, 2 ** 31 - 1))
	if prompt.shape[1] + cap + 1 > cfg.capacity:
	raise ValueError(f"prompt of {prompt.shape[1]} tokens + {cap} frames exceeds KV capacity {cfg.capacity}")
	rp = sampling.get("repetition_penalty")
	knobs = (mx.array(float(sampling.get("temperature") or cfg.cb0_temp)),
	mx.array(min(int(sampling.get("top_k") or cfg.cb0_top_k), 2049)),
	mx.array(float(sampling.get("top_p") or cfg.cb0_top_p)),
	mx.array(float(rp if rp is not None else cfg.cb0_rep)))
	cross = voice.mem
	keys = iter(mx.random.split(mx.random.key(seed), cap + 1))
	backbone = self.backbone

	def issue(h, hazard, hist, frame, bias, key):
	"""Queue one frame on the GPU; nothing is read back here."""
	out, fb, hazard, hist = self.step.step(h, cross, hist, hazard, mx.array(frame >= minimum),
	mx.array(bias, mx.float32), *knobs, key)
	h = backbone(fb[None, None])
	mx.async_eval(out, h, hazard, hist)
	return out, h, hazard, hist, key

	h = self._prefill(prompt, voice)
	hist0 = mx.full((max(1, cfg.cb0_rep_win),), -1, mx.int32)
	cur = issue(h, mx.array(0.0), hist0, 0, 0.0, next(keys))
	cb0_hist = np.full(g.hist, -1, np.int64)
	frames, reason, boost_at, n_boost, bias = 0, "cap", None, 0, 0.0
	while True:
	nxt = issue(*cur[1:4], frames + 1, bias, next(keys)) if frames + 1 < cap else None
	out = np.array(cur[0]) # waits for this frame only; the next one is already queued
	stop = int(out[0])
	if stop:
	reason = STOP_REASONS[stop]
	break
	codes = torch.from_numpy(out[1:].astype(np.int64))
	yield codes
	frames += 1
	cb0_hist[:-1] = cb0_hist[1:]
	cb0_hist[-1] = int(out[1])
	if frames >= cap:
	reason = "max_frames" if max_frames and frames >= max_frames else (
	"cap_guard" if g.enabled and cap < loose else "cap")
	break
	new_bias = bias
	if (g.enabled and g.loop and frames >= g.start and frames % g.every == 0
	and not any(a <= frames <= b for a, b in suspend)):
	hit = G.loop_detect_batch(cb0_hist[None], np.array([frames]), g)[0]
	if hit is not None and g.mode == "boost":
	if boost_at is None:
	boost_at, n_boost = frames, n_boost + 1
	new_bias = g.boost
	if frames - boost_at >= g.backstop:
	reason = "loop_backstop"
	break
	elif hit is not None:
	reason = "loop_" + hit.split(":")[0]
	break
	elif boost_at is not None:
	new_bias, boost_at = 0.0, None
	if new_bias != bias: # the queued frame used the old AEOS bias: redo it (it rewrites the same KV slot)
	bias = new_bias
	nxt = issue(*cur[1:4], frames, bias, nxt[4])
	cur = nxt
	info.update({"reason": reason, "frames": frames, "cap": cap, "n_boost": n_boost})