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 re
import struct
import threading
import wave
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

import onnxruntime
from piper import PiperConfig, PiperVoice, SynthesisConfig

from hermes_jetson_tts_cues import build_cue_cache, resolve_cached_cue, write_cached_cue


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")))
STREAM_WRITE_BYTES = 32 * 1024
STREAM_WRITE_TIMEOUT_SECONDS = max(
    1.0,
    float(os.getenv("HERMES_TTS_STREAM_WRITE_TIMEOUT_SECONDS", "15")),
)
VOICE_LOCK = threading.Lock()
TURN_ID_PATTERN = re.compile(r"[A-Za-z0-9._:-]{1,128}\Z")

# Piper deliberately leaves a generous tail after sentence punctuation. That
# sounds natural when one synthesis result is played in isolation, but hands-
# free mode queues many independently synthesized clauses and compounds those
# tails with the browser's punctuation cadence. Only silence after the last
# audible 5 ms window is shortened; voiced samples are never faded or removed.
SILENCE_WINDOW_MS = 5
SILENCE_ABS_THRESHOLD = 32
SILENCE_TRIM_TOLERANCE_MS = 20
CLAUSE_PAUSE_MS = 90
COLON_PAUSE_MS = 110
QUESTION_PAUSE_MS = 170
PERIOD_PAUSE_MS = 190
ELLIPSIS_PAUSE_MS = 220
PARAGRAPH_PAUSE_MS = 280
UNPUNCTUATED_PAUSE_MS = 70
SENTENCE_BOUNDARY_PATTERN = re.compile(
    r"(?P<punct>\.{1,3}|[!?…]+)(?:[\"'”’)}\]]+)?(?P<gap>\s+|$)"
)

# 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 normalize_turn_id(value: object) -> str | None:
    """Return a bounded header-safe turn identifier, or discard it."""
    if not isinstance(value, str):
        return None
    normalized = value.strip()
    if not TURN_ID_PATTERN.fullmatch(normalized):
        return None
    return normalized


def trailing_pause_ms(text: str) -> int:
    """Return the natural trailing pause budget for one spoken text unit."""
    if re.search(r"\n\s*\n\s*$", text):
        return PARAGRAPH_PAUSE_MS
    terminal = text.rstrip().rstrip("\"'”’)}]")
    if terminal.endswith(("...", "…")):
        return ELLIPSIS_PAUSE_MS
    if terminal.endswith("."):
        return PERIOD_PAUSE_MS
    if terminal.endswith(("?", "!")):
        return QUESTION_PAUSE_MS
    if terminal.endswith((":", ";")):
        return COLON_PAUSE_MS
    if terminal.endswith(","):
        return CLAUSE_PAUSE_MS
    return UNPUNCTUATED_PAUSE_MS


def sentence_pause_targets(text: str) -> list[int]:
    """Map Piper's likely sentence results to punctuation-aware pause budgets."""
    targets: list[int] = []
    cursor = 0
    for match in SENTENCE_BOUNDARY_PATTERN.finditer(text):
        unit = text[cursor : match.end()]
        targets.append(trailing_pause_ms(unit))
        cursor = match.end()
    if text[cursor:].strip() or not targets:
        targets.append(trailing_pause_ms(text[cursor:] or text))
    return targets


