497 lines
19 KiB
Python
497 lines
19 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 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()
|