"""Multi-engine TTS with one-model-in-VRAM-at-a-time eviction. The CUDA box is shared with other GPU workloads, so this module keeps at most ONE engine's model resident in VRAM. Switching engines (or going idle) unloads the current model and returns its VRAM to the driver before loading the next. All load / unload / generate calls MUST run on the server's single GPU executor thread (Metal/MPS/CUDA thread affinity + serialization). The ModelManager does no locking of its own beyond a light guard because the executor already serializes every GPU job, including the idle-eviction sweep. Engines: - QwenEngine : qwen_tts (existing pytorch backend), GPU voice prompts. - VoxcpmEngine : VoxCPM2, voices = reference clip (timbre) + optional tone seed (prompt). Voice state is just file paths -> no GPU tensors to juggle across eviction. - DotsEngine : dots.tts, voices = reference clip (+ ref_text in continuation mode). No tone seed. Also file-path-only voice state. - CosyvoiceEngine : CosyVoice 3 via a sidecar process (own torch pin). """ import os import gc import io import hashlib import time import threading from pathlib import Path import numpy as np import soundfile as sf import requests def pcm16(audio_f32) -> bytes: """float32 [-1,1] mono -> little-endian s16le PCM bytes.""" a = np.clip(np.asarray(audio_f32, dtype=np.float32), -1.0, 1.0) return (a * 32767.0).astype(" bool: raise NotImplementedError def prepare_voice(self, voice_id: str, voice_dir: Path, meta: dict): """Ensure any per-voice GPU artifacts exist (model is resident).""" def generate(self, text, voice_id, voice_dir, meta, language, gen_kwargs): """Return (audio_float32_1d, sample_rate).""" raise NotImplementedError def stream(self, text, voice_id, voice_dir, meta, language, gen_kwargs): """Yield audio_float32_1d chunks. Default: one chunk via generate().""" audio, _sr = self.generate(text, voice_id, voice_dir, meta, language, gen_kwargs) yield audio # --------------------------------------------------------------------------- # Qwen3-TTS engine (pytorch) # --------------------------------------------------------------------------- class QwenEngine(Engine): name = "qwen" sample_rate = 24000 AVAILABLE_MODELS = { "0.6B": "Qwen3-TTS-12Hz-0.6B-Base", "1.7B": "Qwen3-TTS-12Hz-1.7B-Base", } def __init__(self, models_dir: Path): self._models_dir = models_dir self._model = None self.model_name = None self._prompts: dict[str, list] = {} # voice_id -> List[VoiceClonePromptItem] (GPU) @property def loaded(self) -> bool: return self._model is not None def load(self): import torch from qwen_tts import Qwen3TTSModel device = "cuda" if torch.cuda.is_available() else ( "mps" if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available() else "cpu") key = os.environ.get("POLYTTS_MODEL", "1.7B") if key not in self.AVAILABLE_MODELS: key = "1.7B" self.model_name = self.AVAILABLE_MODELS[key] path = self._models_dir / self.model_name if not path.exists(): raise RuntimeError(f"Qwen model not found: {path}") print(f"[qwen] loading {self.model_name} on {device} …", flush=True) self._model = Qwen3TTSModel.from_pretrained( str(path), device_map=device, dtype=torch.float32, attn_implementation="sdpa") print("[qwen] loaded.", flush=True) def unload(self): # GPU voice prompts are tied to the model device; drop them and rebuild # lazily on next use rather than shuttling tensors CPU<->GPU. self._prompts.clear() self._model = None gc.collect() _free_cuda() _trim_ram() print("[qwen] unloaded.", flush=True) def prepare_voice(self, voice_id, voice_dir, meta): if voice_id in self._prompts: return wav_path = voice_dir / "voice.wav" self._prompts[voice_id] = self._model.create_voice_clone_prompt( ref_audio=str(wav_path), ref_text=meta["ref_text"], x_vector_only_mode=bool(meta.get("x_vector_only_mode", False)), ) def generate(self, text, voice_id, voice_dir, meta, language, gen_kwargs): self.prepare_voice(voice_id, voice_dir, meta) # Strip VoxCPM/CosyVoice-only knobs before forwarding to the qwen model. gk = {k: v for k, v in gen_kwargs.items() if k not in ("cfg_value", "inference_timesteps", "denoise", "instruct")} wavs, sr = self._model.generate_voice_clone( text=text, language=language or "Chinese", voice_clone_prompt=self._prompts[voice_id], **gk) return np.asarray(wavs[0], dtype=np.float32), sr # --------------------------------------------------------------------------- # VoxCPM2 engine # --------------------------------------------------------------------------- class VoxcpmEngine(Engine): name = "voxcpm" sample_rate = 48000 def __init__(self): self._model = None self.model_name = os.environ.get("VOXCPM_MODEL_ID", "openbmb/VoxCPM2") self._cfg = float(os.environ.get("VOXCPM_CFG_VALUE", "3.3")) self._steps = int(os.environ.get("VOXCPM_TIMESTEPS", "10")) @property def loaded(self) -> bool: return self._model is not None def load(self): from voxcpm import VoxCPM print(f"[voxcpm] loading {self.model_name} …", flush=True) self._model = VoxCPM.from_pretrained(self.model_name, load_denoiser=False) self.sample_rate = self._model.tts_model.sample_rate # from_pretrained stages the full checkpoint through CPU tensors before # moving the warm model to CUDA. Reclaim those now-unused host buffers # without unloading the GPU-resident model. gc.collect() _trim_ram() print(f"[voxcpm] loaded. sr={self.sample_rate}", flush=True) def unload(self): # Voice state is file paths only -> nothing GPU-resident to drop. self._model = None gc.collect() _free_cuda() _trim_ram() print("[voxcpm] unloaded.", flush=True) def _clone_kwargs(self, voice_dir: Path, meta: dict, gen_kwargs: dict | None = None) -> dict: """reference clip = timbre; optional seed clip = locked tone. Per-request overrides (cfg_value / inference_timesteps / denoise) are read from gen_kwargs; absent -> engine env defaults.""" gk = gen_kwargs or {} kw = dict(reference_wav_path=str(voice_dir / "voice.wav"), cfg_value=gk.get("cfg_value", self._cfg), inference_timesteps=gk.get("inference_timesteps", self._steps), normalize=True) seed = voice_dir / "seed.wav" if seed.exists() and meta.get("seed_text"): kw["prompt_wav_path"] = str(seed) kw["prompt_text"] = meta["seed_text"] if gk.get("denoise"): kw["denoise"] = True return kw def generate(self, text, voice_id, voice_dir, meta, language, gen_kwargs): audio = self._model.generate(text=text, **self._clone_kwargs(voice_dir, meta, gen_kwargs)) return np.asarray(audio, dtype=np.float32), self.sample_rate def stream(self, text, voice_id, voice_dir, meta, language, gen_kwargs): ck = self._clone_kwargs(voice_dir, meta, gen_kwargs) for chunk in self._model.generate_streaming(text=text, **ck): yield np.asarray(chunk, dtype=np.float32) # --------------------------------------------------------------------------- # dots.tts engine # --------------------------------------------------------------------------- class DotsEngine(Engine): """dots.tts — fully-continuous AR TTS (Qwen2.5-1.5B backbone + a flow-matching DiT head over a 48 kHz AudioVAE, with a frozen CAM++ speaker x-vector as a side input). Voice state is file paths only, like VoxCPM. Two cloning paths, and which one a voice gets is DECLARED by its meta, never inferred: ``x_vector_only_mode`` -> timbre embedding alone (prompt audio, no transcript); otherwise *continuation*, which additionally feeds ``ref_text`` as ``prompt_text``. The two measure materially differently on our own voices (0.838 vs 0.762 speaker similarity on a VoxAlert pack), so this engine must not quietly substitute one for the other — a silent downgrade surfaces months later as "the voice sounds different now". See openspec/changes/add-dots-tts-engine/design.md. """ name = "dots" sample_rate = 48000 def __init__(self): self._rt = None self.model_name = os.environ.get("POLYTTS_DOTS_MODEL", "dots-studio/dots.tts-mf") # NFE 4 + optimize=True are LOAD-BEARING, not tuning. The MeanFlow # checkpoint ships no `sampling` block, so the library falls back to # (euler, 10 steps, cfg 1.2) — at which these same weights measured # RTF 1.81-2.02 on an RTX 3090, i.e. SLOWER THAN REAL TIME, against # 0.23-0.37 at 4 steps with optimize=True. Do not "simplify" these # constants away: at the library defaults this engine reads as a # regression and the result is buried. self._steps = int(os.environ.get("POLYTTS_DOTS_NUM_STEPS", "4")) self._optimize = os.environ.get("POLYTTS_DOTS_OPTIMIZE", "1").lower() not in ("0", "false", "no") @property def loaded(self) -> bool: return self._rt is not None def load(self): from dots_tts.runtime import DotsTtsRuntime print(f"[dots] loading {self.model_name} (steps={self._steps}, " f"optimize={self._optimize}) …", flush=True) self._rt = DotsTtsRuntime.from_pretrained( self.model_name, precision="bfloat16", optimize=self._optimize) self.sample_rate = int(self._rt.sample_rate) # from_pretrained stages the checkpoint through host memory before the # warm model lands on CUDA; hand those pages back without unloading the # GPU-resident model (same reclaim as voxcpm). gc.collect() _trim_ram() print(f"[dots] loaded. sr={self.sample_rate}", flush=True) def unload(self): # Voice state is file paths only -> nothing GPU-resident to drop. self._rt = None gc.collect() _free_cuda() _trim_ram() print("[dots] unloaded.", flush=True) def _clone_kwargs(self, voice_dir: Path, meta: dict, language) -> dict: """reference clip = timbre; ref_text = the transcript the model continues from, in continuation mode only. VoxCPM/Qwen per-request knobs (cfg_value, inference_timesteps, temperature, ...) are deliberately NOT forwarded. They are differently scaled here — VoxCPM's cfg_value 3.3 against dots' guidance_scale 1.2 — so passing them through would change quality under a name that means something else on this engine. """ kw = dict(prompt_audio_path=str(voice_dir / "voice.wav"), num_steps=self._steps) if not meta.get("x_vector_only_mode", False): kw["prompt_text"] = meta["ref_text"] if language: # normalize_language_code() accepts our qwen-style names ("Chinese", # "English") via langcodes and returns None for anything it cannot # resolve, so an unknown language degrades to auto-detect. kw["language"] = language return kw # A generation is a COLLAPSE when the model emits EOS almost immediately and # returns a fraction of a second of audio for a whole sentence. English runs # ~14 characters of text per second of speech; under this fraction of that # estimate is not a short reading, it is a failure. # # DETECTION ONLY — there is deliberately no retry. Measured on the reference # that produces it (hl-hev-suit, 30 s of Half-Life HEV announcer audio): the # collapse is reproducible across three different RNG seeds, across # prompt_text truncated to 289/200/120/60 characters, and across four # rewordings of the target line. The seed cannot move it because the seed # feeds the flow-matching noise while the EOS decision is the AR backbone's, # which runs deterministically. A redraw costs a full generation and returns # the same bytes, so this logs and moves on rather than pretending to recover. # # Rate on this host: 2 of 60 generations (3.3%) over ten voice packs and six # realistic narration lines — and BOTH were hl-hev-suit (2 of its 6 lines). # The other nine packs were clean across 54 generations. It is a property of # particular (reference, text) pairs, not a background flake rate. _COLLAPSE_FLOOR_RATIO = 0.25 _COLLAPSE_MIN_S = 0.35 def _seed(self, voice_id: str, text: str): """dots.tts draws its flow-matching noise from the GLOBAL torch RNG and exposes no seed argument, so a repeat of the same (text, voice) would otherwise differ byte-for-byte — which the /tts/stream disk PCM cache assumes it does not. Pin the RNG per (voice, text) before every generation. The text is IN the seed deliberately. Seeding from the voice alone makes every phrase in that voice start from one RNG state, so a voice that draws badly draws badly across the board; per-pair seeding keeps one bad draw from correlating with the next. This is an RNG seed, not a tone lock: the engine has no tone-seed input, so VoxCPM's `seed.wav` / `seed_text` have no meaning here and are ignored. """ import torch h = hashlib.sha256(f"{voice_id}\x00{text}".encode()).hexdigest()[:8] torch.manual_seed(int(h, 16)) def _collapsed(self, text: str, samples: int) -> bool: """Did this generation fail rather than merely run short? Reported so the pair is identifiable. A caller cannot tell 0.16 s of audio from a very short line, and the /tts/stream disk cache will store the short body and serve it forever after, so without this the failure is invisible in the logs and permanent in the cache. """ floor = max(self._COLLAPSE_MIN_S, (len(text) / 14.0) * self._COLLAPSE_FLOOR_RATIO) return (samples / float(self.sample_rate)) < floor def _warn_if_collapsed(self, text, voice_id, samples): if self._collapsed(text, samples): print(f"[dots] COLLAPSED generation: {samples / self.sample_rate:.2f}s of " f"audio for {len(text)} chars, voice={voice_id}. This pair is " f"deterministic — it will not clear on retry, and the PCM cache " f"will keep it. Re-cut the reference or use a different voice.", flush=True) def generate(self, text, voice_id, voice_dir, meta, language, gen_kwargs): self._seed(voice_id, text) out = self._rt.generate(text=text, **self._clone_kwargs(voice_dir, meta, language)) audio = out["audio"].detach().float().cpu().numpy().reshape(-1) self._warn_if_collapsed(text, voice_id, audio.size) return np.asarray(audio, dtype=np.float32), int(out.get("sample_rate", self.sample_rate)) def stream(self, text, voice_id, voice_dir, meta, language, gen_kwargs): self._seed(voice_id, text) n = 0 for chunk in self._rt.generate_stream( text=text, **self._clone_kwargs(voice_dir, meta, language)): a = chunk.detach().float().cpu().numpy().reshape(-1) n += a.size yield a self._warn_if_collapsed(text, voice_id, n) # --------------------------------------------------------------------------- # CosyVoice 3 engine (via an isolated sidecar process) # --------------------------------------------------------------------------- class CosyvoiceEngine(Engine): """CosyVoice 3 — zero-shot voice clone + instruct (emotion/style) control, the one real TONE lever (VoxCPM cannot vary tone). CosyVoice pins torch 2.3.1 which conflicts with this venv (torch 2.12 for qwen/voxcpm), so it runs in a separate `cosyvoice` conda env as a sidecar HTTP service (cosyvoice_worker.py on 127.0.0.1:8101). This engine is a thin client; PolyTTS's ModelManager still orchestrates VRAM — load()/unload() ask the sidecar to load/free the model (evicting voxcpm first since the manager keeps one engine in VRAM).""" name = "cosyvoice" sample_rate = 24000 def __init__(self, base_url="http://127.0.0.1:8101"): self.base_url = base_url self._loaded = False self.model_name = "Fun-CosyVoice3-0.5B" @property def loaded(self) -> bool: return self._loaded def load(self): try: r = requests.post(f"{self.base_url}/load", timeout=600) r.raise_for_status() self.sample_rate = int(r.json().get("sample_rate", self.sample_rate)) self._loaded = True print(f"[cosyvoice] sidecar model loaded (sr={self.sample_rate})", flush=True) except Exception as e: raise RuntimeError( f"cosyvoice sidecar load failed — is cosyvoice_worker.py running on " f"{self.base_url}? (conda run -n cosyvoice python cosyvoice_worker.py): {e}") def unload(self): try: requests.post(f"{self.base_url}/unload", timeout=120) except Exception: pass self._loaded = False print("[cosyvoice] sidecar model unloaded", flush=True) def _synth(self, text, voice_dir, meta, gen_kwargs): instruct = gen_kwargs.get("instruct") body = {"text": text, "voice_wav_path": str(voice_dir / "voice.wav"), "ref_text": meta.get("ref_text", "")} if instruct: body["instruct"] = instruct r = requests.post(f"{self.base_url}/tts", json=body, timeout=900) r.raise_for_status() sr = int(r.headers.get("X-Sample-Rate", self.sample_rate)) audio, _ = sf.read(io.BytesIO(r.content), dtype="float32") return np.asarray(audio, dtype=np.float32), sr def generate(self, text, voice_id, voice_dir, meta, language, gen_kwargs): return self._synth(text, voice_dir, meta, gen_kwargs) def stream(self, text, voice_id, voice_dir, meta, language, gen_kwargs): # CosyVoice streams per-sentence internally; default to one chunk here. audio, sr = self._synth(text, voice_dir, meta, gen_kwargs) self.sample_rate = sr yield audio # --------------------------------------------------------------------------- # One-model-in-VRAM manager # --------------------------------------------------------------------------- class ModelManager: def __init__(self, engines: dict[str, Engine], idle_seconds: int): self.engines = engines self.idle_seconds = idle_seconds self.resident: str | None = None self.last_used = time.monotonic() self._guard = threading.Lock() def ensure(self, name: str) -> Engine: """Make `name` the resident engine, evicting any other. GPU-thread only.""" if name not in self.engines: raise KeyError(f"unknown engine: {name}") with self._guard: if self.resident != name: if self.resident is not None: self.engines[self.resident].unload() self.resident = None self.engines[name].load() self.resident = name self.last_used = time.monotonic() return self.engines[name] def unload_now(self) -> str | None: """Force-evict the resident model now, regardless of idle time. Returns the name of the engine that was unloaded (or None if already empty). GPU-thread only. Used to hand VRAM/RAM to a co-resident workload (e.g. the local renderer) without killing the server process.""" with self._guard: evicted = self.resident if evicted is not None: self.engines[evicted].unload() self.resident = None print(f"[manager] force-unloaded {evicted}", flush=True) else: _trim_ram() return evicted def maybe_evict(self) -> bool: """Evict the resident model if idle past the timeout. GPU-thread only.""" with self._guard: if self.resident and (time.monotonic() - self.last_used) > self.idle_seconds: evicted = self.resident self.engines[evicted].unload() self.resident = None print(f"[manager] idle-evicted {evicted}", flush=True) return True return False def status(self) -> dict: return { "resident": self.resident, "idle_seconds": self.idle_seconds, "idle_for": round(time.monotonic() - self.last_used, 1) if self.resident else None, "engines": list(self.engines.keys()), }