Files
watch_sceen_ai/wsai/backends/melo_worker.py
EJClaw 77d7cd8b56 perf(voice): warm up STT+TTS workers before signalling ready
The first CUDA inference pays a large lazy cost (kernel autotune/cudnn) —
~10s for a cold TTS synth — which would blow the voice loop's ~1s budget on
the very first reply. Each worker now runs one dummy inference (TTS: a short
phrase; STT: 1s of silence) after model load and before emitting "ready", so
"ready" means "hot". Warmup failures are logged and never block startup.

Verified: first real call after startup is now TTS ~238ms / STT ~189ms
(was ~11s cold for TTS). 12 tests pass.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-08-18 21:06:23 +09:00

110 lines
4.0 KiB
Python

"""Persistent MeloTTS worker (Korean).
MeloTTS lives in its own Python (melo311); loading the model takes seconds, so
we load it ONCE here and then serve synthesis requests over stdin/stdout. This
process is launched with the melo311 interpreter by wsai.backends.melo.MeloTTS.
MeloTTS (and its deps) print progress straight to stdout, which would corrupt
the JSON protocol. So on startup we split the streams: a private duplicate of
the original stdout carries the protocol, and fd 1 is redirected to fd 2 so all
library chatter lands on stderr instead.
Protocol (one JSON object per line, on the protocol channel):
<- {"text": "...", "out": "/abs/path.wav", "speed": 1.3}
-> {"ok": true, "out": "/abs/path.wav", "ms": 123}
-> {"ok": false, "error": "..."}
On startup, once the model is ready, it emits exactly one line:
-> {"ready": true, "ms": <load-ms>, "device": "cpu"}
"""
import json
import os
import sys
import time
# Split protocol from library noise BEFORE importing anything heavy.
_proto = os.fdopen(os.dup(1), "w", buffering=1) # private copy of real stdout
os.dup2(2, 1) # fd1 -> stderr, so stray library prints don't hit the protocol
def _emit(obj: dict) -> None:
_proto.write(json.dumps(obj) + "\n")
_proto.flush()
def _log(*a):
print(*a, file=sys.stderr, flush=True)
def main() -> None:
lang = "KR"
requested = os.environ.get("WSAI_MELO_DEVICE", "auto") # cpu | cuda | auto
from melo.api import TTS # heavy import; only in the melo venv
def _has_cuda() -> bool:
try:
import torch
return torch.cuda.is_available()
except Exception:
return False
device = requested
if requested == "auto":
device = "cuda" if _has_cuda() else "cpu"
t0 = time.monotonic()
try:
tts = TTS(language=lang, device=device)
except Exception as exc:
# CUDA picked but unusable (CPU-only torch, missing libs, OOM): fall back
# to CPU rather than leaving the whole voice loop dead.
if device == "cuda":
_log(f"[melo_worker] CUDA load failed ({exc}); falling back to CPU")
device = "cpu"
tts = TTS(language=lang, device=device)
else:
raise
speaker_id = tts.hps.data.spk2id[lang]
load_ms = int((time.monotonic() - t0) * 1000)
# Warm up before signalling ready: the first CUDA synth pays a large lazy
# cost (kernel autotune/cudnn), ~10s cold vs ~130ms hot, which would blow the
# voice loop's ~1s budget on the very first reply. Do that dummy synth here so
# "ready" means "hot". Failures must not block startup.
warmup_ms = None
try:
warm_out = os.path.expanduser("~/.cache/wsai/tts/_warmup.wav")
os.makedirs(os.path.dirname(warm_out), exist_ok=True)
w = time.monotonic()
tts.tts_to_file("워밍업", speaker_id, warm_out, speed=1.3)
warmup_ms = int((time.monotonic() - w) * 1000)
except Exception as exc:
_log(f"[melo_worker] warmup skipped: {exc}")
_emit({"ready": True, "ms": load_ms, "device": device, "warmup_ms": warmup_ms})
_log(f"[melo_worker] model ready in {load_ms} ms on {device} (warmup {warmup_ms} ms)")
for line in sys.stdin:
line = line.strip()
if not line:
continue
try:
req = json.loads(line)
text = req["text"]
out = req["out"]
speed = float(req.get("speed", 1.0))
if out.startswith("/tmp") or out.startswith("/dev/shm"):
raise ValueError(f"refusing RAM-backed tmpfs path: {out}")
s = time.monotonic()
tts.tts_to_file(text, speaker_id, out, speed=speed)
ms = int((time.monotonic() - s) * 1000)
_emit({"ok": True, "out": out, "ms": ms})
except Exception as exc: # keep the worker alive across bad requests
_emit({"ok": False, "error": f"{type(exc).__name__}: {exc}"})
_log(f"[melo_worker] error: {exc}")
if __name__ == "__main__":
main()