"""Pocket TTS wrapper for the gateway — the cloned Jarvis voice, served over WS. Ports the load / synth / mood-clip logic from ``voice/tts_pocket.py`` (Kyutai Pocket TTS, MIT — a 100M CALM model that runs ~7x realtime on two CPU cores, ~40ms to first audio, ~200MB RAM), plus the emotion-tag parsing from ``voice/tts_chatterbox.py``. Reproduced here so ``gateway/`` is a standalone deployable with no dependency on the ``voice/`` package. Key differences from the desktop module: * The model is LAZY-loaded on the first voice synth (``get_tts()``); a chat-only WS session never imports or loads it. ``pocket_tts`` / torch are NOT in the base ``requirements.txt`` — install ``requirements-voice.txt`` on the voice-capable deployment (see README). * Voice refs are read from ``$VOICE_DIR`` — default: the repo's ``voices/`` dir for local dev, a Railway volume in prod. Voiceprints NEVER ship in git. * Output is 16-bit PCM WAV bytes per sentence (``synth_wav``), ready to push straight down the socket after its ``sentence`` frame. Env: VOICE_DIR (voice-ref directory), POCKET_VOICE (catalog fallback voice), VOICE_REF (base clone wav basename, default jarvis_ref.wav), POCKET_MAX_TOKENS (default 300 ≈ 24s cap). """ from __future__ import annotations import io import os import re import threading import wave from pathlib import Path import numpy as np SAMPLE_RATE = 24_000 # pocket-tts native rate # Default voice dir = the repo's voices/ (this file is /gateway/tts.py). _REPO_ROOT = Path(__file__).resolve().parent.parent _DEFAULT_VOICE_DIR = _REPO_ROOT / "voices" def voice_dir() -> Path: return Path(os.environ.get("VOICE_DIR") or _DEFAULT_VOICE_DIR) # --------------------------------------------------------------------------- # Emotion tags (ported from tts_chatterbox). The brain prefixes a sentence with # one tag; it is stripped before synthesis and never spoken. On Pocket the tag # selects a mood *reference clip* rather than an exaggeration dial. # --------------------------------------------------------------------------- _EMOTIONS = ( "calm", "flat", "warm", "happy", "excited", "urgent", "alarmed", "sad", "sarcastic", "dry", "serious", "neutral", ) _TAG_RE = re.compile(r"^\s*\[(" + "|".join(_EMOTIONS) + r")\]\s*", re.IGNORECASE) # models improvise tags outside the palette ("[Mildly excited]") — still strip # them (never spoken); the last word often maps anyway. _ANY_TAG_RE = re.compile(r"^\s*\[[a-z][a-z ,'-]{0,24}\]\s*", re.IGNORECASE) # emotion tag -> reference-clip stem (voices/jarvis_.wav). Several tags # share one mood clip. (ported from tts_pocket._TAG_TO_STEM) _TAG_TO_STEM = { "calm": "calm", "flat": "calm", "serious": "calm", "warm": "warm", "happy": "warm", "excited": "excited", "urgent": "excited", "alarmed": "excited", "sad": "sad", "sarcastic": "dry", "dry": "dry", "neutral": None, } # emotion tag -> the protocol's coarse emotion enum (warm|dry|excited|sad|calm| # null) sent on the `sentence` frame for the UI. _TAG_TO_PROTOCOL = { "calm": "calm", "flat": "calm", "serious": "calm", "warm": "warm", "happy": "warm", "excited": "excited", "urgent": "excited", "alarmed": "excited", "sad": "sad", "sarcastic": "dry", "dry": "dry", "neutral": None, } def split_emotion(text: str) -> tuple[str, str | None]: """Strip a leading [emotion] tag; return (clean_text, tag_name_or_None). Unknown tone tags are stripped too (a TTS must never read one aloud); their last word is tried as a known tag ("[Mildly excited]" -> "excited").""" m = _TAG_RE.match(text or "") if m: return text[m.end():].lstrip(), m.group(1).lower() m = _ANY_TAG_RE.match(text or "") if m: inner = m.group(0).strip()[1:-1].strip().lower() key = inner.split()[-1] if inner else "" return text[m.end():].lstrip(), (key if key in _TAG_TO_STEM else None) return text, None def emotion_label(text: str) -> tuple[str, str | None]: """(clean_text, protocol_emotion) for the `sentence` frame. protocol_emotion is one of warm|dry|excited|sad|calm or None.""" clean, tag = split_emotion(text) return clean, (_TAG_TO_PROTOCOL.get(tag) if tag else None) def stem_for(text: str) -> str | None: """The mood-clip stem to voice this (tagged) sentence with, or None (base).""" _clean, tag = split_emotion(text) return _TAG_TO_STEM.get(tag) if tag else None # --------------------------------------------------------------------------- # WAV encoding — pure, testable, no model. float32 mono @ rate -> 16-bit PCM WAV. # --------------------------------------------------------------------------- def float32_to_wav_bytes(samples: np.ndarray, rate: int = SAMPLE_RATE) -> bytes: """Encode a float32 [-1, 1] mono array as a complete 16-bit PCM WAV blob.""" if samples is None or len(samples) == 0: pcm = b"" else: clipped = np.clip(np.asarray(samples, dtype=np.float32), -1.0, 1.0) pcm = (clipped * 32767.0).astype(" None: from pocket_tts import TTSModel # heavy import — kept lazy self._model = TTSModel.load_model() self._lock = threading.Lock() # model state isn't re-entrant self._states: dict[str, object] = {} self._cloned = False vdir = voice_dir() ref_name = os.environ.get("VOICE_REF", "jarvis_ref.wav") base_ref = ref_name if os.path.isabs(ref_name) else str(vdir / ref_name) self._base_name = self._load_base(base_ref) self._preload_emotions(vdir) def _load_base(self, ref: str) -> str: if os.path.isfile(ref): try: self._states[""] = self._model.get_state_for_audio_prompt(ref) self._cloned = True print(f" [tts] voice cloned from {os.path.basename(ref)}") return ref except Exception as e: # noqa: BLE001 — gated weights / bad wav print(f" [tts] clone failed ({str(e)[:80]}) — catalog " f"'{_CATALOG_FALLBACK}'") self._states[""] = self._model.get_state_for_audio_prompt( _CATALOG_FALLBACK) return _CATALOG_FALLBACK def _preload_emotions(self, vdir: Path) -> None: if not self._cloned: return # mood clips are clones; the catalog voice has no moods for stem in dict.fromkeys(s for s in _TAG_TO_STEM.values() if s): path = vdir / f"jarvis_{stem}.wav" if not path.is_file(): continue try: self._states[stem] = self._model.get_state_for_audio_prompt( str(path)) except Exception as e: # noqa: BLE001 print(f" [tts] mood clip {path.name} failed ({str(e)[:60]})") moods = [k for k in self._states if k] if moods: print(f" [tts] mood clones ready ({', '.join(sorted(moods))})") def _state_for(self, text: str): stem = stem_for(text) clean, _tag = split_emotion(text) return clean, self._states.get(stem or "", self._states[""]) def synth(self, text: str) -> np.ndarray: """(optionally [emotion]-tagged) text -> float32 mono PCM @ 24kHz.""" clean, state = self._state_for(text) if not clean.strip(): return np.zeros(0, dtype=np.float32) with self._lock: wav = self._model.generate_audio(state, clean, max_tokens=_MAX_TOKENS) return wav.detach().cpu().numpy().astype(np.float32).reshape(-1) def synth_wav(self, text: str) -> bytes: """(optionally [emotion]-tagged) text -> complete 16-bit PCM WAV bytes.""" return float32_to_wav_bytes(self.synth(text), SAMPLE_RATE) # ---- lazy singleton ------------------------------------------------------- _INSTANCE: PocketTTS | None = None _INIT_LOCK = threading.Lock() def get_tts() -> PocketTTS: """Load-on-first-use accessor. Raises if pocket-tts isn't installed — the caller turns that into an {"type":"error"} frame; chat still works.""" global _INSTANCE if _INSTANCE is None: with _INIT_LOCK: if _INSTANCE is None: _INSTANCE = PocketTTS() return _INSTANCE