diff --git a/dockerfiles/hermes-jetson-tts-server.py b/dockerfiles/hermes-jetson-tts-server.py index c5ed37799..3abdc28b0 100644 --- a/dockerfiles/hermes-jetson-tts-server.py +++ b/dockerfiles/hermes-jetson-tts-server.py @@ -11,7 +11,8 @@ import wave from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path -from piper import PiperVoice, SynthesisConfig +import onnxruntime +from piper import PiperConfig, PiperVoice, SynthesisConfig HOST = os.getenv("HERMES_TTS_HOST", "0.0.0.0") @@ -19,6 +20,7 @@ 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() @@ -96,8 +98,21 @@ def main() -> None: 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}") - print(f"[tts] loading Piper voice {VOICE_NAME} on CPU", flush=True) - voice = PiperVoice.load(model_path, config_path, use_cuda=False, download_dir=CACHE_DIR) + 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) diff --git a/services/hermes/voice-deployment.yaml b/services/hermes/voice-deployment.yaml index e1f8bf2e0..778685d59 100644 --- a/services/hermes/voice-deployment.yaml +++ b/services/hermes/voice-deployment.yaml @@ -123,7 +123,7 @@ spec: kubernetes.io/hostname: titan-20 containers: - name: tts - image: registry.bstein.dev/bstein/hermes-jetson-tts@sha256:1020a73e8941932dd071e5dcc66b5583ebb9376f9ee2fd1711b50e026eb88467 + image: registry.bstein.dev/bstein/hermes-jetson-tts@sha256:82fbed67bd871e821e4fd745a49c91a3749dc6ad45ea4841dc3fb7cfbcd53e61 imagePullPolicy: IfNotPresent ports: - {name: http, containerPort: 9001, protocol: TCP} @@ -133,6 +133,7 @@ spec: - {name: HERMES_TTS_PORT, value: "9001"} - {name: HERMES_TTS_VOICE, value: en_US-lessac-high} - {name: HERMES_TTS_CACHE, value: /opt/models/piper} + - {name: HERMES_TTS_ONNX_THREADS, value: "4"} startupProbe: httpGet: {path: /health, port: http} periodSeconds: 5 @@ -159,8 +160,8 @@ spec: volumeMounts: - {name: tmp, mountPath: /tmp} resources: - requests: {cpu: 500m, memory: 512Mi} - limits: {cpu: "2", memory: 2Gi} + requests: {cpu: "1", memory: 512Mi} + limits: {cpu: "4", memory: 2Gi} volumes: - name: tmp emptyDir: diff --git a/testing/tests/test_hermes_chat_quality.py b/testing/tests/test_hermes_chat_quality.py index e0259c8dd..849a615ab 100644 --- a/testing/tests/test_hermes_chat_quality.py +++ b/testing/tests/test_hermes_chat_quality.py @@ -200,6 +200,7 @@ def test_voice_models_are_baked_and_runtime_has_no_public_egress(): tts_server = (ROOT / "dockerfiles" / "hermes-jetson-tts-server.py").read_text() assert "download_voice" not in tts_server assert "baked Piper voice is missing" in tts_server + assert "session_options.intra_op_num_threads = ONNX_THREADS" in tts_server policies = _documents(HERMES / "networkpolicy.yaml") voice_policy = next( @@ -237,6 +238,11 @@ def test_voice_workloads_have_deliberate_xavier_placement(): assert stt_resources["requests"]["nvidia.com/gpu.shared"] == 1 assert stt_resources["limits"]["nvidia.com/gpu.shared"] == 1 assert "nvidia.com/gpu.shared" not in tts["containers"][0]["resources"]["requests"] + tts_env = { + item["name"]: item["value"] for item in tts["containers"][0]["env"] + } + assert tts_env["HERMES_TTS_ONNX_THREADS"] == "4" + assert tts["containers"][0]["resources"]["limits"]["cpu"] == "4" assert all("hostPath" not in volume for volume in stt["volumes"]) assert all("hostPath" not in volume for volume in tts["volumes"])