def trim_trailing_pcm_silence(
    pcm: bytes | bytearray | memoryview,
    sample_rate: int,
    pause_ms: int,
) -> memoryview:
    """Shorten only a confirmed PCM silence tail to the requested duration.

    The backward scan is restricted to the trailing envelope. Returning a
    memoryview avoids copying sentence audio on the latency-sensitive stream
    path. All-silent input is retained because it carries no safe phoneme
    boundary from which to measure a pause.
    """
    raw = memoryview(pcm).cast("B")
    if len(raw) < 2 or len(raw) % 2 or sample_rate <= 0:
        return raw

    window_samples = max(1, sample_rate * SILENCE_WINDOW_MS // 1000)
    window_bytes = window_samples * 2
    active_end: int | None = None
    cursor = len(raw)
    while cursor > 0:
        start = max(0, cursor - window_bytes)
        # Piper PCM is explicitly signed 16-bit little endian. A conservative
        # threshold preserves quiet word endings while ignoring quantization
        # noise in the generated tail.
        if any(
            abs(struct.unpack_from("<h", raw, offset)[0]) > SILENCE_ABS_THRESHOLD
            for offset in range(start, cursor, 2)
        ):
            active_end = cursor
            break
        cursor = start

    if active_end is None:
        return raw
    desired_tail_bytes = max(0, sample_rate * pause_ms // 1000) * 2
    cutoff = min(len(raw), active_end + desired_tail_bytes)
    tolerance_bytes = max(0, sample_rate * SILENCE_TRIM_TOLERANCE_MS // 1000) * 2
    if len(raw) - cutoff <= tolerance_bytes:
        return raw
    return raw[:cutoff]


def normalize_wav_trailing_pause(audio: bytes, text: str) -> bytes:
    """Apply the same safe silence-tail policy to the compatibility WAV path."""
    source = io.BytesIO(audio)
    try:
        with wave.open(source, "rb") as reader:
            params = reader.getparams()
            if (
                params.nchannels != 1
                or params.sampwidth != 2
                or params.comptype != "NONE"
            ):
                return audio
            pcm = reader.readframes(params.nframes)
    except (EOFError, wave.Error):
        return audio

    normalized = trim_trailing_pcm_silence(
        pcm,
        params.framerate,
        trailing_pause_ms(text),
    )
    if len(normalized) == len(pcm):
        return audio

    output = io.BytesIO()
    with wave.open(output, "wb") as writer:
        writer.setparams(params)
        writer.writeframes(normalized)
    return output.getvalue()


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"
    protocol_version = "HTTP/1.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",
                "streaming": {
                    "path": "/v1/audio/speech/stream",
                    "format": "pcm_s16le",
                },
                "thinking_cues": len(self.server.cue_cache),  # type: ignore[attr-defined]
            },
        )

    def do_POST(self) -> None:
        if self.path not in {"/v1/audio/speech", "/v1/audio/speech/stream"}:
            _json(self, 404, {"error": "not found"})
            return
        try:
            content_length = int(self.headers.get("Content-Length", "0") or "0")
        except (TypeError, ValueError):
            _json(self, 400, {"error": "invalid Content-Length"})
            return
        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"))
            if not isinstance(payload, dict):
                raise TypeError("request must be an object")
            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))

        try:
            cue = resolve_cached_cue(
                payload,
                getattr(self.server, "cue_cache", {}),
            )
        except ValueError as exc:
            _json(self, 400, {"error": str(exc)})
            return
        if cue is not None:
            if self.path != "/v1/audio/speech/stream":
                _json(self, 400, {"error": "thinking cues require PCM streaming"})
                return
            write_cached_cue(self, cue, normalize_turn_id(payload.get("turn_id")))
            return

        # 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]

        if self.path == "/v1/audio/speech/stream":
            self._stream_pcm(
                text,
                speed,
                voice_name,
                voice,
                normalize_turn_id(payload.get("turn_id")),
            )
            return

        self._synthesize_wav(text, speed, voice_name, voice)

    def _synthesize_wav(self, text: str, speed: float, voice_name: str, voice: PiperVoice) -> None:
        """Return the compatibility WAV response after complete synthesis."""

        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 = normalize_wav_trailing_pause(output.getvalue(), text)
            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 _write_stream_chunk(self, audio: memoryview) -> bool:
        """Write one bounded HTTP chunk, returning false after disconnect."""
        try:
            self.wfile.write(f"{len(audio):X}\r\n".encode("ascii"))
            self.wfile.write(audio)
            self.wfile.write(b"\r\n")
            self.wfile.flush()
            return True
        except (BrokenPipeError, ConnectionResetError, TimeoutError, OSError):
            self.close_connection = True
            return False

    def _stream_pcm(
        self,
        text: str,
        speed: float,
        voice_name: str,
        voice: PiperVoice,
        turn_id: str | None,
    ) -> None:
        """Progressively stream sentence PCM with bounded transport writes.

        Piper yields one PCM result per sentence. Each result is split further
        before writing so the HTTP socket, rather than application buffers,
        provides back-pressure. Aborting the browser fetch closes the socket;
        the server then drops the current turn at the next synthesis yield or
        bounded write and releases the voice lock.
        """
        sample_rate = int(voice.config.sample_rate)
        self.send_response(200)
        self.send_header(
            "Content-Type",
            f"audio/pcm;rate={sample_rate};channels=1;encoding=signed-integer;bits=16;endian=little",
        )
        self.send_header("Transfer-Encoding", "chunked")
        self.send_header("Cache-Control", "no-store")
        self.send_header("X-Audio-Format", "pcm_s16le")
        self.send_header("X-Audio-Sample-Rate", str(sample_rate))
        self.send_header("X-Audio-Channels", "1")
        self.send_header("X-Audio-Sample-Width", "2")
        self.send_header("X-TTS-Voice", voice_name)
        if turn_id is not None:
            self.send_header("X-TTS-Turn-ID", turn_id)
        self.end_headers()

        try:
            self.connection.settimeout(STREAM_WRITE_TIMEOUT_SECONDS)
        except (AttributeError, OSError):
            # Unit-test handlers and already-closed clients may not expose a
            # live socket. The writes below remain the cancellation boundary.
            pass

        connected = True
        try:
            pause_targets = sentence_pause_targets(text)
            audio_chunks = iter(
                voice.synthesize(
                    text,
                    SynthesisConfig(length_scale=1.0 / speed),
                )
            )
            chunk_index = 0
            while connected:
                # Serialize Piper/ONNX access, but never hold the model lock
                # while a slow browser applies network back-pressure.
                with VOICE_LOCK:
                    try:
                        audio_chunk = next(audio_chunks)
                    except StopIteration:
                        break
                if (
                    audio_chunk.sample_rate != sample_rate
                    or audio_chunk.sample_width != 2
                    or audio_chunk.sample_channels != 1
                ):
                    raise RuntimeError("Piper returned an unexpected PCM format")
                pause_ms = pause_targets[min(chunk_index, len(pause_targets) - 1)]
                pcm = trim_trailing_pcm_silence(
                    audio_chunk.audio_int16_bytes,
                    sample_rate,
                    pause_ms,
                )
                chunk_index += 1
                for offset in range(0, len(pcm), STREAM_WRITE_BYTES):
                    if not self._write_stream_chunk(pcm[offset : offset + STREAM_WRITE_BYTES]):
                        connected = False
                        break
        except Exception as exc:
            # Headers are already committed, so a JSON error would corrupt the
            # PCM stream. Closing produces an incomplete chunked response that
            # the browser can discard or replace with the compatibility path.
            print(f"[tts] streaming synthesis failed: {exc}", flush=True)
            self.close_connection = True
            return

        if connected:
            try:
                self.wfile.write(b"0\r\n\r\n")
                self.wfile.flush()
            except (BrokenPipeError, ConnectionResetError, TimeoutError, OSError):
                self.close_connection = True


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)
    cue_cache = build_cue_cache(
        voices,
        SynthesisConfig,
        trim_trailing_pcm_silence,
        trailing_pause_ms,
    )
    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]
    server.cue_cache = cue_cache  # type: ignore[attr-defined]
    print(f"[tts] ready on {HOST}:{PORT}", flush=True)
    server.serve_forever(poll_interval=0.25)


if __name__ == "__main__":
    main()