atlas-iac/dockerfiles/hermes-jetson-tts-server.py

#!/usr/bin/env python3
"""CPU-only Piper service shared with the routing Jetson."""

from __future__ import annotations

import io
import json
import os
import threading
import wave
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

import onnxruntime
from piper import PiperConfig, PiperVoice, SynthesisConfig


HOST = os.getenv("HERMES_TTS_HOST", "0.0.0.0")
PORT = int(os.getenv("HERMES_TTS_PORT", "9001"))
CACHE_DIR = Path(os.getenv("HERMES_TTS_CACHE", "/cache/piper"))
MAX_TEXT_CHARS = 5000
ONNX_THREADS = max(1, int(os.getenv("HERMES_TTS_ONNX_THREADS", "4")))
VOICE_LOCK = threading.Lock()

# Fixed, allow-listed language -> baked voice mapping. This is the ONLY path
# from a client-supplied string to a model name: client input is looked up
# here and never used to build a filesystem path directly. Both "-" and "_"
# separators and any case are accepted; anything not present here falls back
# to DEFAULT_VOICE_NAME (safe English default), never an error and never an
# unbaked model.
LANGUAGE_VOICE_MAP = {
    "en": "en_US-amy-medium",
    "en-us": "en_US-amy-medium",
    "ru": "ru_RU-irina-medium",
    "ru-ru": "ru_RU-irina-medium",
    "es": "es_MX-claude-high",
    "es-mx": "es_MX-claude-high",
    "es-es": "es_MX-claude-high",
}
BAKED_VOICE_NAMES = frozenset(LANGUAGE_VOICE_MAP.values())
DEFAULT_VOICE_NAME = os.getenv("HERMES_TTS_VOICE", "en_US-amy-medium")


def normalize_language(value: object) -> str | None:
    """Lowercase and fold "_"/"-" separators; reject non-string/blank input."""
    if not isinstance(value, str):
        return None
    normalized = value.strip().lower().replace("_", "-")
    return normalized or None


def resolve_voice_name(language: object) -> str:
    """Map a client-supplied language to one of the baked policy voices.

    Unknown, missing, or malformed language always resolves to the safe
    default rather than raising, and the result is always a member of
    BAKED_VOICE_NAMES.
    """
    normalized = normalize_language(language)
    if normalized is None:
        return DEFAULT_VOICE_NAME
    return LANGUAGE_VOICE_MAP.get(normalized, DEFAULT_VOICE_NAME)


def _json(handler: BaseHTTPRequestHandler, status: int, payload: dict) -> None:
    body = json.dumps(payload).encode("utf-8")
    handler.send_response(status)
    handler.send_header("Content-Type", "application/json")
    handler.send_header("Content-Length", str(len(body)))
    handler.send_header("Cache-Control", "no-store")
    handler.end_headers()
    handler.wfile.write(body)


