watch_sceen_ai/wsai/backends/melo_worker.py

"""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}
    <- {"segments": [{"text": "...", "speed": 1.3, "pitch": 2.0}, ...],
        "out": "/abs/path.wav"}   # expressive form: per-segment speed + pitch
    -> {"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"}

``pitch`` is a semitone offset applied to that segment's wav (0 == no shift) so
emotion tags can raise/lower the voice without changing the words. Segments are
synthesised independently and concatenated with a short gap so a single reply can
carry several emotions.
"""

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]
    sr = tts.hps.data.sampling_rate
    load_ms = int((time.monotonic() - t0) * 1000)

    import numpy as np
    import soundfile

    _GAP = np.zeros(int(sr * 0.12), dtype=np.float32)  # 120 ms between segments

    def _pitch_shift(audio, semitones: float):
        if not semitones:
            return audio
        import librosa

        return librosa.effects.pitch_shift(
            audio.astype(np.float32), sr=sr, n_steps=float(semitones)
        )

    def _synth_segments(segments: list[dict], out: str) -> None:
        """Synthesize each segment, pitch-shift it, and concatenate to one wav."""
        pieces = []
        for i, seg in enumerate(segments):
            text = seg["text"]
            if not text.strip():
                continue
            speed = float(seg.get("speed", 1.0))
            pitch = float(seg.get("pitch", 0.0))
            audio = tts.tts_to_file(text, speaker_id, None, speed=speed)
            audio = _pitch_shift(np.asarray(audio, dtype=np.float32), pitch)
            if pieces:
                pieces.append(_GAP)
            pieces.append(audio)
        if not pieces:
            raise ValueError("no speakable segment")
        soundfile.write(out, np.concatenate(pieces), sr)

    # 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)
        # Also JIT-warm librosa's pitch shifter (first call pays ~0.4s numba
        # compile) so the first *emotional* reply doesn't stall.
        import librosa

        librosa.effects.pitch_shift(np.zeros(sr, dtype=np.float32), sr=sr, n_steps=1.0)
        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)
            out = req["out"]
            if out.startswith("/tmp") or out.startswith("/dev/shm"):
                raise ValueError(f"refusing RAM-backed tmpfs path: {out}")
            s = time.monotonic()
            if "segments" in req:
                _synth_segments(req["segments"], out)
            else:  # legacy single-utterance form
                speed = float(req.get("speed", 1.0))
                tts.tts_to_file(req["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()