Compare commits
10
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f270174a35 | ||
|
|
41bf35b444 | ||
|
|
a89e2b8092 | ||
|
|
6be0d0cc0d | ||
|
|
94fbce1adb | ||
|
|
4d180db85f | ||
|
|
60f9a46d31 | ||
|
|
b1d2290a45 | ||
|
|
ee418255ed | ||
|
|
f4009c3040 |
@@ -0,0 +1,111 @@
|
||||
"""pockettts: Kyutai PocketTTS as a crash-isolated, CPU-only TTS sidecar (#1306).
|
||||
|
||||
PocketTTS (kyutai-labs/pocket-tts, 100M params) is hired for the "fastest CPU
|
||||
render / lowest latency" job, the row the engine-acceptance framework leaves
|
||||
unheld: every CPU engine OmniVoice ships is either English-only or a quality
|
||||
engine falling back to CPU. PocketTTS is the complementary opposite end of the
|
||||
spectrum from the quality engines (omnivoice, IndexTTS, Supertonic-3): small,
|
||||
fast, CPU-only, zero-shot cloning from a reference clip. Six languages
|
||||
(en/fr/de/pt/it/es), one model per language, selected via the ``language``
|
||||
kwarg. Measured ~8-9x real-time on an Apple M3 Pro (see
|
||||
scripts/bench_engines_latency.py, PR #1322).
|
||||
|
||||
This engine runs PocketTTS in a child process via :class:`SubprocessBackend`,
|
||||
mirroring engines/omnivoice_subprocess and engines/supertonic3. Crash isolation:
|
||||
a wedged generate is hard-killed by the parent's watchdog, reclaiming the
|
||||
child's memory, the thing an in-process engine structurally cannot do. CPU-only
|
||||
by design (``gpu_compat = ("cpu",)``): Kyutai observes no GPU speedup for this
|
||||
100M, batch-1 model.
|
||||
|
||||
Opt-in (Settings -> Engines, or ``OMNIVOICE_TTS_BACKEND=pockettts``); the
|
||||
default ``omnivoice`` engine is unchanged, so existing users see no behaviour
|
||||
change.
|
||||
|
||||
Licence: MIT (code) + CC-BY-4.0 (weights), both commercial-OK (cleared from
|
||||
primary sources in #1306). The weights are gated on HuggingFace (an access
|
||||
agreement plus an acceptable-use clause); the engine must surface that honestly
|
||||
at first-run rather than failing inside a download (condition 6 of the #1306
|
||||
acceptance). That preflight is built on top of this shape, not in it.
|
||||
|
||||
Streaming note: PocketTTS streams audio (``generate_audio_stream``), but this
|
||||
batch sidecar returns one audio frame per synth, matching the SubprocessBackend
|
||||
contract every other subprocess engine uses. A streaming-aware variant
|
||||
(incremental audio frames) is a documented opportunity to recover PocketTTS's
|
||||
~33 ms time-to-first-audio end-to-end; out of scope for this shape, raised on
|
||||
the PR.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from services.subprocess_backend import SubprocessBackend
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import torch # noqa: F401
|
||||
|
||||
|
||||
class PocketTTSBackend(SubprocessBackend):
|
||||
"""Kyutai PocketTTS in a killable, CPU-only sidecar process."""
|
||||
|
||||
id = "pockettts"
|
||||
display_name = "PocketTTS (Kyutai, 6 langs, CPU-only, MIT/CC-BY-4.0)"
|
||||
_DEFAULT_SAMPLE_RATE = 24_000
|
||||
# CPU-only by design (honest hardware reporting, like supertonic3): Kyutai
|
||||
# ships no CUDA/MPS path and reports no GPU speedup for this model.
|
||||
gpu_compat: tuple[str, ...] = ("cpu",)
|
||||
supports_cloning = True # zero-shot clone from a reference clip
|
||||
|
||||
@classmethod
|
||||
def is_available(cls) -> tuple[bool, str]:
|
||||
# Optional-dep gate: the pocket-tts wheel is installed only when the user
|
||||
# opted in. The interpreter is the parent's own (sys.executable), so
|
||||
# there is no separate venv to validate.
|
||||
try:
|
||||
import pocket_tts # type: ignore[import-not-found] # noqa: F401
|
||||
except Exception as e:
|
||||
return False, (
|
||||
f"pocket_tts package not installed or failed to import ({e}). "
|
||||
f"Enable in Settings -> Engines (pip install pocket-tts)."
|
||||
)
|
||||
return True, "ready (CPU-only)"
|
||||
|
||||
@classmethod
|
||||
def venv_python(cls) -> Path:
|
||||
# Parent interpreter: pocket-tts deps (torch>=2.5, scipy, beartype) sit
|
||||
# happily at the parent's pins, so this isolates for crash recovery, not
|
||||
# dependency pins (same rationale as omnivoice-subprocess).
|
||||
return Path(sys.executable)
|
||||
|
||||
@classmethod
|
||||
def sidecar_script(cls) -> Path:
|
||||
return Path(__file__).resolve().parent / "main.py"
|
||||
|
||||
@property
|
||||
def recv_timeout_s(self) -> float:
|
||||
# A cold load pulls gated weights (a 24-layer model can be hundreds of MB),
|
||||
# so allow a long recv deadline; the sidecar also heartbeats progress frames
|
||||
# during the download (main.py) to keep the watchdog armed.
|
||||
try:
|
||||
v = float(os.environ.get("OMNIVOICE_POCKETTTS_RECV_TIMEOUT_S", "600"))
|
||||
except (ValueError, TypeError):
|
||||
return 600.0
|
||||
if not math.isfinite(v): # reject inf/nan so the deadline can't be disabled
|
||||
return 600.0
|
||||
return max(30.0, v)
|
||||
|
||||
@property
|
||||
def sample_rate(self) -> int:
|
||||
return self._DEFAULT_SAMPLE_RATE
|
||||
|
||||
@property
|
||||
def supported_languages(self) -> list[str]:
|
||||
# Protocol tag; six languages (en/fr/de/pt/it/es), one model per
|
||||
# language, selected via the language kwarg.
|
||||
return ["multi"]
|
||||
|
||||
|
||||
__all__ = ["PocketTTSBackend"]
|
||||
@@ -0,0 +1,310 @@
|
||||
"""pockettts sidecar entry point (#1306).
|
||||
|
||||
Runs Kyutai PocketTTS in a child process under the parent's own interpreter
|
||||
(same pins), so a wedged generate can be hard-killed by the parent to reclaim
|
||||
memory. Mirrors engines/omnivoice_subprocess/main.py.
|
||||
|
||||
Wire protocol: length-prefixed JSON over stdin/stdout, byte-identical to
|
||||
services/subprocess_backend.py::
|
||||
|
||||
[ 4-byte big-endian uint32 length ][ N bytes UTF-8 JSON ]
|
||||
|
||||
Op flow:
|
||||
1. sidecar -> parent: {"op":"ready","engine":"pockettts","sample_rate":24000}
|
||||
2. parent -> sidecar: {"op":"ping"} -> {"op":"pong","vram_mb":0}
|
||||
3. parent -> sidecar: {"op":"synthesize","text":"...",
|
||||
"ref_audio":"/path/to/ref.wav",
|
||||
"language":"it"}
|
||||
-> {"op":"progress",...} (cold load) then
|
||||
-> {"op":"audio","audio_pcm_b64":"...","sample_rate":24000,
|
||||
"n_samples":N}
|
||||
4. parent -> sidecar: {"op":"shutdown"} -> exit 0
|
||||
|
||||
Stdlib-only at import time; torch + pocket_tts are imported lazily on the first
|
||||
synthesize so the ready frame fits the parent's 30s spawn handshake even on a
|
||||
cold filesystem.
|
||||
|
||||
Languages: PocketTTS ships one model per language (en/fr/de/pt/it/es), selected
|
||||
by ``language``. The first synth in a given language cold-loads + caches that
|
||||
model; later calls reuse it. (The HF model card's "English only at the moment"
|
||||
line is stale; the GitHub README and pocket-tts 2.1.0 confirm six languages.)
|
||||
|
||||
Note: ``TTSModel.load_model(language=...)`` pulls the gated kyutai weights from
|
||||
HuggingFace, so it needs HF auth + the access agreement accepted. A failure here
|
||||
currently surfaces as a raw error frame; the typed "weights are gated" preflight
|
||||
(condition 6) is built on top of this shape, not in it.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import struct
|
||||
import sys
|
||||
import threading
|
||||
import traceback
|
||||
from collections import OrderedDict
|
||||
|
||||
# Mirrors services/subprocess_backend.py::MAX_FRAME_BYTES.
|
||||
MAX_FRAME_BYTES = 64 * 1024 * 1024
|
||||
|
||||
#: PocketTTS emits 24 kHz mono. Re-read from the loaded model on each generate.
|
||||
POCKETTTS_SAMPLE_RATE = 24_000
|
||||
|
||||
#: OmniVoice language (ISO code, name, or sentinel) -> pocket-tts model language.
|
||||
#: "auto"/"multi"/"na"/None default to english (the library default).
|
||||
_LANG_MAP = {
|
||||
"en": "english", "eng": "english", "english": "english",
|
||||
"fr": "french", "fra": "french", "french": "french",
|
||||
"de": "german", "deu": "german", "german": "german",
|
||||
"pt": "portuguese", "por": "portuguese", "portuguese": "portuguese",
|
||||
"it": "italian", "ita": "italian", "italian": "italian",
|
||||
"es": "spanish", "esp": "spanish", "spanish": "spanish",
|
||||
}
|
||||
|
||||
#: Default preset voice per language when no reference clip is supplied (public
|
||||
#: presets from kyutai/tts-voices; voice source does not affect synth speed).
|
||||
_DEFAULT_VOICE_BY_LANG = {
|
||||
"english": "alba",
|
||||
"italian": "giovanni",
|
||||
"spanish": "lola",
|
||||
"german": "juergen",
|
||||
"portuguese": "rafael",
|
||||
"french": "estelle",
|
||||
}
|
||||
|
||||
#: Emit a progress frame at least this often during a cold load so the parent's
|
||||
#: recv watchdog doesn't kill a healthy sidecar on a slow first download.
|
||||
_HEARTBEAT_S = 5.0
|
||||
|
||||
#: ref_audio must be a local file path, not a URL (local-first; no SSRF).
|
||||
_URL_RE = re.compile(r"^[a-z][a-z0-9+.\-]*://", re.IGNORECASE)
|
||||
|
||||
#: Bound the per-(language, voice) voice-state cache (LRU) so a long session
|
||||
#: with many distinct reference clips can't grow memory without limit.
|
||||
_VOICE_CACHE_MAX = 8
|
||||
|
||||
# Per-language model cache: load_model(language=...) is slow and PocketTTS ships
|
||||
# one model per language, so cache each. Bounded by the distinct languages used
|
||||
# in a session (at most six).
|
||||
_MODELS: dict[str, object] = {}
|
||||
|
||||
# (language, voice) -> voice_state, LRU-bounded to _VOICE_CACHE_MAX entries.
|
||||
# get_state_for_audio_prompt is relatively slow, so cache per (language, voice)
|
||||
# to avoid re-encoding on every call.
|
||||
_voice_cache: OrderedDict[str, object] = OrderedDict()
|
||||
|
||||
|
||||
# -- wire protocol -----------------------------------------------------------
|
||||
|
||||
#: Serializes _send across threads (the cold-load heartbeat + the main loop) so
|
||||
#: concurrent length+body writes can't interleave and corrupt the framing.
|
||||
_send_lock = threading.Lock()
|
||||
|
||||
|
||||
def _send(stream, obj: dict) -> None:
|
||||
body = json.dumps(obj, separators=(",", ":")).encode("utf-8")
|
||||
with _send_lock:
|
||||
stream.write(struct.pack("!I", len(body)))
|
||||
stream.write(body)
|
||||
stream.flush()
|
||||
|
||||
|
||||
def _recv(stream):
|
||||
header = stream.read(4)
|
||||
if len(header) < 4:
|
||||
return None # EOF
|
||||
(n,) = struct.unpack("!I", header)
|
||||
if n > MAX_FRAME_BYTES:
|
||||
raise IOError(f"frame too large: {n}")
|
||||
body = bytearray()
|
||||
while len(body) < n:
|
||||
chunk = stream.read(n - len(body))
|
||||
if not chunk:
|
||||
raise IOError("short read")
|
||||
body.extend(chunk)
|
||||
return json.loads(bytes(body).decode("utf-8"))
|
||||
|
||||
|
||||
def _measure_vram_mb() -> float:
|
||||
"""CPU-only engine: always 0. Kept for protocol parity with the parent."""
|
||||
return 0.0
|
||||
|
||||
|
||||
# -- model loading (lazy, on first synthesize per language) ------------------
|
||||
|
||||
|
||||
def _pocket_language(raw) -> str:
|
||||
"""Map an OmniVoice language value to a pocket-tts model language. A specific
|
||||
but unsupported language raises rather than silently fall back to English and
|
||||
mispronounce; empty / "auto" / "multi" / "na" default to English."""
|
||||
if not raw:
|
||||
return "english"
|
||||
s = str(raw).strip().lower()
|
||||
if s in ("", "auto", "multi", "na"):
|
||||
return "english"
|
||||
if s in _LANG_MAP:
|
||||
return _LANG_MAP[s]
|
||||
raise ValueError(
|
||||
f"PocketTTS does not support language {raw!r}; supported: en, fr, de, pt, it, es."
|
||||
)
|
||||
|
||||
|
||||
def _load_model(stdout, language: str):
|
||||
"""Cold-construct the PocketTTS model for ``language`` (cached per language).
|
||||
Emits progress frames for the parent watchdog. Raises on failure (e.g.
|
||||
gated-weights access without HF auth); the caller emits an error frame and
|
||||
stays alive for a retry."""
|
||||
model = _MODELS.get(language)
|
||||
if model is not None:
|
||||
return model
|
||||
|
||||
_send(stdout, {"op": "progress", "stage": "loading_model", "percent": 0})
|
||||
|
||||
# Heartbeat: a cold load (gated weights download) can outlast the parent's
|
||||
# recv timeout. Emit a progress frame every few seconds while it runs so the
|
||||
# parent's watchdog sees activity and does not kill a healthy sidecar.
|
||||
stop = threading.Event()
|
||||
|
||||
def _heartbeat() -> None:
|
||||
pct = 1
|
||||
while not stop.wait(_HEARTBEAT_S):
|
||||
pct = min(pct + 1, 99)
|
||||
_send(stdout, {"op": "progress", "stage": "loading_model", "percent": pct})
|
||||
|
||||
hb = threading.Thread(target=_heartbeat, daemon=True)
|
||||
hb.start()
|
||||
try:
|
||||
from pocket_tts import TTSModel # type: ignore[import-not-found] # noqa: PLC0415
|
||||
|
||||
model = TTSModel.load_model(language=language)
|
||||
_MODELS[language] = model
|
||||
finally:
|
||||
stop.set()
|
||||
hb.join(timeout=_HEARTBEAT_S + 1)
|
||||
_send(stdout, {"op": "progress", "stage": "loading_model", "percent": 100})
|
||||
return model
|
||||
|
||||
|
||||
def _voice_state(model, language: str, ref_audio):
|
||||
"""Return a (cached, LRU-bounded) voice state for ``ref_audio`` (a local
|
||||
file path) or the language's default preset voice when none is given. URLs
|
||||
are rejected to keep the sidecar local-first (no SSRF)."""
|
||||
if ref_audio and _URL_RE.match(ref_audio):
|
||||
raise ValueError(
|
||||
"ref_audio must be a local file path; URLs are not accepted (local-first)."
|
||||
)
|
||||
voice = ref_audio or _DEFAULT_VOICE_BY_LANG.get(language, "alba")
|
||||
# For a local file ref, fold mtime+size into the cache key so a file replaced
|
||||
# at the same path does not return a stale voice from the previous contents.
|
||||
fingerprint = ""
|
||||
if ref_audio:
|
||||
try:
|
||||
st = os.stat(ref_audio)
|
||||
fingerprint = f"|m{st.st_mtime_ns}s{st.st_size}"
|
||||
except OSError:
|
||||
fingerprint = ""
|
||||
key = f"{language}|{voice}{fingerprint}"
|
||||
state = _voice_cache.get(key)
|
||||
if state is not None:
|
||||
_voice_cache.move_to_end(key)
|
||||
return state
|
||||
state = model.get_state_for_audio_prompt(voice)
|
||||
_voice_cache[key] = state
|
||||
if len(_voice_cache) > _VOICE_CACHE_MAX:
|
||||
_voice_cache.popitem(last=False) # evict oldest
|
||||
return state
|
||||
|
||||
|
||||
def _tensor_to_pcm_b64(audio, sample_rate: int) -> tuple[str, int, int]:
|
||||
"""Convert a float waveform in [-1, 1] to base64 int16 PCM."""
|
||||
import numpy as np
|
||||
|
||||
arr = np.asarray(audio, dtype=np.float32).squeeze()
|
||||
if arr.ndim > 1:
|
||||
raise ValueError(
|
||||
f"expected mono audio (1-D after squeeze), got shape {arr.shape}; "
|
||||
f"PocketTTS returns mono, so a multi-channel array means an upstream change."
|
||||
)
|
||||
arr = np.clip(arr, -1.0, 1.0)
|
||||
pcm = (arr * 32767.0).astype(np.int16).tobytes()
|
||||
return base64.b64encode(pcm).decode("ascii"), int(sample_rate), int(arr.shape[-1])
|
||||
|
||||
|
||||
def _handle_synthesize(msg: dict, stdout) -> None:
|
||||
"""Dispatch one synthesize request. Emits the audio frame or raises."""
|
||||
text = msg.get("text")
|
||||
if not text or not isinstance(text, str):
|
||||
raise ValueError("synthesize: missing or non-string 'text'")
|
||||
|
||||
language = _pocket_language(msg.get("language"))
|
||||
model = _load_model(stdout, language)
|
||||
ref_audio = msg.get("ref_audio") or None
|
||||
voice_state = _voice_state(model, language, ref_audio)
|
||||
|
||||
audio = model.generate_audio(voice_state, text)
|
||||
sample_rate = int(getattr(model, "sample_rate", POCKETTTS_SAMPLE_RATE))
|
||||
|
||||
pcm_b64, sr, n_samples = _tensor_to_pcm_b64(audio, sample_rate)
|
||||
_send(stdout, {
|
||||
"op": "audio",
|
||||
"audio_pcm_b64": pcm_b64,
|
||||
"sample_rate": sr,
|
||||
"n_samples": n_samples,
|
||||
})
|
||||
|
||||
|
||||
# -- main loop ---------------------------------------------------------------
|
||||
|
||||
|
||||
def main() -> int:
|
||||
stdin = sys.stdin.buffer
|
||||
stdout = sys.stdout.buffer
|
||||
|
||||
# Ready handshake fires BEFORE any heavy import.
|
||||
_send(stdout, {
|
||||
"op": "ready",
|
||||
"engine": "pockettts",
|
||||
"sample_rate": POCKETTTS_SAMPLE_RATE,
|
||||
})
|
||||
|
||||
while True:
|
||||
try:
|
||||
msg = _recv(stdin)
|
||||
except Exception as exc:
|
||||
_send(stdout, {
|
||||
"op": "error",
|
||||
"stage": "recv",
|
||||
"message": f"{type(exc).__name__}: {exc}",
|
||||
"traceback": traceback.format_exc(),
|
||||
})
|
||||
return 1
|
||||
if msg is None:
|
||||
return 0
|
||||
|
||||
op = msg.get("op") if isinstance(msg, dict) else None
|
||||
try:
|
||||
if op == "ping":
|
||||
_send(stdout, {"op": "pong", "vram_mb": _measure_vram_mb()})
|
||||
elif op == "synthesize":
|
||||
_handle_synthesize(msg, stdout)
|
||||
elif op == "shutdown":
|
||||
return 0
|
||||
else:
|
||||
_send(stdout, {
|
||||
"op": "error",
|
||||
"stage": "dispatch",
|
||||
"message": f"unknown op: {op!r}",
|
||||
})
|
||||
except Exception as exc:
|
||||
_send(stdout, {
|
||||
"op": "error",
|
||||
"stage": op or "unknown",
|
||||
"message": f"{type(exc).__name__}: {exc}",
|
||||
"traceback": traceback.format_exc(),
|
||||
})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -1902,6 +1902,11 @@ _LAZY_REGISTRY: dict[str, tuple[str, str]] = {
|
||||
# engine stays the default). Unlike the entries above it runs under the
|
||||
# parent interpreter (crash isolation, not dependency isolation).
|
||||
"omnivoice-subprocess": ("engines.omnivoice_subprocess", "OmniVoiceSubprocessBackend"),
|
||||
# Issue #1306: Kyutai PocketTTS, CPU-only, low-latency TTS hired for the
|
||||
# "fastest CPU render / lowest latency" job. Opt-in, subprocess-isolated
|
||||
# under the parent interpreter (crash isolation, not dependency isolation,
|
||||
# same as omnivoice-subprocess: pocket-tts deps sit at the parent's pins).
|
||||
"pockettts": ("engines.pockettts", "PocketTTSBackend"),
|
||||
# Issue #590: Confucius4-TTS (netease-youdao) — LLM-based, 14-language
|
||||
# cross-lingual zero-shot cloning, Apache-2.0. Opt-in + subprocess-isolated
|
||||
# (own Python 3.10 venv) like the entries above. Validated end-to-end
|
||||
@@ -2008,6 +2013,7 @@ _INSTALL_HINTS: dict[str, str] = {
|
||||
"sherpa-onnx": "pip install sherpa-onnx (universal ONNX runtime, WASM-ready)",
|
||||
"omnivoice-gguf":"Bundled — runs the C++ omnivoice-tts binary in bin/. Quants download lazily from Serveurperso/OmniVoice-GGUF on first generate.",
|
||||
"supertonic3": "uv sync --extra supertonic (CPU-only ONNX, 31 langs, ~400 MB model on first use; OpenRAIL-M model license)",
|
||||
"pockettts": "pip install pocket-tts (Kyutai, CPU-only, ~100 MB model on first use; MIT code + CC-BY-4.0 weights; weights are HF-gated, set HF_TOKEN)",
|
||||
"moss-tts-v15": "git clone OpenMOSS/MOSS-TTS + set OMNIVOICE_MOSS_TTS_V15_DIR (own venv, transformers==5.0; 8B, ~16 GB weights; CUDA/CPU, no MPS; Apache-2.0)",
|
||||
"dots-tts": "git clone rednote-hilab/dots.tts + set OMNIVOICE_DOTS_TTS_DIR (own venv, transformers==4.57; 2B, ~9 GB weights; CUDA/CPU, Linux/macOS only — no Windows; Apache-2.0)",
|
||||
"confucius4-tts":"git clone netease-youdao/Confucius4-TTS + set OMNIVOICE_CONFUCIUS4_TTS_DIR (own Python 3.10 venv; 14-lang cross-lingual zero-shot clone; ~5 GB weights auto-download; CUDA/CPU, no MPS; Apache-2.0)",
|
||||
|
||||
@@ -57,6 +57,7 @@ tts_engines:
|
||||
- id: confucius4-tts
|
||||
readme: "**Confucius4-TTS**"
|
||||
doc: docs/engines/confucius4-tts.md
|
||||
- id: pockettts
|
||||
|
||||
# Same contract against backend/services/asr_backend.py _REGISTRY.
|
||||
asr_engines:
|
||||
|
||||
@@ -0,0 +1,392 @@
|
||||
"""PocketTTS sidecar unit tests (#1306 / #1328).
|
||||
|
||||
The sidecar is stdlib-only at import time (torch + pocket_tts load lazily on the
|
||||
first synthesize), so everything below runs without the optional `pocket-tts`
|
||||
wheel installed and without spawning a child process — the model is mocked.
|
||||
|
||||
These pin the four review findings that were fixed on the PR, each of which is
|
||||
silent-by-construction and would regress without a test:
|
||||
|
||||
* an unsupported language used to fall back to English and mispronounce
|
||||
(greptile P1) — it now raises;
|
||||
* the defensive downmix assumed channels-first and destroyed a channels-last
|
||||
waveform (CodeRabbit) — it now raises;
|
||||
* the cold-load heartbeat thread and the main loop both write frames, so a
|
||||
heartbeat firing mid-write interleaved the length and body segments and
|
||||
corrupted the wire (greptile P1) — writes are now serialized;
|
||||
* a reference clip replaced at the same path served the previous voice from
|
||||
cache (greptile P1) — the cache key carries an mtime+size fingerprint.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import importlib.util
|
||||
import io
|
||||
import os
|
||||
import struct
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
_SIDECAR = (
|
||||
Path(__file__).resolve().parent.parent
|
||||
/ "backend" / "engines" / "pockettts" / "main.py"
|
||||
)
|
||||
|
||||
|
||||
def _load_sidecar():
|
||||
spec = importlib.util.spec_from_file_location("pockettts_sidecar_main", _SIDECAR)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sc():
|
||||
"""Fresh module per test — the caches are module-level globals."""
|
||||
return _load_sidecar()
|
||||
|
||||
|
||||
# ── Language selection ────────────────────────────────────────────────────
|
||||
|
||||
@pytest.mark.parametrize("raw,expected", [
|
||||
("en", "english"), ("EN", "english"), ("eng", "english"), ("english", "english"),
|
||||
("fr", "french"), ("French", "french"),
|
||||
("de", "german"), ("pt", "portuguese"), ("it", "italian"), ("es", "spanish"),
|
||||
(" it ", "italian"),
|
||||
])
|
||||
def test_language_mapping(sc, raw, expected):
|
||||
assert sc._pocket_language(raw) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["", None, "auto", "AUTO", "multi", "na"])
|
||||
def test_absent_or_sentinel_language_defaults_to_english(sc, raw):
|
||||
""""auto" means the caller expressed no preference, which is not the same as
|
||||
asking for a language this engine cannot speak."""
|
||||
assert sc._pocket_language(raw) == "english"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["ja", "zh", "hi", "ru", "korean"])
|
||||
def test_an_unsupported_language_raises_instead_of_speaking_english(sc, raw):
|
||||
"""PocketTTS ships six models. Quietly handing a Japanese request to the
|
||||
English model returns confident, fluent, wrong audio — the user hears their
|
||||
text mispronounced by an English speaker and nothing reports a problem
|
||||
(greptile P1)."""
|
||||
with pytest.raises(ValueError) as e:
|
||||
sc._pocket_language(raw)
|
||||
assert "does not support" in str(e.value)
|
||||
assert raw in str(e.value), "the error must name the language that was asked for"
|
||||
# ...and say which ones do work, or the user cannot act on it.
|
||||
for code in ("en", "fr", "de", "pt", "it", "es"):
|
||||
assert code in str(e.value)
|
||||
|
||||
|
||||
# ── Waveform → PCM ────────────────────────────────────────────────────────
|
||||
|
||||
def test_pcm_roundtrip_mono(sc):
|
||||
arr = np.array([0.0, 0.5, -0.5, 1.0, -1.0], dtype=np.float32)
|
||||
b64, sr, n = sc._tensor_to_pcm_b64(arr, 24000)
|
||||
assert (sr, n) == (24000, 5)
|
||||
back = np.frombuffer(base64.b64decode(b64), dtype=np.int16)
|
||||
assert back.tolist() == [0, 16383, -16383, 32767, -32767]
|
||||
|
||||
|
||||
def test_pcm_clips_out_of_range(sc):
|
||||
"""Anything beyond [-1, 1] would wrap to the opposite sign as int16 — a loud
|
||||
click in the output rather than a clipped peak."""
|
||||
arr = np.array([2.0, -2.0], dtype=np.float32)
|
||||
b64, _, _ = sc._tensor_to_pcm_b64(arr, 24000)
|
||||
assert np.frombuffer(base64.b64decode(b64), dtype=np.int16).tolist() == [32767, -32767]
|
||||
|
||||
|
||||
def test_pcm_squeezes_a_leading_batch_axis(sc):
|
||||
"""(1, N) is the ordinary shape a batch-1 model returns; it must not trip the
|
||||
multi-channel guard below."""
|
||||
b64, _, n = sc._tensor_to_pcm_b64(np.zeros((1, 8), dtype=np.float32), 24000)
|
||||
assert n == 8
|
||||
|
||||
|
||||
def test_multichannel_audio_raises_rather_than_being_downmixed_wrongly(sc):
|
||||
"""The original `arr.mean(axis=0)` assumed channels-first. For a
|
||||
channels-last (N, 2) array it averages across TIME, not across channels —
|
||||
every output sample becomes the mean of two neighbouring samples, which is
|
||||
not a downmix but a destroyed waveform played back as noise.
|
||||
|
||||
PocketTTS returns mono, so this is unreachable today; the point is that if
|
||||
that ever changes it surfaces as an error frame instead of as garbage audio
|
||||
nobody can trace (CodeRabbit)."""
|
||||
for shape in [(100, 2), (2, 100)]:
|
||||
with pytest.raises(ValueError) as e:
|
||||
sc._tensor_to_pcm_b64(np.zeros(shape, dtype=np.float32), 24000)
|
||||
assert "mono" in str(e.value)
|
||||
assert str(shape[0]) in str(e.value), "the error must report the shape it got"
|
||||
|
||||
|
||||
# ── Wire framing ──────────────────────────────────────────────────────────
|
||||
|
||||
def test_send_recv_roundtrip(sc):
|
||||
buf = io.BytesIO()
|
||||
sc._send(buf, {"op": "audio", "n_samples": 5})
|
||||
buf.seek(0)
|
||||
assert sc._recv(buf) == {"op": "audio", "n_samples": 5}
|
||||
|
||||
|
||||
def test_recv_returns_none_at_eof(sc):
|
||||
"""A closed pipe is an orderly parent shutdown, not an error."""
|
||||
assert sc._recv(io.BytesIO(b"")) is None
|
||||
|
||||
|
||||
def test_recv_rejects_an_oversized_frame(sc):
|
||||
"""Without the cap a corrupt length header allocates unbounded memory."""
|
||||
with pytest.raises(IOError, match="frame too large"):
|
||||
sc._recv(io.BytesIO(struct.pack("!I", sc.MAX_FRAME_BYTES + 1)))
|
||||
|
||||
|
||||
def test_recv_raises_on_a_truncated_body(sc):
|
||||
"""A body shorter than its header means the child died mid-write; looping on
|
||||
a stream that will never yield more would hang the parent instead."""
|
||||
with pytest.raises(IOError, match="short read"):
|
||||
sc._recv(io.BytesIO(struct.pack("!I", 100) + b"{}"))
|
||||
|
||||
|
||||
def test_concurrent_sends_do_not_interleave_frames(sc):
|
||||
"""The cold-load heartbeat thread emits progress frames while the main loop
|
||||
may emit the audio frame. `_send` writes the length and the body as two
|
||||
separate calls, so without serialization one thread's header can land
|
||||
between another's header and body — the parent then reads a length that
|
||||
belongs to a different frame and the pipe is desynchronized for good
|
||||
(greptile P1).
|
||||
|
||||
This fails without the lock: the writer sleeps between the two writes, which
|
||||
is exactly the window the real code has.
|
||||
"""
|
||||
chunks: list[bytes] = []
|
||||
|
||||
class _SlowStream:
|
||||
"""Records writes in arrival order and yields between them."""
|
||||
|
||||
def write(self, b):
|
||||
chunks.append(bytes(b))
|
||||
# Force a thread switch in the gap the lock exists to close.
|
||||
threading.Event().wait(0.001)
|
||||
|
||||
def flush(self):
|
||||
pass
|
||||
|
||||
threads = [
|
||||
threading.Thread(target=sc._send, args=(_SlowStream(), {"op": "progress", "i": i}))
|
||||
for i in range(8)
|
||||
]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
# Every frame must decode back out cleanly and in one piece.
|
||||
stream = io.BytesIO(b"".join(chunks))
|
||||
seen = []
|
||||
while (msg := sc._recv(stream)) is not None:
|
||||
seen.append(msg)
|
||||
assert sorted(m["i"] for m in seen) == list(range(8)), (
|
||||
f"frames interleaved on the wire; decoded {len(seen)} of 8"
|
||||
)
|
||||
|
||||
|
||||
# ── Voice state cache ─────────────────────────────────────────────────────
|
||||
|
||||
class _FakeModel:
|
||||
"""Counts encodes so cache hits are observable."""
|
||||
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def get_state_for_audio_prompt(self, voice):
|
||||
self.calls.append(voice)
|
||||
return f"state:{voice}:{len(self.calls)}"
|
||||
|
||||
|
||||
def test_a_url_reference_is_refused(sc):
|
||||
"""The sidecar is local-first; handing a URL to the model would make it
|
||||
fetch on the user's behalf (SSRF, and a silent network call from an app that
|
||||
promises not to make them)."""
|
||||
for url in ["http://x/a.wav", "HTTPS://x/a.wav", "file:///etc/passwd", "ftp://x/a.wav"]:
|
||||
with pytest.raises(ValueError, match="local file path"):
|
||||
sc._voice_state(_FakeModel(), "english", url)
|
||||
|
||||
|
||||
def test_no_reference_uses_the_languages_default_voice(sc):
|
||||
"""Falling back to the English preset for an Italian request would clone an
|
||||
English speaker onto Italian text."""
|
||||
for lang, voice in [("italian", "giovanni"), ("spanish", "lola"),
|
||||
("german", "juergen"), ("french", "estelle"),
|
||||
("portuguese", "rafael"), ("english", "alba")]:
|
||||
model = _FakeModel()
|
||||
sc._voice_state(model, lang, None)
|
||||
assert model.calls == [voice]
|
||||
|
||||
|
||||
def test_the_same_reference_is_encoded_once(sc, tmp_path):
|
||||
ref = tmp_path / "ref.wav"
|
||||
ref.write_bytes(b"RIFF")
|
||||
model = _FakeModel()
|
||||
a = sc._voice_state(model, "english", str(ref))
|
||||
b = sc._voice_state(model, "english", str(ref))
|
||||
assert a == b and len(model.calls) == 1
|
||||
|
||||
|
||||
def test_a_replaced_reference_file_is_re_encoded(sc, tmp_path):
|
||||
"""Re-recording a clip and saving over the same filename is the ordinary way
|
||||
a user iterates on a voice. Keyed on the path alone, the cache kept serving
|
||||
the old recording and no amount of re-recording changed the output
|
||||
(greptile P1)."""
|
||||
ref = tmp_path / "ref.wav"
|
||||
ref.write_bytes(b"first recording")
|
||||
model = _FakeModel()
|
||||
sc._voice_state(model, "english", str(ref))
|
||||
|
||||
ref.write_bytes(b"second recording, different length")
|
||||
os.utime(ref, (1_000_000, 1_000_000))
|
||||
sc._voice_state(model, "english", str(ref))
|
||||
assert len(model.calls) == 2, "the replaced clip was served from cache"
|
||||
|
||||
|
||||
def test_a_same_size_replacement_is_caught_by_nanosecond_mtime(sc, tmp_path):
|
||||
"""Size alone misses a re-record of identical length, and whole-second mtime
|
||||
misses one written within the same second — which is precisely what a script
|
||||
or a fast save does."""
|
||||
ref = tmp_path / "ref.wav"
|
||||
ref.write_bytes(b"AAAA")
|
||||
model = _FakeModel()
|
||||
sc._voice_state(model, "english", str(ref))
|
||||
|
||||
st = os.stat(ref)
|
||||
ref.write_bytes(b"BBBB") # same size
|
||||
if os.stat(ref).st_mtime_ns == st.st_mtime_ns:
|
||||
pytest.skip("filesystem mtime resolution too coarse to distinguish")
|
||||
sc._voice_state(model, "english", str(ref))
|
||||
assert len(model.calls) == 2
|
||||
|
||||
|
||||
def test_a_missing_reference_still_reaches_the_model(sc, tmp_path):
|
||||
"""stat() failing is not this function's call to make — the model owns what
|
||||
it can resolve, and a path inside the sidecar's own namespace may be valid
|
||||
even when it cannot be stat'd from here."""
|
||||
model = _FakeModel()
|
||||
sc._voice_state(model, "english", str(tmp_path / "gone.wav"))
|
||||
assert len(model.calls) == 1
|
||||
|
||||
|
||||
def test_the_voice_cache_is_bounded(sc, tmp_path):
|
||||
"""A 50-speaker dub would otherwise hold every encoded voice state for the
|
||||
life of the process."""
|
||||
model = _FakeModel()
|
||||
for i in range(sc._VOICE_CACHE_MAX + 5):
|
||||
p = tmp_path / f"r{i}.wav"
|
||||
p.write_bytes(b"x")
|
||||
sc._voice_state(model, "english", str(p))
|
||||
assert len(sc._voice_cache) <= sc._VOICE_CACHE_MAX
|
||||
|
||||
|
||||
def test_the_cache_evicts_least_recently_used(sc, tmp_path):
|
||||
"""LRU, not FIFO: the voice being used on every line is the one that must
|
||||
survive a burst of one-off speakers."""
|
||||
model = _FakeModel()
|
||||
refs = []
|
||||
for i in range(sc._VOICE_CACHE_MAX):
|
||||
p = tmp_path / f"r{i}.wav"
|
||||
p.write_bytes(b"x")
|
||||
refs.append(str(p))
|
||||
sc._voice_state(model, "english", p and str(p))
|
||||
|
||||
sc._voice_state(model, "english", refs[0]) # touch the oldest
|
||||
n_before = len(model.calls)
|
||||
newcomer = tmp_path / "new.wav"
|
||||
newcomer.write_bytes(b"x")
|
||||
sc._voice_state(model, "english", str(newcomer)) # forces one eviction
|
||||
|
||||
sc._voice_state(model, "english", refs[0])
|
||||
assert len(model.calls) == n_before + 1, "the recently-used voice was evicted"
|
||||
|
||||
|
||||
def test_languages_do_not_share_cache_entries(sc, tmp_path):
|
||||
"""The same clip encoded for the Italian model is not the Italian model's
|
||||
state — the key has to carry the language."""
|
||||
ref = tmp_path / "ref.wav"
|
||||
ref.write_bytes(b"x")
|
||||
model = _FakeModel()
|
||||
sc._voice_state(model, "english", str(ref))
|
||||
sc._voice_state(model, "italian", str(ref))
|
||||
assert len(model.calls) == 2
|
||||
|
||||
|
||||
# ── Backend surface ───────────────────────────────────────────────────────
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def backend():
|
||||
import sys as _sys
|
||||
_sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "backend"))
|
||||
from engines.pockettts import PocketTTSBackend
|
||||
return PocketTTSBackend
|
||||
|
||||
|
||||
def test_recv_timeout_rejects_non_finite_values(backend, monkeypatch):
|
||||
"""inf would disable the deadline entirely, so a wedged sidecar would never
|
||||
be reaped — the watchdog is the whole reason this engine is a subprocess."""
|
||||
for raw in ("inf", "-inf", "nan", "NaN"):
|
||||
monkeypatch.setenv("OMNIVOICE_POCKETTTS_RECV_TIMEOUT_S", raw)
|
||||
assert backend().recv_timeout_s == 600.0
|
||||
|
||||
|
||||
def test_recv_timeout_rejects_garbage(backend, monkeypatch):
|
||||
monkeypatch.setenv("OMNIVOICE_POCKETTTS_RECV_TIMEOUT_S", "soon")
|
||||
assert backend().recv_timeout_s == 600.0
|
||||
|
||||
|
||||
def test_recv_timeout_has_a_floor(backend, monkeypatch):
|
||||
"""A 1s deadline kills every cold load before it can finish."""
|
||||
monkeypatch.setenv("OMNIVOICE_POCKETTTS_RECV_TIMEOUT_S", "1")
|
||||
assert backend().recv_timeout_s == 30.0
|
||||
|
||||
|
||||
def test_recv_timeout_honours_a_sane_override(backend, monkeypatch):
|
||||
monkeypatch.setenv("OMNIVOICE_POCKETTTS_RECV_TIMEOUT_S", "900")
|
||||
assert backend().recv_timeout_s == 900.0
|
||||
|
||||
|
||||
def test_is_available_reports_why_the_import_failed(backend):
|
||||
""""not installed" sends a user with a torch ABI mismatch or a half-written
|
||||
wheel to reinstall a package they already have (CodeRabbit)."""
|
||||
ok, msg = backend.is_available()
|
||||
if ok:
|
||||
pytest.skip("pocket-tts is installed in this environment")
|
||||
assert "pocket_tts" in msg
|
||||
# The bare message would end after the install hint; the cause has to be in it.
|
||||
assert "(" in msg and ")" in msg
|
||||
|
||||
|
||||
def test_engine_is_cpu_only_and_advertises_it(backend):
|
||||
"""Kyutai reports no GPU speedup for this 100M batch-1 model, so claiming
|
||||
CUDA would send the scheduler looking for a device it cannot use."""
|
||||
assert backend.gpu_compat == ("cpu",)
|
||||
assert backend.supports_cloning is True
|
||||
|
||||
|
||||
def test_sample_rate_is_in_lockstep_with_the_sidecar(sc, backend):
|
||||
"""The parent sizes buffers from its own constant and the sidecar stamps the
|
||||
frame with its; a drift between them resamples every render."""
|
||||
assert backend().sample_rate == sc.POCKETTTS_SAMPLE_RATE
|
||||
|
||||
|
||||
def test_the_engine_is_registered_lazily(backend):
|
||||
"""Registered eagerly, the optional pocket-tts import would run for every
|
||||
user on every startup."""
|
||||
import sys as _sys
|
||||
_sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "backend"))
|
||||
from services.tts_backend import _LAZY_REGISTRY
|
||||
assert _LAZY_REGISTRY["pockettts"] == ("engines.pockettts", "PocketTTSBackend")
|
||||
|
||||
|
||||
def test_the_sidecar_script_path_resolves(backend):
|
||||
assert backend.sidecar_script().is_file()
|
||||
Reference in New Issue
Block a user