class SpeechHandler(BaseHTTPRequestHandler):
    """Serve health and bounded local speech synthesis."""

    server_version = "AtlasPiper/1"

    def log_message(self, message: str, *args: object) -> None:
        print(f"[tts] {self.address_string()} {message % args}", flush=True)

    def do_GET(self) -> None:
        if self.path != "/health":
            _json(self, 404, {"error": "not found"})
            return
        _json(
            self,
            200,
            {
                "ok": True,
                "voices": sorted(self.server.voices),  # type: ignore[attr-defined]
                "default_voice": self.server.default_voice_name,  # type: ignore[attr-defined]
                "device": "cpu",
            },
        )

    def do_POST(self) -> None:
        if self.path != "/v1/audio/speech":
            _json(self, 404, {"error": "not found"})
            return
        content_length = int(self.headers.get("Content-Length", "0") or "0")
        if content_length <= 0 or content_length > 64 * 1024:
            _json(self, 413, {"error": "request is missing or too large"})
            return
        try:
            payload = json.loads(self.rfile.read(content_length).decode("utf-8"))
            text = str(payload.get("input") or payload.get("text") or "").strip()
            speed = float(payload.get("speed") or 1.0)
        except (TypeError, ValueError, json.JSONDecodeError):
            _json(self, 400, {"error": "invalid JSON request"})
            return
        if not text:
            _json(self, 400, {"error": "input is required"})
            return
        if len(text) > MAX_TEXT_CHARS:
            _json(self, 400, {"error": "input is too long"})
            return
        speed = min(2.0, max(0.5, speed))

        # Policy is driven ONLY by "language". A client-supplied "voice"
        # field is deliberately never read here; it cannot override the
        # allow-listed mapping.
        voice_name = resolve_voice_name(payload.get("language"))
        voice = self.server.voices[voice_name]  # type: ignore[attr-defined]

        output = io.BytesIO()
        try:
            with VOICE_LOCK, wave.open(output, "wb") as wav_file:
                voice.synthesize_wav(
                    text,
                    wav_file,
                    SynthesisConfig(length_scale=1.0 / speed),
                )
            audio = output.getvalue()
            self.send_response(200)
            self.send_header("Content-Type", "audio/wav")
            self.send_header("Content-Length", str(len(audio)))
            self.send_header("Cache-Control", "no-store")
            self.send_header("X-TTS-Voice", voice_name)
            self.end_headers()
            self.wfile.write(audio)
        except Exception as exc:
            print(f"[tts] synthesis failed: {exc}", flush=True)
            _json(self, 500, {"error": "speech synthesis failed"})


def _load_voice(cache_dir: Path, voice_name: str, threads: int) -> PiperVoice:
    model_path = cache_dir / f"{voice_name}.onnx"
    config_path = cache_dir / f"{voice_name}.onnx.json"
    if not model_path.exists() or not config_path.exists():
        raise RuntimeError(f"baked Piper voice is missing: {voice_name}")
    with config_path.open("r", encoding="utf-8") as config_file:
        config = PiperConfig.from_dict(json.load(config_file))
    session_options = onnxruntime.SessionOptions()
    session_options.intra_op_num_threads = threads
    session_options.inter_op_num_threads = 1
    session = onnxruntime.InferenceSession(
        str(model_path),
        sess_options=session_options,
        providers=["CPUExecutionProvider"],
    )
    return PiperVoice(session=session, config=config, download_dir=cache_dir)


def load_voices(cache_dir: Path, threads: int) -> dict[str, PiperVoice]:
    """Eagerly load all three policy voices.

    Preload (not lazy-load-on-first-use) was chosen deliberately: measured
    RSS on this model set is ~88MB for one voice and ~243MB for all three
    (~+155MB versus the previous single-voice baseline), which comfortably
    fits the pod's memory budget on the CPU-only voice node. Preloading
    avoids a slow, request-serializing first synthesis per language and
    keeps the fail-closed missing-model check (below) at process start
    rather than deferring a possible crash to a live user request.
    """
    return {name: _load_voice(cache_dir, name, threads) for name in sorted(BAKED_VOICE_NAMES)}


def main() -> None:
    """Load the checksum-pinned policy voices from the image and serve them on CPU."""
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    if DEFAULT_VOICE_NAME not in BAKED_VOICE_NAMES:
        raise RuntimeError(
            f"HERMES_TTS_VOICE must name one of the baked policy voices: {sorted(BAKED_VOICE_NAMES)}"
        )
    voices = load_voices(CACHE_DIR, ONNX_THREADS)
    print(
        f"[tts] loaded {len(voices)} Piper voices on CPU with {ONNX_THREADS} ONNX threads each: "
        + ", ".join(sorted(voices)),
        flush=True,
    )
    server = ThreadingHTTPServer((HOST, PORT), SpeechHandler)
    server.voices = voices  # type: ignore[attr-defined]
    server.default_voice_name = DEFAULT_VOICE_NAME  # type: ignore[attr-defined]
    print(f"[tts] ready on {HOST}:{PORT}", flush=True)
    server.serve_forever(poll_interval=0.25)


if __name__ == "__main__":
    main()