bililive-touhou-tts/tts.py

"""Node.js TTS Bridge — Persistent Process Manager.

Starts a long-lived Node.js subprocess running tts_bridge.js. Communication uses
stdin/stdout with temporary files for text and WAV data.

Supports PyInstaller frozen mode: uses local node.exe and tts_bridge.js.
"""

import asyncio
import os
import pathlib
import sys
import tempfile
import logging

logger = logging.getLogger(__name__)


def _get_portable_root() -> pathlib.Path:
    return pathlib.Path(sys.executable).parent.parent.parent


def _get_node_exe() -> str:
    if getattr(sys, "frozen", False):
        p = _get_portable_root() / "node.exe"
        if p.exists():
            return str(p)
    return "node"


def _resolve_path(rel: str) -> str:
    if getattr(sys, "frozen", False):
        p = _get_portable_root() / rel
        if p.exists():
            return str(p)
    return str(pathlib.Path(__file__).parent / rel)


class TTSBridge:
    """Manages a persistent Node.js subprocess for aquestalk.js TTS synthesis."""

    def __init__(self, voice: str = "f1", speed: int = 100, bridge_script: str | None = None):
        self._voice = voice
        self._speed = speed
        self._process: asyncio.subprocess.Process | None = None
        self._bridge_script = bridge_script or _resolve_path("tts_bridge.js")
        self._node_exe = _get_node_exe()
        self._stderr_task: asyncio.Task | None = None
        self._io_lock = asyncio.Lock()

    async def start(self) -> None:
        if not pathlib.Path(self._bridge_script).exists():
            raise FileNotFoundError(f"Bridge script not found: {self._bridge_script}")

        logger.info("Starting Node.js TTS bridge (voice=%s, speed=%d)...", self._voice, self._speed)

        self._process = await asyncio.create_subprocess_exec(
            self._node_exe,
            self._bridge_script,
            "--voice", self._voice,
            "--speed", str(self._speed),
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        self._stderr_task = asyncio.create_task(self._drain_stderr())

        line = await asyncio.wait_for(self._process.stdout.readline(), timeout=60)
        line_str = line.decode("utf-8").strip()

        if "READY" not in line_str:
            stderr_data = await self._process.stderr.read()
            raise RuntimeError(
                f"Bridge failed to start. Got: {line_str}. Stderr: {stderr_data.decode('utf-8', errors='replace')}"
            )

        logger.info("Node.js TTS bridge ready.")

    async def _drain_stderr(self) -> None:
        """Continuously read bridge stderr and forward to Python logs."""
        if self._process is None or self._process.stderr is None:
            return
        try:
            while True:
                raw = await self._process.stderr.readline()
                if not raw:
                    break
                line = raw.decode("utf-8", errors="replace").rstrip()
                if line:
                    logger.info("[bridge] %s", line)
        except asyncio.CancelledError:
            pass
        except Exception:
            logger.debug("Bridge stderr drain stopped", exc_info=True)

    async def synthesize(self, kana_text: str) -> bytes:
        if self._process is None or self._process.stdin is None:
            raise RuntimeError("Bridge not started. Call start() first.")

        async with self._io_lock:
            tmpdir = pathlib.Path(tempfile.mkdtemp(prefix="tts_"))
            try:
                input_path = tmpdir / "input.txt"
                output_path = tmpdir / "output.wav"

                input_path.write_text(kana_text, encoding="utf-8")

                cmd_line = f"{input_path}|{output_path}\n"
                self._process.stdin.write(cmd_line.encode("utf-8"))
                await self._process.stdin.drain()

                line = await asyncio.wait_for(self._process.stdout.readline(), timeout=120)
                line_str = line.decode("utf-8").strip()

                if line_str.startswith("ERR:"):
                    raise RuntimeError(f"TTS synthesis error: {line_str[4:]}")

                if not line_str.startswith("OK:") or line_str[3:] != str(output_path):
                    raise RuntimeError(f"Unexpected bridge response: {line_str}")

                wav_data = output_path.read_bytes()
                return wav_data

            finally:
                for f in tmpdir.iterdir():
                    try:
                        f.unlink()
                    except OSError:
                        pass
                try:
                    tmpdir.rmdir()
                except OSError:
                    pass

    async def stop(self) -> None:
        if self._process is not None:
            logger.info("Stopping TTS bridge...")
            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
            try:
                if self._process.stdin is not None:
                    self._process.stdin.close()
            except OSError:
                pass
            try:
                self._process.terminate()
                await asyncio.wait_for(self._process.wait(), timeout=5)
            except asyncio.TimeoutError:
                self._process.kill()
                await self._process.wait()
            logger.info("TTS bridge stopped.")
            self._process = None