| """Queued engine for the server: the interface maya_mac.server expects (upstream maya3.engine.Engine's), on top of |
| MayaTTS. |
| |
| One worker thread owns the model and runs requests one at a time in arrival order (MayaTTS generates a single stream, |
| faster than real time on Apple Silicon with MLX; there is no batching). Voice registration also runs on that |
| thread, between requests, so the GPU is only ever driven from one thread. |
| |
| engine = ServingEngine(MayaTTS.from_pretrained()) |
| req = engine.submit("Hello there.", "en", "Ahana", on_event=lambda kind, payload: ...) |
| req.cancel() |
| Events, called from the worker thread (keep them cheap and non-blocking): |
| ('audio', np.int16 [n]) PCM 24 kHz mono, a few frames at a time |
| ('end', dict) reason, frames, audio_s, ttfa_s, wall_s, pieces, stops |
| ('error', dict) code, error |
| """ |
| from __future__ import annotations |
|
|
| import collections |
| import itertools |
| import threading |
| import time |
|
|
| import numpy as np |
|
|
| SAMPLING_KEYS = {"seed", "temperature", "top_k", "top_p", "repetition_penalty"} |
|
|
|
|
| def _check_sampling(sp): |
| """Same ranges as upstream maya3.engine._check_sampling.""" |
| if not sp: |
| return {} |
| bad = set(sp) - SAMPLING_KEYS |
| if bad: |
| raise ValueError(f"unknown sampling keys {sorted(bad)}; allowed {sorted(SAMPLING_KEYS)}") |
| out = {k: v for k, v in sp.items() if v is not None} |
| if "temperature" in out and not 0.05 <= float(out["temperature"]) <= 2.0: |
| raise ValueError("temperature must be in [0.05, 2.0]") |
| if "top_k" in out and not 1 <= int(out["top_k"]) <= 2049: |
| raise ValueError("top_k must be in [1, 2049]") |
| if "top_p" in out and not 0.0 < float(out["top_p"]) <= 1.0: |
| raise ValueError("top_p must be in (0, 1]") |
| if "repetition_penalty" in out and not 0.0 <= float(out["repetition_penalty"]) <= 5.0: |
| raise ValueError("repetition_penalty must be in [0, 5]") |
| if "seed" in out: |
| out["seed"] = int(out["seed"]) |
| return out |
|
|
|
|
| class Request: |
| _ids = itertools.count() |
|
|
| def __init__(self, text, lang, voice, on_event, max_frames=None, rid=None, sampling=None): |
| self.id = rid or f"r{next(self._ids)}" |
| self.text, self.lang, self.voice = text, lang, voice |
| self.on_event = on_event |
| self.max_frames = max_frames |
| self.sampling = sampling or {} |
| self.cancelled = False |
| self.t_submit = time.perf_counter() |
|
|
| def cancel(self): |
| self.cancelled = True |
|
|
|
|
| class ServingEngine: |
| def __init__(self, tts, log=None): |
| self.tts = tts |
| self.registry = tts.registry |
| self.failed = None |
| self._q = collections.deque() |
| self._cv = threading.Condition() |
| self._stop = False |
| self._busy = False |
| self.stats = collections.Counter() |
| self._thread = threading.Thread(target=self._loop, name="maya-mac-engine", daemon=True) |
| self._thread.start() |
|
|
| |
| def submit(self, text, lang, voice="plain", on_event=None, max_frames=None, rid=None, split=True, sampling=None): |
| if not isinstance(text, str) or not text.strip(): |
| raise ValueError("text must be a non-empty string") |
| sampling = _check_sampling(sampling) |
| vid = voice if isinstance(voice, str) else voice.id |
| self.registry.get(vid) |
| req = Request(text.strip(), lang, vid, on_event or (lambda k, p: None), max_frames, rid, sampling) |
| with self._cv: |
| if self.failed: |
| raise RuntimeError(f"engine failed: {self.failed}") |
| if self._stop: |
| raise RuntimeError("engine stopped") |
| self._q.append(req) |
| self._cv.notify() |
| return req |
|
|
| def call(self, fn, timeout=300): |
| """Run fn() on the worker thread (between requests) and return its result.""" |
| if threading.current_thread() is self._thread: |
| return fn() |
| box, ev = {}, threading.Event() |
| with self._cv: |
| self._q.append((fn, box, ev)) |
| self._cv.notify() |
| if not ev.wait(timeout): |
| raise TimeoutError("engine call timed out") |
| if "error" in box: |
| raise box["error"] |
| return box["result"] |
|
|
| def health(self): |
| with self._cv: |
| pending = sum(isinstance(x, Request) for x in self._q) |
| out = {"ready": self.failed is None and not self._stop, "failed": self.failed, "device": self.tts.device, |
| "active": int(self._busy), "pending": pending, "max_seats": 1, **self.stats} |
| try: |
| if self.tts.device == "mps": |
| import torch |
| out["mem_alloc_gb"] = round(torch.mps.current_allocated_memory() / 1e9, 2) |
| elif self.tts.device == "mlx": |
| import mlx.core as mx |
| out["mem_alloc_gb"] = round(mx.get_active_memory() / 1e9, 2) |
| except Exception: |
| pass |
| return out |
|
|
| def close(self): |
| with self._cv: |
| self._stop = True |
| self._cv.notify() |
| self._thread.join(timeout=30) |
|
|
| |
| def _loop(self): |
| try: |
| while True: |
| with self._cv: |
| while not self._q and not self._stop: |
| self._cv.wait() |
| if self._stop: |
| break |
| item = self._q.popleft() |
| self._busy = isinstance(item, Request) |
| if isinstance(item, Request): |
| self._run(item) |
| else: |
| fn, box, ev = item |
| try: |
| box["result"] = fn() |
| except Exception as exc: |
| box["error"] = exc |
| ev.set() |
| self._busy = False |
| except BaseException as exc: |
| import traceback |
| traceback.print_exc() |
| self.failed = f"{type(exc).__name__}: {exc}" |
| finally: |
| with self._cv: |
| left = list(self._q) |
| self._q.clear() |
| for item in left: |
| if isinstance(item, Request): |
| self._emit(item, "error", {"code": "unavailable", "error": self.failed or "engine stopped"}) |
| else: |
| item[1]["error"] = RuntimeError("engine stopped") |
| item[2].set() |
|
|
| def _emit(self, req, kind, payload): |
| try: |
| req.on_event(kind, payload) |
| except Exception: |
| req.cancelled = True |
|
|
| def _run(self, req): |
| if req.cancelled: |
| self._emit(req, "error", {"code": "cancelled", "error": "cancelled"}) |
| return |
| self.stats["admitted"] += 1 |
| info, t_first = {}, None |
| gen = self.tts.stream(req.text, lang=req.lang, voice=req.voice, max_frames=req.max_frames, info=info, |
| **req.sampling) |
| try: |
| for chunk in gen: |
| if req.cancelled: |
| break |
| if t_first is None: |
| t_first = time.perf_counter() |
| self._emit(req, "audio", (np.clip(chunk, -1, 1) * 32767).round().astype(np.int16)) |
| except (ValueError, KeyError) as exc: |
| self._emit(req, "error", {"code": "invalid_request", "error": f"{type(exc).__name__}: {exc}"}) |
| return |
| except Exception as exc: |
| import traceback |
| traceback.print_exc() |
| self.stats["errors"] += 1 |
| self._emit(req, "error", {"code": "internal", "error": f"{type(exc).__name__}: {exc}"}) |
| return |
| finally: |
| gen.close() |
| if req.cancelled: |
| self._emit(req, "error", {"code": "cancelled", "error": "cancelled"}) |
| return |
| self.stats["completed"] += 1 |
| self._emit(req, "end", {"reason": info.get("reason"), "frames": info.get("frames", 0), |
| "audio_s": info.get("audio_s", 0.0), |
| "ttfa_s": (t_first - req.t_submit) if t_first else None, |
| "wall_s": time.perf_counter() - req.t_submit, "pieces": info.get("pieces"), |
| "stops": info.get("stops", [])}) |
|
|
|
|
| class ServingTTS: |
| """What maya_mac.server needs from a model object: .engine, register_voice, remove_voice, close.""" |
|
|
| def __init__(self, tts): |
| self.tts = tts |
| self.engine = ServingEngine(tts) |
|
|
| def register_voice(self, voice_id, ref_audio=None, ref_text=None, lang=None, mode="icl_spk", meta=None): |
| return self.engine.call(lambda: self.tts.register_voice(voice_id, ref_audio, ref_text, lang, mode, meta)) |
|
|
| def remove_voice(self, voice_id): |
| return self.tts.remove_voice(voice_id) |
|
|
| def close(self): |
| self.engine.close() |
| self.tts.close() |
|
|
