"""PolyTTS FastAPI server — multi-engine, dual MLX / manager backend. Voices are uploaded via POST /voices (content-hashed, deduplicated) and referenced by voice_id in the POST /tts endpoint. MLX path: a single persistent model with pre-computed voice embeddings and an incremental PCM streaming endpoint backed by a disk L2 cache. Manager path (non-MLX): a multi-engine ModelManager (qwen + voxcpm) that keeps one model in VRAM at a time and evicts it after idle so a co-resident workload can reclaim the GPU. """ import os import gc import io import json import hashlib import asyncio import concurrent.futures import threading from collections import OrderedDict from contextlib import contextmanager from pathlib import Path import numpy as np import soundfile as sf import uvicorn from fastapi import FastAPI, File, Form, HTTPException, UploadFile from fastapi.responses import Response, StreamingResponse from pydantic import BaseModel import cache import pcm_cache from engines import QwenEngine, VoxcpmEngine, CosyvoiceEngine, DotsEngine, pcm16 RUNTIME = os.environ.get("POLYTTS_RUNTIME", "mlx").lower() # Non-MLX runtimes use the multi-engine manager (one model in VRAM at a time, # evicted on engine-switch and after idle so the shared GPU is freed). MLX uses # a separate native-MLX inference path (NOT engines.py, which is the # PyTorch/CUDA backend) — so the manager path is PyTorch-only. Making MLX a # livestack node requires wrapping the native-MLX load/unload in ManagedUnits, # not flipping this flag. See docs/ — deferred. _MANAGER_PATH = RUNTIME != "mlx" IDLE_EVICT_SECONDS = int(os.environ.get("POLYTTS_IDLE_EVICT_SECONDS", "120")) # Whether two engines may hold the card at once. Off by default: the smallest # host this runs on cannot afford it. See the residency note below. COLOAD = os.environ.get("POLYTTS_COLOAD", "").strip().lower() in ("1", "true", "yes") DEFAULT_ENGINE = os.environ.get("POLYTTS_DEFAULT_ENGINE", "qwen") # Engines whose non-x-vector cloning path consumes the reference TRANSCRIPT, # so a blank ref_text cannot serve it. cosyvoice is absent deliberately: its # sidecar passes ref_text straight through and tolerates an empty one. _NEEDS_REF_TEXT = frozenset({"qwen", "dots"}) # Cap in-memory voice caches so that registering thousands of unique voices # cannot grow memory without bound. Evicted voices remain on disk and are # reloaded on next use. Default 100 is generous for normal usage. _MAX_VOICES_IN_MEMORY = int(os.environ.get("POLYTTS_MAX_VOICES_MEM", "100")) VOICES_DIR = Path(os.environ.get( "POLYTTS_VOICES_DIR", str(Path(__file__).resolve().parent / "voices"), )) MODELS_DIR = Path(__file__).resolve().parent / "models" PORT = int(os.environ.get("POLYTTS_PORT", "8100")) TTS_TIMEOUT = int(os.environ.get("POLYTTS_TIMEOUT_SECONDS", "600")) # Audio seconds per streamed chunk. Lower = faster first-audio, more overhead. # 0.5 s gives sub-second time-to-first-chunk while keeping per-chunk work small. STREAM_INTERVAL = float(os.environ.get("POLYTTS_STREAM_INTERVAL", "0.5")) # MLX voxcpm engine (optional second model on the MLX runtime). When voices # registered under engine="voxcpm" are present, the MLX path lazily loads this # model and serves them alongside qwen. Voice ids are hashed exactly like the # manager path, so a voxcpm voice has the SAME voice_id on the mac (MLX) and the # CUDA nodes — the caller can hit either without re-registering. MLX_VOXCPM_MODEL_ID = os.environ.get("POLYTTS_MLX_VOXCPM_MODEL", "mlx-community/VoxCPM2-8bit") MLX_VOXCPM_CFG = float(os.environ.get("POLYTTS_MLX_VOXCPM_CFG", "2.0")) MLX_VOXCPM_STEPS = int(os.environ.get("POLYTTS_MLX_VOXCPM_STEPS", "10")) # Per-engine cap on chars (and sentences) per SINGLE model generation. VoxCPM # quality degrades on long continuous generation (~>3 sentences / ~80 CJK # chars): long input is split into multiple generations and the audio # concatenated into one seamless response. Engines not listed here have NO # server-side cap (qwen etc. are robust to long text). A client may override # per-request via max_chars_per_gen / max_sentences_per_gen (0 = unlimited), # at its own risk. MAX_CHARS_PER_GEN = { "voxcpm": int(os.environ.get("POLYTTS_VOXCPM_MAX_CHARS", "80")), } MAX_SENTENCES_PER_GEN = { "voxcpm": int(os.environ.get("POLYTTS_VOXCPM_MAX_SENTENCES", "3")), } app = FastAPI(title="PolyTTS Server") # Filled at startup model = None model_name = None voice_meta: OrderedDict[str, dict] = OrderedDict() # MLX: voice_id -> {ref_text} # MLX voice prompt cache: voice_id -> {speaker_embed, ref_codes, ref_text, audio} # Pre-computed at registration time so that generation avoids redundant # speaker-encoder and speech-tokenizer work on every cache-miss TTS call. # The "audio" field is a minimal stub — the full waveform is not retained # because the monkey-patched methods ignore their input. _mlx_prompt_cache: OrderedDict[str, dict] = OrderedDict() # Manager path (non-MLX): the one-model-in-VRAM manager + a GPU-free registry of # voices (voice_id -> {engine, dir, meta}) scanned from disk at startup. manager = None # polycore.ModelManager, set by attach() below residence = None # LivestackCoordinator (None in standalone / mlx mode) voice_registry: "OrderedDict[str, dict]" = OrderedDict() # MLX runtime, voxcpm engine: a lazily-loaded second model + a GPU-free registry # of voxcpm voices (voice_id -> {dir, ref_text, meta}). The model loads on the # first voxcpm request (or at startup if such voices exist) and stays resident — # the MLX path never evicts, so there is no per-request cold start. _voxcpm_mlx = None _mlx_voxcpm_voices: "OrderedDict[str, dict]" = OrderedDict() # Single-thread executor — keeps all GPU work on ONE thread to respect # Metal thread affinity (MLX) and MPS requirements (PyTorch). _gpu_executor = concurrent.futures.ThreadPoolExecutor( max_workers=1, thread_name_prefix="gpu", ) def _inventory() -> dict: """What this server HAS, for the fleet broker: the voice ids it can serve. Published so a caller can ask Harmony for "a polytts in North America that has voice X" in one request. Without it, placement knows only that both North American nodes serve `polytts`, sends the synthesis to the nearer one, and gets `404 Unknown voice_id` because the voice was cloned on the other — a failure the placement layer had no way to avoid, having been asked the wrong question. Computed per call rather than snapshotted: cloning a voice adds one while the process runs, and a stale inventory is how work is sent to a node that no longer matches. """ ids = list(voice_registry) or list(_mlx_voxcpm_voices) or list(voice_meta) # Only what this server can answer for. The engines it can run are already # in `units`, and advertising an engine list derived from the voices would # be wrong on the MLX path, where the registry the loop reads is a # different one. return {"voice": ids} if ids else {} # --- polycore + livestack residency (same wiring as polyasr) ---------------- # Build polycore ManagedUnits — the default engine is HARD_PIN (benchday needs TTS # hot at all times) — then attach() builds the manager+coordinator and mounts # /livestack. Exclusive by default (coload=False): one engine in VRAM at a time, # the other evicted on switch. Without livestack, polycore's LocalCoordinator # reproduces the standalone one-model-in-VRAM + idle-evict behaviour. No hard # dependency. # # `POLYTTS_COLOAD=1` lets them share the card, because "one engine at a time" is # a claim about a MACHINE's VRAM and not about this code — the same reasoning # `POLYTTS_DEFAULT_RESIDENCY` already follows. It matters the moment one feed # speaks two languages: Chinese runs on VoxCPM and English on Qwen, so exclusive # residency evicts and reloads an engine on every language switch. Measured on # xc-tower-ubuntu 2026-09-19, the same English sentence took 30-45 s through the # resident Qwen engine and 70-82 s through a VoxCPM voice, the difference being # the swap back. On that host VoxCPM (~5.5 GB) and Qwen 1.7B (~4 GB) sit inside # 12.6 GB of free card. HOST_ID = os.environ.get("POLYTTS_HOST_ID", os.environ.get("HOST_ID", "zz-tower0")) if _MANAGER_PATH: from livestack_node import ManagedUnit, ResidencyPolicy, free_cuda # A PyTorch runtime with no accelerator is not a slow TTS server; it is a # broken one that nothing downstream can detect. # # Measured: xc-tower-ubuntu ran this service for a week reporting # `"device": "cpu"` on a box with two RTX 3090s. torch decides # `cuda.is_available()` once, at import, so a process that starts while the # driver is reloading keeps that answer until it is restarted — and every # engine it loads afterwards goes to system memory. Nothing failed. The # residency manager warmed models, leases were granted, syntheses returned, # and each one took an order of magnitude longer than it should have. The # only symptom was a field in /health that nobody reads. # # Harmony loads and unloads MODELS; it cannot repair a PROCESS that has # decided there is no GPU. Refusing to start is what makes that visible, at # the one moment somebody is watching — a failed unit — instead of at the # worst one, which is never. # # `POLYTTS_ALLOW_CPU=1` is the escape hatch, for a box that genuinely has # no accelerator and wants a slow server on purpose. import torch as _torch_probe _has_cuda = _torch_probe.cuda.is_available() _mps = getattr(_torch_probe.backends, "mps", None) _has_mps = bool(_mps and _mps.is_available()) if not _has_cuda and not _has_mps and os.environ.get("POLYTTS_ALLOW_CPU") != "1": raise SystemExit( f"polytts: POLYTTS_RUNTIME={RUNTIME} needs CUDA or MPS and torch sees neither " f"(torch {_torch_probe.__version__}, cuda build {_torch_probe.version.cuda}, " f"CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '(unset)')}). " "A CPU synthesis is ~10x slower and reports success, so this refuses to start " "rather than serve it. If the GPU is there, the driver was probably not ready " "when this process began: restart the unit. If this box has no accelerator, set " "POLYTTS_ALLOW_CPU=1." ) # VRAM footprints (bytes; estimates from nvidia-smi, refine with measure_footprint). # `dots` is MEASURED, not estimated, and the number is driven by the CLONE # MODE, not by the engine: on an RTX 3090 the same checkpoint reserved # 6.3 GB serving an x-vector voice and 7.6 GB serving continuation against a # 60 s reference clip (687-char ref_text) — the prompt audio is encoded in # full on that path, so the longest registered reference sets the peak. 8 GB # is the 7.6 GB worst case rounded UP. Do not tune this down to the pretty # x-vector number: the planner sizes eviction from it, an earlier 6.5 GB # guess OOM'd mid-synthesis at 7.1 GB, and under-reporting is the same defect # that disqualified the heavier dots.tts-soar checkpoint beside polyasr. # See openspec/changes/add-dots-tts-engine. _FOOTPRINTS = {"qwen": 8_700_000_000, "voxcpm": 5_000_000_000, "cosyvoice": 3_000_000_000, "dots": 8_000_000_000} _DEFAULT_RESIDENCY = os.environ.get("POLYTTS_DEFAULT_RESIDENCY", "SOFT_PIN").strip().upper() def _engine_unit(name, engine, pin): def loader(): engine.load() return engine # ensure() returns the Engine, as callers expect def freer(): engine.unload() free_cuda() # Default engine is SOFT_PIN: preferred-warm but PREEMPTIBLE under GPU # pressure (e.g. an ASR/align burst) — the residence planner restores it # when pressure settles. The other engine is pure demand (UNPINNED). # # POLYTTS_DEFAULT_RESIDENCY overrides the default engine's policy per # host, because "keep TTS hot" is a claim about a machine, not about # this code. On a box whose GPUs are also serving LLMs, a pinned 5 GB # engine can be the reason a model that IS being asked for has nowhere # to go — SOFT_PIN only yields to something strictly more important, so # an idle engine outranks live demand. Unset, behaviour is unchanged. return ManagedUnit(name, loader, freer, footprint=_FOOTPRINTS.get(name, 0), residency_policy=(getattr(ResidencyPolicy, _DEFAULT_RESIDENCY) if pin else ResidencyPolicy.UNPINNED)) _ENGINES = {"qwen": QwenEngine(MODELS_DIR), "voxcpm": VoxcpmEngine(), "cosyvoice": CosyvoiceEngine(), "dots": DotsEngine()} _UNITS = {n: _engine_unit(n, e, n == DEFAULT_ENGINE) for n, e in _ENGINES.items()} def _gpu_call(fn): """Run a thunk on the single GPU executor (warm/evict from /livestack).""" return _gpu_executor.submit(fn).result() try: from livestack_node import attach # `port` is what lets this node report for duty to the host broker on # its own — no LIVESTACK_PEERS entry, no broker restart. It is the one # fact the broker cannot infer (a POST shows it our source address, not # what we listen on), and we already have it right here. # `in_flight` is this server's OWN count of concurrent syntheses — the # admission-control counter `_inflight` defined below, read lazily. A # synthesis does not take a livestack lease, so the facade's lease-derived # fallback would report 0 while the GPU thread is saturated; supplying it # here also gets the report labelled `in_flight_source: "server"`. manager, residence = attach(app, host_id=HOST_ID, kind="polytts", units=_UNITS, idle_seconds=IDLE_EVICT_SECONDS, coload=COLOAD, gpu_call=_gpu_call, port=PORT, in_flight=lambda: _inflight, inventory=_inventory, # After the bind, never before: see startup(). preload=lambda: _warm_manager_path()) except ImportError: from livestack_node import ModelManager manager = ModelManager(_UNITS, IDLE_EVICT_SECONDS, coload=COLOAD) # MLX runtime: the native-MLX models (qwen `model`, voxcpm `_voxcpm_mlx`) are held # as globals and served by a separate native path (NOT engines.py, the PyTorch # backend). Wrap their load/free in ManagedUnits so polytts becomes a livestack # node: a host-broker can SEE its Metal footprint and EVICT the heavy voxcpm # engine when idle to relieve Metal pressure (e.g. an ASR align-chunk spike). # qwen is HARD_PIN (its `model` global is read pervasively in the serving path, so # it is never evicted); voxcpm is SOFT_PIN + reached only through the # _load_voxcpm_mlx() accessor above, so it reloads safely on next use. elif RUNTIME == "mlx": from livestack_node import ManagedUnit, ResidencyPolicy def _qwen_unit_load(): _load_mlx() return model def _qwen_unit_free(): global model model = None import mlx.core as mx mx.clear_cache() gc.collect() # Deferred wrappers: the voxcpm raw load/free are defined later in the file. def _voxcpm_unit_load(): return _load_voxcpm_mlx_raw() def _voxcpm_unit_free(): _free_voxcpm_mlx() _MLX_UNITS = { "qwen": ManagedUnit("qwen", _qwen_unit_load, _qwen_unit_free, footprint=3_000_000_000, residency_policy=ResidencyPolicy.HARD_PIN), "voxcpm": ManagedUnit("voxcpm", _voxcpm_unit_load, _voxcpm_unit_free, footprint=8_000_000_000, residency_policy=ResidencyPolicy.SOFT_PIN), } def _gpu_call(fn): """Run a thunk on the single MLX GPU executor (facade warm/evict).""" return _gpu_executor.submit(fn).result() try: from livestack_node import attach manager, residence = attach(app, host_id=HOST_ID, kind="polytts", units=_MLX_UNITS, idle_seconds=IDLE_EVICT_SECONDS, coload=True, gpu_call=_gpu_call, port=PORT, in_flight=lambda: _inflight, inventory=_inventory, # After the bind, never before: see startup(). # On the GPU executor, like every other MLX call: MLX # streams are per-thread, so weights loaded on the # attach thread abort the first synthesis on the # executor ("There is no Stream(cpu, 1) in current # thread" -> SIGABRT, launchd restart loop). preload=lambda: _gpu_call(_warm_mlx)) except ImportError: manager = None residence = None # --------------------------------------------------------------------------- # Admission control — bound concurrent in-flight streaming syntheses so a burst # cannot pile up unbounded on the single GPU thread. Under cap: no behavior # change; over cap: honest 429 + Retry-After. (Per-account fairness needs the # account threaded from the hub capability; this is the global guard.) # # On the manager path the high-level generate_voice_clone returns whole clips, # so /tts/stream streams at sentence granularity: synthesize each sentence and # emit its PCM as soon as it is ready -> first-audio latency is one sentence, # not the whole digest. # --------------------------------------------------------------------------- _MAX_INFLIGHT = max(1, int(os.environ.get("POLYTTS_MAX_INFLIGHT", "3"))) _RETRY_AFTER_SECONDS = max(1, int(os.environ.get("POLYTTS_RETRY_AFTER", "2"))) _inflight = 0 # single-threaded event loop -> a plain counter is race-free _SENTENCE_ENDERS = ".!?\n。！？" def _try_admit() -> bool: """Reserve an in-flight slot if under cap; caller must _release() when done.""" global _inflight if _inflight >= _MAX_INFLIGHT: return False _inflight += 1 return True def _release() -> None: global _inflight if _inflight > 0: _inflight -= 1 def _split_sentences(text): """Split into sentence chunks (ASCII + CJK enders + newline). No regex so there are no unicode-escape pitfalls.""" text = (text or "").strip() if not text: return [] out, buf = [], [] for ch in text: buf.append(ch) if ch in _SENTENCE_ENDERS: seg = "".join(buf).strip() if seg: out.append(seg) buf = [] seg = "".join(buf).strip() if seg: out.append(seg) return out or [text] # Clause-level punctuation used to hard-split a single over-long sentence. _CLAUSE_ENDERS = ",;，；、、" def _hard_split_clause(sentence, max_chars): """Split one over-long sentence into pieces each <= max_chars, breaking on clause punctuation first, then hard on character count as a last resort.""" if not max_chars or len(sentence) <= max_chars: return [sentence] pieces, buf = [], "" for ch in sentence: buf += ch if ch in _CLAUSE_ENDERS and len(buf) >= max_chars: pieces.append(buf) buf = "" if buf: pieces.append(buf) out = [] for p in pieces: if len(p) <= max_chars: out.append(p) else: for i in range(0, len(p), max_chars): out.append(p[i:i + max_chars]) return out or [sentence] def _chunk_text(text, max_chars, max_sentences): """Greedily pack sentences into generation chunks, each <= max_chars chars and <= max_sentences sentences, so a single model generation never runs long enough for quality to degrade. max_chars=0 / max_sentences=0 disables that limit; both 0 -> one chunk (no chunking). Over-long single sentences are hard-split on clause boundaries.""" text = (text or "").strip() if not text: return [] no_char_cap = not max_chars no_sent_cap = not max_sentences if no_char_cap and no_sent_cap: return [text] sentences = _split_sentences(text) if (no_char_cap or len(text) <= max_chars) and (no_sent_cap or len(sentences) <= max_sentences): return [text] chunks, cur, cur_sents = [], "", 0 for s in sentences: if max_chars and len(s) > max_chars: if cur: chunks.append(cur) cur, cur_sents = "", 0 for piece in _hard_split_clause(s, max_chars): chunks.append(piece) continue candidate = (cur + s) if cur else s overflow_chars = max_chars and len(candidate) > max_chars overflow_sents = max_sentences and cur_sents + 1 > max_sentences if cur and (overflow_chars or overflow_sents): chunks.append(cur) cur, cur_sents = s, 1 else: cur, cur_sents = candidate, cur_sents + 1 if cur: chunks.append(cur) return chunks or [text] _LANG_NAMES = { "en": "English", "zh": "Chinese", "ja": "Japanese", "ko": "Korean", "es": "Spanish", "fr": "French", "de": "German", "it": "Italian", "pt": "Portuguese", "ru": "Russian", } def _lang_name(code): """Base ISO code (what the phone sends) -> the model's expected language name. Pass through an already-named value; default English for terminal content (the model's own default is Chinese, wrong for code).""" if not code: return "English" c = str(code).strip() return _LANG_NAMES.get(c.lower(), c) # --------------------------------------------------------------------------- # MLX backend # --------------------------------------------------------------------------- # Default: 1.7B-8bit for best voice quality (~2.9 GB: 2.25 GB backbone + 651 MB codec). # Override with POLYTTS_MLX_MODEL for lower memory at reduced quality: # 0.6B-4bit: ~1.63 GB (mlx-community/Qwen3-TTS-12Hz-0.6B-Base-4bit) # 0.6B-6bit: ~1.65 GB (mlx-community/Qwen3-TTS-12Hz-0.6B-Base-6bit) MLX_MODEL_ID = os.environ.get( "POLYTTS_MLX_MODEL", "mlx-community/Qwen3-TTS-12Hz-1.7B-Base-8bit", ) # How much freed Metal memory the MLX allocator may hoard for reuse. # Lower = tighter steady-state footprint; higher = fewer re-allocations. # 0 disables the cache entirely. Default: 256 MB. _MLX_CACHE_LIMIT = int(os.environ.get("POLYTTS_MLX_CACHE_MB", "256")) * 1024 * 1024 def _trim_voice_caches() -> None: """Evict oldest voices from in-memory caches when over capacity. Oldest = least-recently registered/loaded. Evicted voices remain on disk and are reloaded automatically on next use. """ while len(voice_meta) > _MAX_VOICES_IN_MEMORY: oldest, _ = voice_meta.popitem(last=False) _mlx_prompt_cache.pop(oldest, None) def _load_mlx(): global model, model_name import mlx.core as mx from mlx_audio.tts.utils import load_model model_name = MLX_MODEL_ID.split("/")[-1] print(f"Loading MLX model {model_name} …") model = load_model(MLX_MODEL_ID) # Cap the Metal allocator's free-buffer cache so freed memory returns # to the OS instead of being held indefinitely for potential reuse. mx.set_cache_limit(_MLX_CACHE_LIMIT) active_mb = mx.get_active_memory() / 1024 / 1024 print(f"MLX model loaded. Active Metal memory: {active_mb:.0f} MB " f"(cache limit: {_MLX_CACHE_LIMIT // 1024 // 1024} MB)") mx.reset_peak_memory() def _precompute_voice_mlx(wav_path: str, ref_text: str) -> dict: """Extract and return voice-specific artifacts that are reused across every TTS call for the same voice. Returns dict with: speaker_embed: mx.array — x-vector from the speaker encoder ref_codes: mx.array or None — speech tokenizer codes (ICL mode) ref_text: str audio: mx.array — minimal stub (monkey-patched methods ignore it) """ import mlx.core as mx from mlx_audio.utils import load_audio audio = load_audio(wav_path, sample_rate=model.sample_rate) # Speaker embedding (x-vector) speaker_embed = None if getattr(model, "speaker_encoder", None) is not None: speaker_embed = model.extract_speaker_embedding(audio) mx.eval(speaker_embed) # Reference codec codes (for ICL voice cloning) ref_codes = None st = getattr(model, "speech_tokenizer", None) if st is not None and getattr(st, "has_encoder", False): audio_enc = audio[None, None, :] if audio.ndim == 1 else audio[None, :] ref_codes = st.encode(audio_enc) mx.eval(ref_codes) # Replace full waveform with a tiny stub. The generate() pipeline # requires ref_audio to be a non-None mx.array with ndim >= 1, but # both extract_speaker_embedding and speech_tokenizer.encode are # monkey-patched to return cached values — so the stub is never # actually processed. This avoids retaining ~0.5-1 MB per voice. audio_stub = mx.zeros((model.sample_rate,)) # Release the full waveform from Metal memory del audio mx.clear_cache() return { "speaker_embed": speaker_embed, "ref_codes": ref_codes, "ref_text": ref_text, "audio": audio_stub, } def _prompt_cache_dir(voice_dir: Path) -> Path: """Return the model-specific prompt cache directory for a voice. Layout: voices//prompts// Each model gets its own namespace so switching models (e.g. 0.6B <-> 1.7B) never requires recomputation — both caches coexist on disk. """ model_key = hashlib.sha256(MLX_MODEL_ID.encode()).hexdigest()[:12] return voice_dir / "prompts" / model_key def _save_voice_cache_mlx(voice_dir: Path, prompt: dict): """Persist pre-computed voice embeddings as .npy so they survive restarts.""" cache_dir = _prompt_cache_dir(voice_dir) cache_dir.mkdir(parents=True, exist_ok=True) if prompt["speaker_embed"] is not None: np.save(str(cache_dir / "speaker_embed.npy"), np.array(prompt["speaker_embed"])) if prompt["ref_codes"] is not None: np.save(str(cache_dir / "ref_codes.npy"), np.array(prompt["ref_codes"])) def _load_voice_cache_mlx(voice_dir: Path, ref_text: str) -> dict | None: """Load persisted voice embeddings from disk for the current model. Returns None if no cache exists for this model — a different model's cache is simply ignored, not deleted.""" import mlx.core as mx cache_dir = _prompt_cache_dir(voice_dir) embed_path = cache_dir / "speaker_embed.npy" if not embed_path.exists(): return None codes_path = cache_dir / "ref_codes.npy" speaker_embed = mx.array(np.load(str(embed_path))) ref_codes = mx.array(np.load(str(codes_path))) if codes_path.exists() else None # Minimal stub — see _precompute_voice_mlx for rationale audio_stub = mx.zeros((model.sample_rate,)) return { "speaker_embed": speaker_embed, "ref_codes": ref_codes, "ref_text": ref_text, "audio": audio_stub, } def _load_voices_mlx(): """Load previously-uploaded voices from disk (MLX backend). Tries to load pre-computed embeddings (.npy) first; falls back to full recomputation from the WAV file if cached embeddings are missing. """ if not VOICES_DIR.exists(): return for voice_dir in sorted(VOICES_DIR.iterdir()): if not voice_dir.is_dir(): continue meta_path = voice_dir / "meta.json" wav_path = voice_dir / "voice.wav" if not meta_path.exists() or not wav_path.exists(): continue voice_id = voice_dir.name try: meta = json.loads(meta_path.read_text()) ref_text = meta["ref_text"] # voxcpm voices are served by the MLX voxcpm model, not the qwen # speaker-embedding cache — register and skip the qwen precompute. if meta.get("engine") == "voxcpm": _mlx_voxcpm_voices[voice_id] = { "dir": voice_dir, "ref_text": ref_text, "meta": meta} print(f" registered voxcpm voice: {voice_id}") continue # Try cached embeddings first (fast path — no model inference) prompt = _load_voice_cache_mlx(voice_dir, ref_text) if prompt is None: # Cold start: compute from WAV and persist for next time print(f" computing voice prompt: {voice_id}") prompt = _precompute_voice_mlx(str(wav_path), ref_text) _save_voice_cache_mlx(voice_dir, prompt) else: print(f" loaded cached voice: {voice_id}") voice_meta[voice_id] = {"ref_text": ref_text} _mlx_prompt_cache[voice_id] = prompt except Exception as e: print(f" WARNING: skipping voice {voice_id}: {e}") _trim_voice_caches() print(f"Loaded {len(voice_meta)} voices") def _register_voice_mlx(voice_id: str, wav_path: str, ref_text: str): """Register a new voice — pre-compute and cache embeddings. Must be called from the GPU executor thread. """ prompt = _precompute_voice_mlx(wav_path, ref_text) _save_voice_cache_mlx(VOICES_DIR / voice_id, prompt) voice_meta[voice_id] = {"ref_text": ref_text} voice_meta.move_to_end(voice_id) _mlx_prompt_cache[voice_id] = prompt _mlx_prompt_cache.move_to_end(voice_id) _trim_voice_caches() @contextmanager def _mlx_voice_context(text: str, voice_id: str | None, extra_kwargs: dict | None = None): """Yield generate() kwargs with cached voice embeddings injected. When a voice_id has pre-computed data, we temporarily replace the model's extract_speaker_embedding and speech_tokenizer.encode with lambdas that return the cached values, restoring the originals on exit. This is safe because the GPU executor is single-threaded — only one generation runs at a time. Must be entered on the GPU executor thread. """ kwargs: dict = {"text": text} if extra_kwargs: kwargs.update(extra_kwargs) original_extract = None original_encode = None if voice_id and voice_id in _mlx_prompt_cache: cached = _mlx_prompt_cache[voice_id] kwargs["ref_audio"] = cached["audio"] kwargs["ref_text"] = cached["ref_text"] # Inject cached speaker embedding — skip mel spectrogram + encoder CNN if cached["speaker_embed"] is not None: original_extract = model.extract_speaker_embedding _embed = cached["speaker_embed"] model.extract_speaker_embedding = lambda *_a, **_kw: _embed # Inject cached ref codes — skip speech tokenizer encoding if cached["ref_codes"] is not None: st = model.speech_tokenizer original_encode = st.encode _codes = cached["ref_codes"] st.encode = lambda *_a, **_kw: _codes try: yield kwargs finally: # Always restore originals, even on error if original_extract is not None: model.extract_speaker_embedding = original_extract if original_encode is not None: model.speech_tokenizer.encode = original_encode def _generate_mlx(text: str, voice_id: str | None) -> bytes: """Generate one complete WAV, injecting cached voice embeddings.""" import mlx.core as mx with _mlx_voice_context(text, voice_id) as kwargs: chunks = [] sample_rate = None for result in model.generate(**kwargs): chunks.append(np.array(result.audio)) if sample_rate is None: sample_rate = result.sample_rate audio = np.concatenate(chunks) del chunks mx.clear_cache() gc.collect() buf = io.BytesIO() sf.write(buf, audio, sample_rate, subtype="PCM_16", format="WAV") return buf.getvalue() # Sentinel pushed onto the stream queue to signal clean end-of-generation. _STREAM_DONE = object() def _float_to_pcm16(audio_arr) -> bytes: """Convert a float waveform in [-1, 1] to little-endian signed 16-bit PCM.""" a = np.asarray(audio_arr, dtype=np.float32).reshape(-1) np.clip(a, -1.0, 1.0, out=a) return (a * 32767.0).astype(" int: """Output sample rate for the voxcpm header. VoxCPM2 is 48 kHz; falls back to that until the model is resident (the audio is 48 kHz regardless).""" return int(getattr(_voxcpm_mlx, "sample_rate", 48000)) if _voxcpm_mlx is not None else 48000 def _voxcpm_mlx_kwargs(voice: dict) -> dict: return dict( ref_audio=str(voice["dir"] / "voice.wav"), ref_text=voice["ref_text"], inference_timesteps=MLX_VOXCPM_STEPS, cfg_value=MLX_VOXCPM_CFG, ) def _generate_voxcpm_mlx(text: str, voice: dict, chunks=None) -> bytes: """One complete WAV from the MLX voxcpm model (GPU executor thread). Long text is synthesized one chunk at a time (see _chunk_text) and the audio concatenated, so a single generation never runs past the VoxCPM cap.""" import mlx.core as mx m = _load_voxcpm_mlx() for_chunks = chunks if chunks else [text] parts, sample_rate = [], None for ch in for_chunks: for result in m.generate(text=ch, **_voxcpm_mlx_kwargs(voice)): parts.append(np.array(result.audio)) if sample_rate is None: sample_rate = result.sample_rate audio = np.concatenate(parts) del parts mx.clear_cache() gc.collect() buf = io.BytesIO() sf.write(buf, audio, sample_rate, subtype="PCM_16", format="WAV") return buf.getvalue() def _generate_voxcpm_mlx_stream(text, voice, queue, loop, chunks=None): """Push s16le PCM from the MLX voxcpm model onto `queue` (GPU thread). VoxCPM2 decodes a clip per generate() call (one yield), so /tts/stream emits it at sentence granularity — same shape as the manager path, but with the model resident (no cold start). Long input is pre-split into cap-bounded chunks (one generation each) so quality does not degrade on long text.""" import mlx.core as mx def push(item): loop.call_soon_threadsafe(queue.put_nowait, item) try: m = _load_voxcpm_mlx() for_chunks = chunks if chunks else [text] for ch in for_chunks: for result in m.generate(text=ch, **_voxcpm_mlx_kwargs(voice)): pcm = _float_to_pcm16(result.audio) if pcm: push(pcm) mx.clear_cache() gc.collect() push(_STREAM_DONE) except Exception as exc: push(exc) # --------------------------------------------------------------------------- # Manager path (non-MLX): GPU-free voice registry + one-model-in-VRAM manager # --------------------------------------------------------------------------- def _scan_voice_registry(): """Populate voice_registry from disk WITHOUT any GPU work (no model load).""" voice_registry.clear() if not VOICES_DIR.exists(): return for d in sorted(VOICES_DIR.iterdir()): if not d.is_dir(): continue meta_path, wav_path = d / "meta.json", d / "voice.wav" if not meta_path.exists() or not wav_path.exists(): continue try: meta = json.loads(meta_path.read_text()) except Exception as e: print(f" skipping voice {d.name}: {e}") continue # Legacy voices (pre-engine field) are Qwen. voice_registry[d.name] = {"engine": meta.get("engine", "qwen"), "dir": d, "meta": meta} print(f"Registered {len(voice_registry)} voices (models load lazily)") def _resolve_engine(req_engine, voice_id): """(engine_name, registry_entry) for a request, or (None, None) if unknown.""" entry = voice_registry.get(voice_id) if entry is None: return None, None return (req_engine or entry["engine"]), entry # --------------------------------------------------------------------------- # FastAPI # --------------------------------------------------------------------------- def _warm_mlx() -> None: """Load what MLX serving needs. Runs AFTER the port is bound — see below.""" global manager if manager is not None: manager.ensure("qwen") else: _load_mlx() _load_voices_mlx() # Pre-load the voxcpm model when voxcpm voices exist so the first request # isn't slow. Never fail the warm if it can't load. if _mlx_voxcpm_voices: try: _load_voxcpm_mlx() except Exception as e: print(f"[voxcpm-mlx] preload failed (voxcpm voices will 503): {e}") def _warm_manager_path() -> None: """Load the pinned default engine on the GPU executor.""" _gpu_executor.submit(manager.ensure, DEFAULT_ENGINE).result() @app.on_event("startup") async def startup(): global manager # **Nothing that loads a model happens here any more, and the reason is a # deadlock rather than a preference.** # # A FastAPI startup hook runs BEFORE uvicorn binds. A load in here means the # port stays closed until it finishes, so the broker cannot snapshot this # node, so it never learns this process hosts `polytts` — and `admit` for an # unknown kind is a refusal, by design, because a refusal and an outage must # not look alike. `manager.ensure` then waits for a grant that cannot come. # # Measured on xc-mac-studio, 2026-09-18: the local broker answered # `defer polytts-708943429 (unknown kind)` and this process sat with no # listener, no open sockets and no CPU. The fleet showed `mia`, # `Connection refused`, for forty-six hours — and a restart landed straight # back in it. # # `attach(preload=...)` warms after the facade answers, which is the only # place the cycle can be broken: bound, then announced, then snapshotted, # then warmed. While it warms the node reports `ready: false`, which the # fleet view models and the ranker filters on, so an unwarmed node is not # chosen rather than chosen and slow. if RUNTIME == "mlx": if manager is None: # No livestack: nobody will warm this for us, and there is no # arbitration to deadlock against either. Own thread, so the bind # still happens now. # The thread only waits; the load itself runs on the GPU executor # (MLX streams are per-thread — see the attach preload above). threading.Thread(target=lambda: _gpu_executor.submit(_warm_mlx).result(), name="polytts-warm", daemon=True).start() return _scan_voice_registry() if manager is not None and residence is None: # Same case on the manager path: no livestack attach, no preload hook. threading.Thread(target=_warm_manager_path, name="polytts-warm", daemon=True).start() async def _idle_loop(): loop = asyncio.get_running_loop() interval = max(5, IDLE_EVICT_SECONDS // 4) while True: await asyncio.sleep(interval) try: await loop.run_in_executor(_gpu_executor, manager.maybe_evict) except Exception as e: print(f"[idle] {e}") asyncio.create_task(_idle_loop()) print(f"Manager ready (units={list(manager.units)}, default={DEFAULT_ENGINE} " f"(pinned), livestack={residence is not None}, idle_evict={IDLE_EVICT_SECONDS}s)") def _hash_audio(data: bytes) -> str: return hashlib.sha256(data).hexdigest()[:16] @app.post("/voices") async def upload_voice( audio: UploadFile = File(...), ref_text: str = Form(...), x_vector_only_mode: bool = Form(False), engine: str = Form(DEFAULT_ENGINE), seed_audio: UploadFile | None = File(None), seed_text: str | None = Form(None), ): """Register a voice. Qwen voices = reference clip (timbre). VoxCPM voices may additionally carry a tone *seed* (``seed_audio`` + ``seed_text``) so every generation continues that locked tone. ``engine``/``seed_*`` are ignored on the MLX path. """ audio_bytes = await audio.read() if not audio_bytes: raise HTTPException(400, "Empty audio file") # ----- Legacy MLX path (unchanged) ----- if not _MANAGER_PATH: # voxcpm on the MLX runtime: hash the voice id EXACTLY like the manager # path so the same reference produces the same voice_id on every node. if engine == "voxcpm": seed_bytes = await seed_audio.read() if seed_audio is not None else b"" h = audio_bytes + b"|" + engine.encode() if x_vector_only_mode: h += b"|xvec" if seed_bytes: h += b"|seed:" + hashlib.sha256(seed_bytes).digest() if seed_text: h += b"|st:" + seed_text.encode() voice_id = _hash_audio(h) if voice_id in _mlx_voxcpm_voices: return {"voice_id": voice_id} voice_dir = VOICES_DIR / voice_id voice_dir.mkdir(parents=True, exist_ok=True) (voice_dir / "voice.wav").write_bytes(audio_bytes) meta = {"engine": engine, "ref_text": ref_text, "x_vector_only_mode": x_vector_only_mode} if seed_bytes and seed_text: (voice_dir / "seed.wav").write_bytes(seed_bytes) meta["seed_text"] = seed_text (voice_dir / "meta.json").write_text(json.dumps(meta)) _mlx_voxcpm_voices[voice_id] = { "dir": voice_dir, "ref_text": ref_text, "meta": meta} print(f"Registered voxcpm voice {voice_id} (mlx)") return {"voice_id": voice_id} voice_id = _hash_audio(audio_bytes + b"|xvec") if x_vector_only_mode else _hash_audio(audio_bytes) if voice_id in _mlx_prompt_cache: return {"voice_id": voice_id} voice_dir = VOICES_DIR / voice_id voice_dir.mkdir(parents=True, exist_ok=True) (voice_dir / "voice.wav").write_bytes(audio_bytes) (voice_dir / "meta.json").write_text(json.dumps( {"ref_text": ref_text, "x_vector_only_mode": x_vector_only_mode})) loop = asyncio.get_running_loop() await loop.run_in_executor( _gpu_executor, _register_voice_mlx, voice_id, str(voice_dir / "voice.wav"), ref_text) print(f"Registered voice {voice_id}") return {"voice_id": voice_id} # ----- Manager path ----- # Validated against the live registry, not a hand-copied tuple — a list here # rots the moment an engine is added and rejects it as "unknown". if engine not in _ENGINES: raise HTTPException(400, f"Unknown engine: {engine}") # Continuation cloning needs the reference transcript. An engine that cannot # clone without one must FAIL the registration rather than fall back to its # x-vector path: a silent downgrade produces a voice that sounds wrong months # later with nothing in the logs to say why. Checked BEFORE any directory is # created, so a rejected registration leaves nothing behind. if engine in _NEEDS_REF_TEXT and not x_vector_only_mode and not ref_text.strip(): raise HTTPException( 400, f"engine '{engine}' clones by continuation when x_vector_only_mode is " f"false, and that requires a non-empty ref_text (the transcript of the " f"reference clip). Supply ref_text, or register with " f"x_vector_only_mode=true to clone from timbre alone.") seed_bytes = await seed_audio.read() if seed_audio is not None else b"" # Content-hash so identical (audio, engine, mode, seed) dedupes to one id. h = audio_bytes + b"|" + engine.encode() if x_vector_only_mode: h += b"|xvec" if seed_bytes: h += b"|seed:" + hashlib.sha256(seed_bytes).digest() if seed_text: h += b"|st:" + seed_text.encode() voice_id = _hash_audio(h) voice_dir = VOICES_DIR / voice_id voice_dir.mkdir(parents=True, exist_ok=True) (voice_dir / "voice.wav").write_bytes(audio_bytes) meta = {"engine": engine, "ref_text": ref_text, "x_vector_only_mode": x_vector_only_mode} if seed_bytes and seed_text: (voice_dir / "seed.wav").write_bytes(seed_bytes) meta["seed_text"] = seed_text (voice_dir / "meta.json").write_text(json.dumps(meta)) # No GPU work here — voice artifacts are built lazily on first /tts when the # owning engine is resident, so registration never forces a model load. voice_registry[voice_id] = {"engine": engine, "dir": voice_dir, "meta": meta} print(f"Registered {engine} voice {voice_id}") return {"voice_id": voice_id} @app.get("/voices") def list_voices(): if not _MANAGER_PATH: return {"voices": sorted(set(voice_meta.keys()) | set(_mlx_voxcpm_voices.keys()))} return {"voices": sorted(voice_registry.keys())} class TTSRequest(BaseModel): text: str voice_id: str | None = None language: str | None = None engine: str | None = None # manager path: override the voice's engine temperature: float | None = None top_k: int | None = None top_p: float | None = None repetition_penalty: float | None = None subtalker_temperature: float | None = None subtalker_top_k: int | None = None subtalker_top_p: float | None = None max_new_tokens: int | None = None non_streaming_mode: bool | None = None # Server-side generation chunking override (at the caller's own risk). # None = use the engine's default cap; 0 = unlimited (one generation). max_chars_per_gen: int | None = None max_sentences_per_gen: int | None = None # VoxCPM generation-quality overrides (engine default when None): # cfg_value = classifier-free-guidance / clone strength (lower = more natural) # inference_timesteps = denoising steps (higher = cleaner, slower) # denoise = denoise the reference/seed audio before cloning cfg_value: float | None = None inference_timesteps: int | None = None denoise: bool | None = None # CosyVoice instruct (emotion/style) — e.g. "请用温暖的语气" or "用沉稳严肃的语调". # Only the cosyvoice engine honors this; ignored by qwen/voxcpm. instruct: str | None = None def _generation_kwargs(req: TTSRequest) -> dict: fields = ( "temperature", "top_k", "top_p", "repetition_penalty", "subtalker_temperature", "subtalker_top_k", "subtalker_top_p", "max_new_tokens", "cfg_value", "inference_timesteps", "denoise", "instruct", ) return {key: getattr(req, key) for key in fields if getattr(req, key) is not None} def _gen_chunks(engine_name, req): """Split req.text into per-generation chunks bounded by the engine's server-side cap (overridable per-request). Returns [] only for empty text.""" mc = req.max_chars_per_gen if mc is None: mc = MAX_CHARS_PER_GEN.get(engine_name, 0) ms = req.max_sentences_per_gen if ms is None: ms = MAX_SENTENCES_PER_GEN.get(engine_name, 0) return _chunk_text(req.text, mc, ms) @app.post("/tts") async def tts(req: TTSRequest): # `_lang_name`, exactly as the streaming path below does it. This one took # the field raw, so the ISO code every other surface accepts — the phone # sends `en`, and so does anything that speaks BCP-47 — reached the engine # unmapped and came back a 500: `Unsupported languages: ['en']. Supported: # ['auto', 'chinese', 'english', ...]`. Measured 2026-09-18 on the first # English synthesis this fleet ever asked for. One request shape working on # one endpoint and 500ing on the other is not a language problem. # # The default stays Chinese here (the streaming path defaults to English # for terminal content); only the MAPPING is shared. language = _lang_name(req.language or "Chinese") gen_kwargs = _generation_kwargs(req) non_streaming_mode = req.non_streaming_mode if req.non_streaming_mode is not None else False loop = asyncio.get_running_loop() if not _MANAGER_PATH: # voxcpm voice → MLX voxcpm model (its own cache namespace + 48 kHz). vox = _mlx_voxcpm_voices.get(req.voice_id) if req.voice_id else None if vox is not None: vox_key = ("voxcpm", req.text, req.voice_id, language, req.max_chars_per_gen, req.max_sentences_per_gen) cached = cache.get(vox_key) if cached is not None: return Response(content=cached, media_type="audio/wav") vox_chunks = _gen_chunks("voxcpm", req) try: wav_bytes = await asyncio.wait_for( loop.run_in_executor(_gpu_executor, _generate_voxcpm_mlx, req.text, vox, vox_chunks), timeout=TTS_TIMEOUT) except asyncio.TimeoutError: raise HTTPException(504, "TTS generation timed out") cache.put(vox_key, wav_bytes) return Response(content=wav_bytes, media_type="audio/wav") cache_key = (req.text, req.voice_id, language, non_streaming_mode, json.dumps(gen_kwargs, sort_keys=True)) cached = cache.get(cache_key) if cached is not None: return Response(content=cached, media_type="audio/wav") if req.voice_id and req.voice_id not in _mlx_prompt_cache: raise HTTPException(404, f"Unknown voice_id: {req.voice_id}") try: wav_bytes = await asyncio.wait_for( loop.run_in_executor(_gpu_executor, _generate_mlx, req.text, req.voice_id), timeout=TTS_TIMEOUT) except asyncio.TimeoutError: raise HTTPException(504, "TTS generation timed out") cache.put(cache_key, wav_bytes) return Response(content=wav_bytes, media_type="audio/wav") # ----- Manager path ----- if not req.voice_id: raise HTTPException(400, "voice_id is required") engine_name, entry = _resolve_engine(req.engine, req.voice_id) if entry is None: raise HTTPException(404, f"Unknown voice_id: {req.voice_id}") cache_key = (engine_name, req.text, req.voice_id, language, non_streaming_mode, json.dumps(gen_kwargs, sort_keys=True), req.max_chars_per_gen, req.max_sentences_per_gen) cached = cache.get(cache_key) if cached is not None: return Response(content=cached, media_type="audio/wav") def job(): eng = manager.ensure(engine_name) # VoxCPM etc. degrade on long continuous generation: synthesize one # chunk at a time and concatenate, so a single model call never runs # past the engine's cap. chunks = _gen_chunks(engine_name, req) parts, sr = [], None for ch in chunks: audio, _sr = eng.generate( ch, req.voice_id, entry["dir"], entry["meta"], language, {**gen_kwargs, "non_streaming_mode": non_streaming_mode}) parts.append(np.asarray(audio, dtype=np.float32)) if sr is None: sr = _sr audio = np.concatenate(parts) if parts else np.zeros(0, dtype=np.float32) buf = io.BytesIO() sf.write(buf, audio, sr or eng.sample_rate, subtype="PCM_16", format="WAV") return buf.getvalue() try: wav_bytes = await asyncio.wait_for( loop.run_in_executor(_gpu_executor, job), timeout=TTS_TIMEOUT) except asyncio.TimeoutError: raise HTTPException(504, "TTS generation timed out") cache.put(cache_key, wav_bytes) return Response(content=wav_bytes, media_type="audio/wav") @app.post("/tts/stream") async def tts_stream(req: TTSRequest): """Stream raw PCM as the model produces it. Body: little-endian signed 16-bit mono PCM, chunked as generated. Headers advertise the format so the caller can wrap/transcode (e.g. to Opus). No WAV header is emitted — the length is unknown mid-stream. MLX path: the generation runs on the single GPU executor thread; chunks flow out through an asyncio queue. First-audio latency ≈ model time-to-first-chunk (tuned by POLYTTS_STREAM_INTERVAL), not the full clip duration. Identical (text, voice_id) requests are served from a disk L2 PCM cache, skipping the GPU entirely. Manager path: sentence-chunked across engines; within a sentence, engines with native streaming (VoxCPM) emit sub-sentence chunks for lower first-audio latency. Qwen emits one chunk per sentence. """ if not _MANAGER_PATH: vox = _mlx_voxcpm_voices.get(req.voice_id) if req.voice_id else None if req.voice_id and vox is None and req.voice_id not in _mlx_prompt_cache: raise HTTPException(404, f"Unknown voice_id: {req.voice_id}") sr = _voxcpm_mlx_sr() if vox is not None else int(getattr(model, "sample_rate", 24000)) base_headers = { "X-Sample-Rate": str(sr), "X-Audio-Format": "s16le", "X-Channels": "1", "Cache-Control": "no-store", } # L2 cache: identical (text, voice_id) ⇒ identical audio. Serve from disk and # skip the GPU entirely (no admission slot needed — no synthesis runs). cached = pcm_cache.get(req.text, req.voice_id) if cached is not None: cached_sr, cached_bytes = cached async def cached_body(): chunk = 32768 for i in range(0, len(cached_bytes), chunk): yield cached_bytes[i:i + chunk] return StreamingResponse( cached_body(), media_type="application/octet-stream", headers={**base_headers, "X-Sample-Rate": str(cached_sr), "X-Cache": "hit"}, ) if not _try_admit(): raise HTTPException( 429, "TTS server busy", headers={"Retry-After": str(_RETRY_AFTER_SECONDS)}, ) loop = asyncio.get_running_loop() queue: asyncio.Queue = asyncio.Queue() # Kick off generation on the GPU thread; do NOT await — chunks arrive via # the queue. The future is intentionally not held: if the client # disconnects, generation runs to completion and its chunks are dropped. vox_chunks = _gen_chunks("voxcpm", req) if vox is not None else None if vox is not None: loop.run_in_executor( _gpu_executor, _generate_voxcpm_mlx_stream, req.text, vox, queue, loop, vox_chunks, ) else: loop.run_in_executor( _gpu_executor, _generate_mlx_stream, req.text, req.voice_id, queue, loop, ) async def body(): chunks: list[bytes] = [] completed = False try: while True: try: item = await asyncio.wait_for(queue.get(), timeout=TTS_TIMEOUT) except asyncio.TimeoutError: print("[tts/stream] timed out waiting for audio") return if item is _STREAM_DONE: completed = True return if isinstance(item, Exception): # Mid-stream failure: end the response. Bytes already sent # are valid PCM; the caller treats a short stream as a soft # failure. print(f"[tts/stream] generation error: {item}") return chunks.append(item) yield item finally: _release() # Only cache a CLEANLY completed synthesis — never a partial stream # (client disconnect / timeout / error), which would poison the cache # with truncated audio. if completed and chunks: try: pcm_cache.put(req.text, req.voice_id, sr, b"".join(chunks)) except Exception as e: # noqa: BLE001 print(f"[pcm_cache] put failed: {e}") return StreamingResponse( body(), media_type="application/octet-stream", headers={**base_headers, "X-Cache": "miss"}, ) # ----- Manager path ----- if not req.voice_id: raise HTTPException(400, "voice_id is required") engine_name, entry = _resolve_engine(req.engine, req.voice_id) if entry is None: raise HTTPException(404, f"Unknown voice_id: {req.voice_id}") if not _try_admit(): raise HTTPException( 429, "TTS server busy", headers={"Retry-After": str(_RETRY_AFTER_SECONDS)}, ) try: language = _lang_name(req.language) gen_kwargs = _generation_kwargs(req) non_streaming_mode = req.non_streaming_mode if req.non_streaming_mode is not None else False gen_chunks = _gen_chunks(engine_name, req) loop = asyncio.get_running_loop() queue: asyncio.Queue = asyncio.Queue() # Make the engine resident first so the response sample-rate header is # correct (cold start may load the model here). eng = await loop.run_in_executor(_gpu_executor, manager.ensure, engine_name) sample_rate = eng.sample_rate def worker(): try: e = manager.ensure(engine_name) for gch in gen_chunks: for chunk in e.stream( gch, req.voice_id, entry["dir"], entry["meta"], language, {**gen_kwargs, "non_streaming_mode": non_streaming_mode}, ): pcm = pcm16(chunk) if pcm: loop.call_soon_threadsafe(queue.put_nowait, pcm) loop.call_soon_threadsafe(queue.put_nowait, _STREAM_DONE) except Exception as exc: loop.call_soon_threadsafe(queue.put_nowait, exc) loop.run_in_executor(_gpu_executor, worker) async def body(): try: while True: try: item = await asyncio.wait_for(queue.get(), timeout=TTS_TIMEOUT) except asyncio.TimeoutError: print("[tts/stream] timed out waiting for audio") return if item is _STREAM_DONE: return if isinstance(item, Exception): print(f"[tts/stream] generation error: {item}") return yield item finally: _release() headers = { "X-Sample-Rate": str(sample_rate), "X-Audio-Format": "s16le", "X-Channels": "1", "X-Engine": engine_name, "Cache-Control": "no-store", } return StreamingResponse(body(), media_type="application/octet-stream", headers=headers) except Exception: _release() raise @app.post("/model/unload") async def model_unload(): """Gracefully evict the resident model from VRAM (and return freed heap to the OS) without stopping the server. Lets a co-resident workload — e.g. the local renderer — reclaim GPU+RAM; the model reloads lazily on the next /tts. No-op on the MLX path (single persistent model).""" if not _MANAGER_PATH or manager is None: raise HTTPException(400, "unload requires the manager (non-MLX) runtime") loop = asyncio.get_running_loop() evicted = await loop.run_in_executor(_gpu_executor, manager.unload_now) return {"unloaded": evicted, "manager": manager.status()} @app.get("/health") def health(): info = { "model": model_name, "runtime": RUNTIME, "tts_timeout_seconds": TTS_TIMEOUT, "inflight": _inflight, "max_inflight": _MAX_INFLIGHT, } if not _MANAGER_PATH: import mlx.core as mx info["device"] = "apple-silicon-mlx" info["voices"] = sorted(voice_meta.keys()) info["memory_mb"] = { "active": round(mx.get_active_memory() / 1024 / 1024), "peak": round(mx.get_peak_memory() / 1024 / 1024), "cache": round(mx.get_cache_memory() / 1024 / 1024), } info["wav_cache_entries"] = len(cache._store) else: import torch if torch.cuda.is_available(): info["device"] = "cuda" info["vram_mb"] = { "allocated": round(torch.cuda.memory_allocated() / 1e6), "reserved": round(torch.cuda.memory_reserved() / 1e6), } elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): info["device"] = "mps" else: info["device"] = "cpu" info["manager"] = manager.status() if manager else None info["model"] = sorted(manager.resident) if manager else None # Which CHECKPOINT each resident engine actually loaded. The engine NAME # cannot answer "did POLYTTS_DOTS_MODEL / VOXCPM_MODEL_ID take effect?", # and a server quietly running a different checkpoint than the operator # set is otherwise invisible until somebody listens to the output. info["checkpoint"] = ({n: getattr(_ENGINES[n], "model_name", None) for n in sorted(manager.resident) if n in _ENGINES} if manager else None) info["voices"] = sorted(voice_registry.keys()) info["wav_cache_entries"] = len(cache._store) return info def _refuse_if_port_taken() -> None: """A port already in use is a failure, and uvicorn reports it as success. Measured on xc-mac-studio: two instances raced, the loser logged `[Errno 48] address already in use`, shut down gracefully and **exited 0**. `run.sh` reads exit 0 as "server exited cleanly", breaks its restart loop, and leaves the launchd job holding a wrapper with no server inside it. launchd saw a healthy job; the fleet saw the node `mia` with `Connection refused` — for forty-six hours. Checked before uvicorn starts, so the exit code says what happened. There is a race between this probe and the bind; it does not matter, because the bind failure after it is the same non-zero exit, which is all `run.sh` needed in the first place. """ import socket probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM) # Deliberately NOT SO_REUSEADDR: the question is "is somebody serving # here", and the option exists to make that question answer no. try: probe.bind(("0.0.0.0", PORT)) except OSError as exc: import errno as _errno what = ("another process is already serving it" if exc.errno == _errno.EADDRINUSE else "this process cannot bind it") raise SystemExit( f"polytts: port {PORT} is unavailable — {what} ({exc}). " f"Check `lsof -nP -iTCP:{PORT} -sTCP:LISTEN`. Refusing to start: a second " "instance exits 0 after its failed bind, and every supervisor above reads " "that as a clean shutdown." ) finally: probe.close() if __name__ == "__main__": _refuse_if_port_taken() uvicorn.run(app, host="0.0.0.0", port=PORT)