"""Real Korean TTS via MeloTTS, run as a persistent out-of-venv worker. MeloTTS needs its own interpreter (melo311). Loading the model costs several seconds, so we keep one worker process alive and stream synthesis requests to it (see melo_worker.py for the protocol). Each `speak()` writes a wav to `out_dir` and hands the path to a sink (default: log it). The Discord voice integration later swaps the sink for "play this wav into the call". Env: WSAI_MELO_PYTHON interpreter with melo installed (default: /home/claude/jarvis-tts/melo311/bin/python) WSAI_MELO_DEVICE cpu | cuda | auto (default auto: GPU if torch sees one, else CPU; the worker falls back to CPU if CUDA fails) WSAI_TTS_OUT_DIR where wavs are written (default ~/.cache/wsai/tts) WSAI_TTS_SPEED synthesis speed multiplier (default 1.3) """ from __future__ import annotations import asyncio import collections import json import logging import os import time from pathlib import Path from typing import Awaitable, Callable from ..interfaces import Reply log = logging.getLogger("wsai.tts.melo") _DEFAULT_PYTHON = "/home/claude/jarvis-tts/melo311/bin/python" # A sink receives the finished wav path plus the reply it voices. Sink = Callable[[str, Reply], Awaitable[None]] async def _log_sink(path: str, reply: Reply) -> None: log.info("TTS wav ready: %s (%s)", path, reply.text[:40]) class MeloTTS: def __init__( self, *, python: str | None = None, device: str | None = None, out_dir: str | None = None, speed: float | None = None, sink: Sink | None = None, ) -> None: self.python = python or os.environ.get("WSAI_MELO_PYTHON", _DEFAULT_PYTHON) self.device = device or os.environ.get("WSAI_MELO_DEVICE", "auto") self.out_dir = Path(out_dir or os.environ.get("WSAI_TTS_OUT_DIR") or (Path.home() / ".cache/wsai/tts")) self.speed = float(speed if speed is not None else os.environ.get("WSAI_TTS_SPEED", "1.3")) self.sink = sink or _log_sink self._proc: asyncio.subprocess.Process | None = None self._lock = asyncio.Lock() self._n = 0 self.load_ms: int | None = None # Keep the worker's most recent stderr lines so a crash reports its real # cause instead of a bare JSONDecodeError. Bounded so it can't grow. self._stderr_tail: collections.deque[str] = collections.deque(maxlen=40) self._stderr_task: asyncio.Task | None = None async def _drain_stderr(self, stream: asyncio.StreamReader) -> None: # The worker redirects fd1 -> fd2, so ALL library chatter lands on # stderr. If we PIPE stderr but never read it, the OS pipe buffer fills # and the worker blocks forever. So we continuously drain it and keep # only the last few lines for diagnostics. try: while True: line = await stream.readline() if not line: return self._stderr_tail.append(line.decode(errors="replace").rstrip()) except asyncio.CancelledError: raise except Exception: # draining must never crash the caller return def _stderr_hint(self) -> str: tail = "\n".join(self._stderr_tail) return f" worker stderr tail:\n{tail}" if tail else " (worker produced no stderr)" async def warmup(self) -> None: """Start and load the worker now so the first real utterance is warm. Called at pipeline startup so users don't wait ~7s (CPU model load) for the very first spoken reply.""" await self._ensure() async def _ensure(self) -> None: if self._proc is not None and self._proc.returncode is None: return self.out_dir.mkdir(parents=True, exist_ok=True) env = {**os.environ, "WSAI_MELO_DEVICE": self.device} # Run the worker module from the wsai source tree with the melo venv. repo_root = str(Path(__file__).resolve().parents[2]) self._proc = await asyncio.create_subprocess_exec( self.python, "-m", "wsai.backends.melo_worker", cwd=repo_root, env=env, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) self._stderr_tail.clear() assert self._proc.stderr is not None self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr)) ready = await self._proc.stdout.readline() if not ready: # worker died before signalling ready await self._proc.wait() raise RuntimeError( f"melo worker exited before ready (code {self._proc.returncode})." f"{self._stderr_hint()}" ) try: info = json.loads(ready.decode()) except json.JSONDecodeError as exc: raise RuntimeError( f"melo worker sent invalid ready line {ready!r}: {exc}." f"{self._stderr_hint()}" ) from exc if not info.get("ready"): raise RuntimeError( f"melo worker failed to start: {info}.{self._stderr_hint()}" ) self.load_ms = info.get("ms") log.info("melo worker ready in %s ms on %s", self.load_ms, info.get("device")) async def speak(self, reply: Reply) -> None: await self._ensure() self._n += 1 out = str(self.out_dir / f"tts-{self._n:06d}.wav") req = json.dumps({"text": reply.text, "out": out, "speed": self.speed}) s = time.monotonic() async with self._lock: assert self._proc and self._proc.stdin and self._proc.stdout self._proc.stdin.write((req + "\n").encode()) await self._proc.stdin.drain() resp = await self._proc.stdout.readline() if not resp: raise RuntimeError(f"melo worker closed unexpectedly.{self._stderr_hint()}") res = json.loads(resp.decode()) if not res.get("ok"): raise RuntimeError(f"melo synth failed: {res.get('error')}") log.debug("synth %d ms (worker %s ms)", int((time.monotonic() - s) * 1000), res.get("ms")) await self.sink(res["out"], reply) async def aclose(self) -> None: if self._proc is not None and self._proc.returncode is None: try: self._proc.terminate() await asyncio.wait_for(self._proc.wait(), timeout=5) except (ProcessLookupError, asyncio.TimeoutError): pass if self._stderr_task is not None: self._stderr_task.cancel() try: await self._stderr_task except (asyncio.CancelledError, Exception): pass self._stderr_task = None self._proc = None