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

124 lines
4.6 KiB
Python

#!/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"))
VOICE_NAME = os.getenv("HERMES_TTS_VOICE", "en_US-lessac-high")
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()
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, "voice": VOICE_NAME, "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))
output = io.BytesIO()
try:
with VOICE_LOCK, wave.open(output, "wb") as wav_file:
self.server.voice.synthesize_wav( # type: ignore[attr-defined]
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.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 main() -> None:
"""Load the checksum-pinned voice from the image and serve it on CPU."""
CACHE_DIR.mkdir(parents=True, exist_ok=True)
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 = ONNX_THREADS
session_options.inter_op_num_threads = 1
session = onnxruntime.InferenceSession(
str(model_path),
sess_options=session_options,
providers=["CPUExecutionProvider"],
)
voice = PiperVoice(session=session, config=config, download_dir=CACHE_DIR)
print(
f"[tts] loaded Piper voice {VOICE_NAME} on CPU with {ONNX_THREADS} ONNX threads",
flush=True,
)
server = ThreadingHTTPServer((HOST, PORT), SpeechHandler)
server.voice = voice # type: ignore[attr-defined]
print(f"[tts] ready on {HOST}:{PORT}", flush=True)
server.serve_forever(poll_interval=0.25)
if __name__ == "__main__":
main()