feat(omnivoice): streaming /tts + language-safe sanitizer

Add a live-consumer streaming path and text sanitation to the OmniVoice
wrapper, so it can front speech-to-speech chat engines (not just the
asset-engine's batch WAV use).

- POST /tts: chunked 24 kHz mono s16le PCM (or open-ended WAV), driven by
  the adaptive buffer-ratchet scheduler. Emits the first sentence
  immediately, then ratchets chunk size up on OmniVoice's ~40x realtime
  headroom -> sub-second time-to-first-audio. Wire-compatible with
  chatterbox-fast /tts (both 24 kHz mono PCM). Batch /v1/audio/speech is
  unchanged for asset/file callers.

- scheduler.py: VENDORED byte-faithful copy of chatterbox-fast's pure-
  Python (torch-free) scheduler, pinned to commit 7631462 (v0.1.0/v0.1.1).
  Vendor-copy over a shared package (operator call 2026-06-19): the module
  has no GPU deps, so reuse it without dragging chatterbox-fast's torch
  tree into this image. Promote to a shared package only on a 3rd consumer
  or real drift.

- sanitize.py: language-safe TTS sanitizer run on both endpoints. Strips
  markdown, <think> blocks, HTML, and model control tokens; deliberately
  SKIPS the fork's English-only number/phone normalization that would
  corrupt OmniVoice's 600-language input. Preserves [laughter]-style tags.

- Refactor: shared GenParams base for SpeechRequest + TTSStreamRequest;
  single GEN_LOCK serializes generation (single-stream interactive).

- Dockerfile/playbook: copy + upload the two new modules; build-time
  `import app` smoke; correct stale "Gradio demo / no FastAPI" comments.
This commit is contained in:
2026-06-19 22:47:15 -07:00
parent 826c2a6a64
commit 288d085236
7 changed files with 671 additions and 67 deletions
+4
View File
@@ -35,3 +35,7 @@ htpasswd-new
# graphify: commit only the lightweight labeled map; ignore heavy/regenerable artifacts
graphify-out/*
!graphify-out/GRAPH_REPORT.md
# Python bytecode (e.g. from local py_compile of stack wrappers)
__pycache__/
*.pyc
+18 -5
View File
@@ -1,11 +1,12 @@
# Deploy OmniVoice (https://github.com/k2-fsa/OmniVoice) to irv-ml1, GPU 0
# (RTX 3090). Apache-2.0 zero-shot multilingual voice-cloning TTS, served
# via upstream's own Gradio demo (no FastAPI wrapper).
# behind our OWN FastAPI wrapper (app.py): batch /v1/audio/speech plus a
# streaming /tts driven by the vendored buffer-ratchet scheduler.
#
# Builds the image locally from stacks/omnivoice/Dockerfile (CUDA 12.8 +
# torch 2.8.0 + omnivoice from PyPI), stages the build context under
# /opt/docker/compose/omnivoice/, brings it up, and waits for the Gradio
# UI on :8199.
# torch 2.8.0 + omnivoice from PyPI + vendored scheduler.py/sanitize.py),
# stages the build context under /opt/docker/compose/omnivoice/, brings it
# up, and waits for /healthz on :8199.
#
# First run is slow: ~5-10 min docker build + a one-time HF weight pre-warm
# (k2-fsa/OmniVoice) on first container start (entrypoint.sh). The wait loop
@@ -61,12 +62,24 @@ steps:
dest: "{{ compose_dir }}/Dockerfile"
mode: "0644"
- name: Upload app.py (asset-engine FastAPI wrapper)
- name: Upload app.py (batch + streaming FastAPI wrapper)
upload:
src: stacks/omnivoice/app.py
dest: "{{ compose_dir }}/app.py"
mode: "0644"
- name: Upload scheduler.py (vendored buffer-ratchet streaming scheduler)
upload:
src: stacks/omnivoice/scheduler.py
dest: "{{ compose_dir }}/scheduler.py"
mode: "0644"
- name: Upload sanitize.py (language-safe TTS text sanitizer)
upload:
src: stacks/omnivoice/sanitize.py
dest: "{{ compose_dir }}/sanitize.py"
mode: "0644"
- name: Stage chatterbox reference voices for cloning (skip _*.wav artifacts)
shell: |
set -e
+9 -2
View File
@@ -4,8 +4,9 @@
# languages) voice-cloning + voice-design TTS, diffusion-LM architecture,
# Apache-2.0. Upstream ships a pip package + its own Gradio demo
# (`omnivoice-demo`); there's no official image, so we build a thin CUDA
# container around the pip package and run its Gradio server directly.
# Unlike index-tts we DON'T need a FastAPI wrapper — OmniVoice serves itself.
# container around the pip package and run our OWN FastAPI wrapper (app.py):
# a batch OpenAI-style /v1/audio/speech plus a streaming /tts driven by the
# vendored buffer-ratchet scheduler (scheduler.py) for live chat consumers.
ARG CUDA_BASE=nvidia/cuda:12.8.0-cudnn-runtime-ubuntu22.04
FROM ${CUDA_BASE}
@@ -41,9 +42,15 @@ RUN python -c "import omnivoice, fastapi, soundfile, uvicorn; print('omnivoice w
WORKDIR /app
COPY app.py /app/app.py
COPY scheduler.py /app/scheduler.py
COPY sanitize.py /app/sanitize.py
COPY entrypoint.sh /usr/local/bin/entrypoint.sh
RUN chmod +x /usr/local/bin/entrypoint.sh
# Fail the build if the wrapper (incl. the vendored scheduler + sanitizer) won't
# import. Model load is lazy (startup event), so this is a cheap CPU-only check.
RUN python -c "import app; print('omnivoice app import OK')"
EXPOSE 8001
# Our wrapper exposes /healthz (200 once model + >=1 voice are loaded). Generous
+33 -5
View File
@@ -18,19 +18,47 @@ RTF as low as ~0.025 (≈40× real-time). **Apache-2.0** — commercially clean
## How it's served
Behind our own thin **FastAPI wrapper** ([`app.py`](app.py)) — upstream ships
only a Gradio demo, which we replaced (2026-06-19) so the **asset-engine** can
consume it. Endpoints on `http://10.100.79.3:8199`:
only a Gradio demo, which we replaced (2026-06-19). The wrapper serves **two
consumption modes** on `http://10.100.79.3:8199`:
| Endpoint | Purpose |
|---|---|
| `POST /v1/audio/speech` | OpenAI-style `{input, voice, response_format=wav}` → 24 kHz PCM_16 mono |
| `POST /v1/audio/speech` | **Batch** OpenAI-style `{input, voice, instruct, language, …}` one 24 kHz PCM_16 mono WAV. For the **asset-engine** (form-driven asset generation). |
| `POST /tts` | **Streaming** chunked 24 kHz mono `s16le` PCM (`format=pcm`, default) or open-ended WAV — for **live speech-to-speech chat engines**. Wire-compatible with chatterbox-fast's `/tts`. |
| `GET /v1/audio/voices` | `{"voices": [...]}` — the staged clone targets |
| `GET /v1/audio/languages` | `{"languages": ["Auto", …]}` — 600+ |
| `GET /v1/audio/instruct-items` | `{"instruct_items": [...]}` — controlled voice-DESIGN tags |
| `GET /healthz` | readiness (200 once model + ≥1 voice loaded) |
The wrapper loads OmniVoice + a Whisper ASR and **precomputes a voice-clone
prompt per staged reference WAV at startup** (Whisper auto-transcribes each
reference), so per-request latency is just generation. v1 is **clone-only**
OmniVoice's voice-*design* / language / instruct controls aren't exposed yet.
reference), so per-request latency is just generation. The full generation
surface is exposed: zero-shot **clone** (`voice`) and/or voice-**design**
(`instruct`), plus `language` / `speed` / `duration` and the diffusion knobs.
### Streaming — sub-second time-to-first-audio
`POST /tts` (`stream=true`, default) runs the **adaptive buffer-ratchet
scheduler** vendored from chatterbox-fast ([`scheduler.py`](scheduler.py)): it
emits the first sentence immediately and ratchets chunk size up on OmniVoice's
~40× realtime headroom, so a live consumer hears speech start in ~tens of ms
instead of waiting for the whole utterance. `stream=false` is a whole-text
one-shot for A/B. Scheduler tunables (`margin`, `margin_first`, `rtf_prior`,
`sec_per_char_prior`) are per-request overrides.
`scheduler.py` is a **vendored byte-faithful copy** (not a dependency) of
chatterbox-fast's pure-Python, torch-free scheduler — see its header for the
pinned commit. It is reused without dragging chatterbox-fast's GPU dependency
tree into this image; re-vendor on upstream change rather than editing in place.
### Text sanitization
Both endpoints run `input` through a **language-safe sanitizer**
([`sanitize.py`](sanitize.py)) before synthesis: it strips markdown, LLM
artifacts (`<think>` blocks), HTML, and model control tokens, but deliberately
**skips** English-only number/phone/email normalization that would corrupt
OmniVoice's multilingual input. OmniVoice's own `[laughter]`-style symbols are
preserved.
### Voices — reused from chatterbox
+219 -55
View File
@@ -1,5 +1,4 @@
"""Thin FastAPI wrapper exposing OmniVoice (k2-fsa/OmniVoice) as an OpenAI-style
TTS for the asset-engine.
"""Thin FastAPI wrapper exposing OmniVoice (k2-fsa/OmniVoice) for the fleet.
Upstream ships only a Gradio demo; we own this wrapper (same pattern as
stacks/index-tts/app.py). Voices are reference WAVs staged in
@@ -7,35 +6,55 @@ ${OMNIVOICE_VOICES_DIR} (the reused chatterbox /refs/*.wav). A voice-clone promp
is precomputed once per voice at startup (the loaded Whisper ASR auto-transcribes
each reference) and cached, so per-request latency is just generation.
Two consumption modes:
- BATCH (asset-engine / OpenAI-compat): POST /v1/audio/speech -> one WAV blob.
- STREAM (live speech-to-speech chat engines): POST /tts -> chunked PCM, driven
by the vendored adaptive buffer-ratchet scheduler (scheduler.py, from
chatterbox-fast). Emits the first sentence immediately for sub-second
time-to-first-audio, then ratchets chunk size up on OmniVoice's ~40x realtime
headroom. Wire-compatible with chatterbox-fast's /tts (both 24 kHz mono s16le).
All text is run through the language-safe sanitizer (sanitize.py) before synthesis
on BOTH endpoints — strips markdown / LLM artifacts / control tokens without the
English-only normalization that would corrupt OmniVoice's multilingual input.
Exposes OmniVoice's full generation surface:
- clone (voice=<staged ref>) and/or voice-DESIGN (instruct=<free-text style>)
- clone (voice=<staged ref>) and/or voice-DESIGN (instruct=<controlled tags>)
- language (Auto + 600+), speed, duration
- diffusion controls: num_step, guidance_scale, denoise, preprocess_prompt,
postprocess_output, plus a generation_overrides passthrough for expert knobs
(t_shift, layer_penalty_factor, position_temperature, class_temperature, ...).
Endpoints:
GET /healthz -> readiness (200 once model + >=1 voice are loaded)
GET /v1/audio/voices -> {"voices": [<name>, ...]}
GET /v1/audio/languages -> {"languages": ["Auto", <display name>, ...]}
POST /v1/audio/speech -> audio/wav
GET /healthz -> readiness (200 once model + >=1 voice are loaded)
GET /v1/audio/voices -> {"voices": [<name>, ...]}
GET /v1/audio/languages -> {"languages": ["Auto", <display name>, ...]}
GET /v1/audio/instruct-items -> {"instruct_items": [<tag>, ...]}
POST /v1/audio/speech -> audio/wav (batch, OpenAI-style)
POST /tts -> streaming PCM/WAV (chatterbox-fast-compatible)
"""
import glob
import io
import logging
import os
import struct
import threading
import time
from pathlib import Path
from typing import Any, Dict, Optional
from typing import Any, Dict, Literal, Optional
import numpy as np
import soundfile as sf
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel
from fastapi.responses import Response, StreamingResponse
from pydantic import BaseModel, Field
from omnivoice import OmniVoice, OmniVoiceGenerationConfig
from sanitize import sanitize_tts_text
from scheduler import ChunkConfig, ChunkResult, stream_chunks
try:
from omnivoice.utils.lang_map import LANG_NAMES, lang_display_name
LANGUAGES = ["Auto"] + sorted(lang_display_name(n) for n in LANG_NAMES)
@@ -57,18 +76,25 @@ CKPT = os.environ.get("OMNIVOICE_CKPT", "k2-fsa/OmniVoice")
VOICES_DIR = os.environ.get("OMNIVOICE_VOICES_DIR", "/app/voices")
ASR_MODEL = os.environ.get("OMNIVOICE_ASR_MODEL", "openai/whisper-large-v3-turbo")
app = FastAPI(title="OmniVoice TTS (asset-engine wrapper)")
app = FastAPI(title="OmniVoice TTS (asset-engine + streaming wrapper)")
MODEL: Optional[OmniVoice] = None
PROMPTS: dict = {} # voice name -> VoiceClonePrompt
SR: int = 24000
# Generation is serialized: the workload is single-stream interactive and a
# streaming request holds the model for the duration of its stream. Concurrent
# callers queue rather than interleave on the GPU.
GEN_LOCK = threading.Lock()
class GenParams(BaseModel):
"""OmniVoice generation parameters shared by the batch and streaming endpoints."""
class SpeechRequest(BaseModel):
input: str
# Voice source — at least one of voice (clone) / instruct (design) is required.
voice: Optional[str] = None # staged reference clip -> clone timbre
instruct: Optional[str] = None # free-text voice DESIGN / style
instruct: Optional[str] = None # voice DESIGN / style (controlled tags)
# Generation controls (defaults mirror the upstream demo).
language: Optional[str] = "Auto" # "Auto" -> auto-detect
speed: Optional[float] = None # 0.51.5; ignored if duration set
@@ -82,10 +108,27 @@ class SpeechRequest(BaseModel):
# layer_penalty_factor, position_temperature, class_temperature,
# audio_chunk_duration, audio_chunk_threshold). Unknown keys are dropped.
generation_overrides: Optional[Dict[str, Any]] = None
class SpeechRequest(GenParams):
"""OpenAI-style /v1/audio/speech (batch) request."""
response_format: str = "wav"
model: Optional[str] = None # ignored (single model); OpenAI-compat
class TTSStreamRequest(GenParams):
"""Streaming /tts request — chatterbox-fast-compatible wire protocol."""
format: Literal["pcm", "wav"] = "pcm" # raw s16le PCM (default) or open-ended WAV
stream: bool = True # False -> whole-text one-shot (A/B vs stream)
# Scheduler overrides (None -> ChunkConfig defaults; see scheduler.py).
margin: Optional[float] = Field(default=None)
margin_first: Optional[float] = Field(default=None)
rtf_prior: Optional[float] = Field(default=None)
sec_per_char_prior: Optional[float] = Field(default=None)
@app.on_event("startup")
def _load() -> None:
global MODEL, SR
@@ -108,6 +151,104 @@ def _load() -> None:
log.info("%d voices loaded; %d languages", len(PROMPTS), len(LANGUAGES))
# ── generation helpers (shared by batch + streaming) ────────────────────────
def _base_gen_kwargs(req: GenParams) -> Dict[str, Any]:
"""Validate the voice source and build the MODEL.generate kwargs minus `text`.
Raises HTTPException (400/404) for caller errors — call this BEFORE a stream
starts so those land as proper status codes, not mid-stream failures.
"""
has_instruct = bool(req.instruct and req.instruct.strip())
if not req.voice and not has_instruct:
raise HTTPException(
status_code=400,
detail="provide a 'voice' (clone a staged reference) and/or 'instruct' (design a voice)",
)
cfg = {
"num_step": req.num_step,
"guidance_scale": req.guidance_scale,
"denoise": req.denoise,
"preprocess_prompt": req.preprocess_prompt,
"postprocess_output": req.postprocess_output,
**(req.generation_overrides or {}),
}
kw: Dict[str, Any] = {"generation_config": OmniVoiceGenerationConfig.from_dict(cfg)}
if req.voice:
prompt = PROMPTS.get(req.voice)
if prompt is None:
raise HTTPException(
status_code=404,
detail=f"unknown voice '{req.voice}'; have {sorted(PROMPTS)}",
)
kw["voice_clone_prompt"] = prompt
if has_instruct:
kw["instruct"] = req.instruct.strip()
if req.language and req.language != "Auto":
kw["language"] = req.language
if req.speed is not None:
kw["speed"] = req.speed
if req.duration is not None:
kw["duration"] = req.duration
return kw
def _synth(text: str, base_kw: Dict[str, Any]) -> "tuple[np.ndarray, float]":
"""Synthesize one text span -> (float32 audio [-1,1], audio_seconds)."""
out = MODEL.generate(text=text, **base_kw)
audio = out[0] if isinstance(out, (list, tuple)) else out
audio = np.asarray(audio, dtype=np.float32)
return audio, (len(audio) / SR if SR else 0.0)
def _pcm16(audio: np.ndarray) -> bytes:
"""float32 [-1,1] -> little-endian s16 PCM bytes (24 kHz mono on the wire)."""
a = np.clip(np.asarray(audio, dtype=np.float32).reshape(-1), -1.0, 1.0)
return (a * 32767.0).astype("<i2").tobytes()
def _wav_header(sr: int, data_len: Optional[int] = None) -> bytes:
"""WAV header. data_len=None -> streaming (0xFFFFFFFF sizes, read to EOF);
an int -> correct RIFF/data sizes for a complete file."""
data_size = 0xFFFFFFFF if data_len is None else data_len
riff_size = 0xFFFFFFFF if data_len is None else 36 + data_len
return (
b"RIFF" + struct.pack("<I", riff_size) + b"WAVE"
+ b"fmt " + struct.pack("<IHHIIHH", 16, 1, 1, sr, sr * 2, 2, 16)
+ b"data" + struct.pack("<I", data_size)
)
def _chunk_config(req: TTSStreamRequest) -> ChunkConfig:
cfg = ChunkConfig()
if req.margin is not None:
cfg.margin = req.margin
if req.margin_first is not None:
cfg.margin_first = req.margin_first
if req.rtf_prior is not None:
cfg.rtf_prior = req.rtf_prior
if req.sec_per_char_prior is not None:
cfg.sec_per_char_prior = req.sec_per_char_prior
return cfg
def _log_chunk(r: ChunkResult, ttfa_ms: float) -> None:
tag = " STARVED" if r.starved else ""
if r.index == 0:
log.info("chunk 0: ttfa=%.0fms gen=%.0fms audio=%.2fs rtf=%.2f%s",
ttfa_ms, r.gen_time * 1000, r.audio_sec, r.rtf, tag)
else:
log.info("chunk %d: gen=%.0fms audio=%.2fs buf=%.2f->%.2f rtf=%.2f%s",
r.index, r.gen_time * 1000, r.audio_sec,
r.buffer_before, r.buffer_after, r.rtf, tag)
# ── read-only discovery endpoints ───────────────────────────────────────────
@app.get("/healthz")
def healthz():
if MODEL is None or not PROMPTS:
@@ -132,59 +273,82 @@ def instruct_items():
return {"instruct_items": INSTRUCT_ITEMS}
# ── synthesis endpoints ─────────────────────────────────────────────────────
@app.post("/v1/audio/speech")
def speech(req: SpeechRequest):
"""Batch (OpenAI-style): generate the whole utterance, return one WAV blob."""
if MODEL is None:
raise HTTPException(status_code=503, detail="model not loaded")
if not req.input or not req.input.strip():
raise HTTPException(status_code=400, detail="input is empty")
if req.response_format not in ("wav", "", None):
raise HTTPException(status_code=400, detail="only response_format=wav is supported")
text = sanitize_tts_text(req.input)
if not text:
raise HTTPException(status_code=400, detail="input is empty after sanitization")
has_instruct = bool(req.instruct and req.instruct.strip())
if not req.voice and not has_instruct:
raise HTTPException(
status_code=400,
detail="provide a 'voice' (clone a staged reference) and/or 'instruct' (design a voice)",
)
cfg = {
"num_step": req.num_step,
"guidance_scale": req.guidance_scale,
"denoise": req.denoise,
"preprocess_prompt": req.preprocess_prompt,
"postprocess_output": req.postprocess_output,
**(req.generation_overrides or {}),
}
gen = OmniVoiceGenerationConfig.from_dict(cfg)
kw: Dict[str, Any] = {"text": req.input.strip(), "generation_config": gen}
if req.voice:
prompt = PROMPTS.get(req.voice)
if prompt is None:
raise HTTPException(
status_code=404,
detail=f"unknown voice '{req.voice}'; have {sorted(PROMPTS)}",
)
kw["voice_clone_prompt"] = prompt
if has_instruct:
kw["instruct"] = req.instruct.strip()
if req.language and req.language != "Auto":
kw["language"] = req.language
if req.speed is not None:
kw["speed"] = req.speed
if req.duration is not None:
kw["duration"] = req.duration
base_kw = _base_gen_kwargs(req) # 400/404 propagate as-is
try:
out = MODEL.generate(**kw)
with GEN_LOCK:
audio, _ = _synth(text, base_kw)
except HTTPException:
raise
except Exception as exc: # noqa: BLE001
raise HTTPException(status_code=400, detail=f"{type(exc).__name__}: {exc}")
audio = out[0] if isinstance(out, (list, tuple)) else out
audio = np.asarray(audio, dtype=np.float32)
buf = io.BytesIO()
sf.write(buf, audio, SR, format="WAV", subtype="PCM_16")
return Response(content=buf.getvalue(), media_type="audio/wav")
@app.post("/tts")
def tts(req: TTSStreamRequest):
"""Streaming: chunked 24 kHz mono PCM (or open-ended WAV) for live consumers.
Wire-compatible with chatterbox-fast /tts. `stream=true` (default) runs the
buffer-ratchet schedule for sub-second time-to-first-audio; `stream=false`
is a whole-text one-shot for A/B comparison.
"""
if MODEL is None:
raise HTTPException(status_code=503, detail="model not loaded")
text = sanitize_tts_text(req.input)
if not text:
raise HTTPException(status_code=400, detail="input is empty after sanitization")
base_kw = _base_gen_kwargs(req) # validate up front (pre-stream)
cfg = _chunk_config(req)
media = "audio/wav" if req.format == "wav" else "application/octet-stream"
def body():
# One request holds the lock for its whole stream (single-stream workload).
with GEN_LOCK:
t_req = time.perf_counter()
ttfa_ms: Optional[float] = None
if not req.stream:
audio, audio_sec = _synth(text, base_kw)
ttfa_ms = (time.perf_counter() - t_req) * 1000
log.info("oneshot: %.0fms gen, %.2fs audio", ttfa_ms, audio_sec)
pcm = _pcm16(audio)
if req.format == "wav":
yield _wav_header(SR, len(pcm))
yield pcm
return
if req.format == "wav":
yield _wav_header(SR) # open-ended; length unknown up front
def _gen(t: str):
return _synth(t, base_kw)
total_audio = 0.0
for r in stream_chunks(text, generate=_gen, clock=time.perf_counter, cfg=cfg):
if ttfa_ms is None:
ttfa_ms = (time.perf_counter() - t_req) * 1000
total_audio += r.audio_sec
_log_chunk(r, ttfa_ms)
yield _pcm16(r.audio)
log.info("stream done: ttfa=%.0fms total_audio=%.2fs",
ttfa_ms or 0.0, total_audio)
return StreamingResponse(body(), media_type=media)
+92
View File
@@ -0,0 +1,92 @@
"""Language-safe text sanitizer for the OmniVoice TTS wrapper.
Strips artifacts that are wrong to read aloud in ANY language — LLM
reasoning/thinking blocks, markdown, HTML/XML tags, model control/special
tokens, stray control characters — and normalizes Unicode (NFKC) + whitespace.
DELIBERATELY does NOT do English-specific normalization (number / phone /
email / currency expansion). OmniVoice is a 600+-language model; those rewrites
are correct only for English and would corrupt non-English input. The fork we
took the idea from (groxaxo/omnivoice-streaming) chains that English-only step
in — we keep only the language-neutral subset (eshpfi session 2026-06-19).
OmniVoice's own inline non-verbal symbols ([laughter], [sigh], [breath], …) are
PRESERVED — they use square brackets and contain no "](...)" link tail, so none
of the markdown/HTML rules below touch them.
"""
from __future__ import annotations
import re
import unicodedata
# LLM reasoning/thinking blocks: <think>…</think>, <reasoning>…</reasoning>, etc.
# Tag + content both removed (the content is private chain-of-thought, not speech).
_THINK_RE = re.compile(
r"<(think|thinking|reasoning|reflection|scratchpad)>.*?</\1>",
re.IGNORECASE | re.DOTALL,
)
# Markdown — applied while line breaks still exist (the (?m)^ anchors need them).
_CODE_FENCE_RE = re.compile(r"```.*?```", re.DOTALL) # fenced code block → drop
_IMAGE_RE = re.compile(r"!\[[^\]]*\]\([^)]*\)") # ![alt](url) → drop
_LINK_RE = re.compile(r"\[([^\]]+)\]\([^)]*\)") # [text](url) → text
_INLINE_CODE_RE = re.compile(r"`([^`]*)`") # `code` → code
# Emphasis: ***, **, *, ~~, __ (single _ deliberately EXCLUDED so snake_case
# words like "voice_clone_prompt" survive intact).
_EMPHASIS_RE = re.compile(r"(\*\*\*|\*\*|\*|~~|__)(.*?)\1", re.DOTALL)
_HEADING_RE = re.compile(r"(?m)^[ \t]{0,3}#{1,6}[ \t]+") # "## H" → ""
_BLOCKQUOTE_RE = re.compile(r"(?m)^[ \t]{0,3}>[ \t]?") # "> q" → "q"
_LIST_MARKER_RE = re.compile(r"(?m)^[ \t]*(?:[-*+]|\d+[.)])[ \t]+") # "- x" / "1. x" → "x"
_HR_RE = re.compile(r"(?m)^[ \t]{0,3}([-*_])(?:[ \t]*\1){2,}[ \t]*$") # "---" rule → drop
# Model / chat-template control + special tokens.
_CONTROL_TOKEN_RE = re.compile(
r"<\|[^>]*?\|>" # <|im_start|>, <|eot_id|>, …
r"|</?s>" # <s> </s>
r"|<(?:pad|unk|mask|bos|eos|sep|cls)>" # common HF special tokens
r"|\[/?INST\]" # [INST] [/INST]
r"|<</?SYS>>", # <<SYS>> <</SYS>>
re.IGNORECASE,
)
# Any remaining HTML/XML tags (run AFTER think-blocks so their content is gone).
_HTML_TAG_RE = re.compile(r"<[^>]+>")
# Stray control characters (keep \n and \t; collapsed in the final whitespace pass).
_CTRL_CHARS_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")
def sanitize_tts_text(text: str) -> str:
"""Return ``text`` cleaned of artifacts that shouldn't be spoken, language-safe.
Idempotent and Unicode-aware. Collapses all whitespace to single spaces at the
end (the scheduler's splitter does the same, so paragraph structure is moot for
synthesis). Returns ``""`` for falsy input.
"""
if not text:
return ""
# 1. LLM reasoning/thinking blocks (tag + content).
text = _THINK_RE.sub(" ", text)
# 2. Markdown constructs (order matters: fences/images/links before inline/emphasis).
text = _CODE_FENCE_RE.sub(" ", text)
text = _IMAGE_RE.sub(" ", text)
text = _LINK_RE.sub(r"\1", text)
text = _INLINE_CODE_RE.sub(r"\1", text)
text = _EMPHASIS_RE.sub(r"\2", text)
text = _HEADING_RE.sub("", text)
text = _BLOCKQUOTE_RE.sub("", text)
text = _LIST_MARKER_RE.sub("", text)
text = _HR_RE.sub(" ", text)
# 3. Model control/special tokens, then any leftover HTML/XML tags.
text = _CONTROL_TOKEN_RE.sub(" ", text)
text = _HTML_TAG_RE.sub(" ", text)
# 4. Unicode normalize (NFKC) + drop stray control characters.
text = unicodedata.normalize("NFKC", text)
text = _CTRL_CHARS_RE.sub("", text)
# 5. Collapse all whitespace (incl. newlines) to single spaces, trim.
return " ".join(text.split())
+296
View File
@@ -0,0 +1,296 @@
# ─────────────────────────────────────────────────────────────────────────
# VENDORED COPY — do not edit here; edit upstream and re-vendor.
#
# Source: chatterbox-fast · chatterbox_fast/scheduler.py
# repo: gitea.phasefinal.com/vh/chatterbox-fast (MIT)
# commit: 76314624bb131d7336629d87eb8e08557d833479 ("chatterbox-fast v0.1.0")
# tag: v0.1.1
# vendored: 2026-06-19 (eshpfi-management · stacks/omnivoice)
#
# Why a vendored copy and not a dependency: this module is pure-Python
# (no torch / chatterbox / FastAPI imports — see the docstring), so OmniVoice's
# wrapper reuses the buffer-ratchet streaming scheduler WITHOUT inheriting
# chatterbox-fast's GPU dependency tree. Operator-approved vendor-copy over a
# shared micro-package (2026-06-19): promote to a shared package only once a
# THIRD consumer appears or these copies actually drift. Keep this file
# byte-identical to upstream; re-vendor on upstream change rather than editing
# in place.
# ─────────────────────────────────────────────────────────────────────────
"""Adaptive buffer-ratchet chunk scheduler for chatterbox-fast.
This module is the *meat* of the streaming engine and is deliberately free of
any GPU / torch / chatterbox / FastAPI imports so the no-starvation guarantee
can be validated as a pure simulation (see ``test_scheduler.py``).
The design (docs/design/chatterbox-fast-plan.md §1):
* Chunk 1 = the first sentence alone — generated and emitted immediately so
first-audio latency is minimal.
* While chunk N plays, generate chunk N+1 by greedily accumulating WHOLE
sentences until the next sentence would push estimated gen-time past
``margin × audio_buffered_remaining``. Never split mid-sentence — each
chunk stays prosodically self-coherent and joins land at natural pauses.
* Because Chatterbox runs faster than realtime (RTF > 1), every chunk's
playback buys wall-clock for a larger next chunk, so chunks ratchet up
~3× and after 2-3 joins the rest of the paragraph is one big near-full-
context chunk. Context loss is confined to those 2-3 sentence-boundary
joins.
* Drive off the *measured* realtime factor and sec-per-char, not constants —
track them live and self-correct.
The scheduler is *online*: chunk N+1's boundary depends on the RTF measured
while generating chunk N, so it cannot be precomputed. ``stream_chunks`` runs
the loop with ``generate`` and ``clock`` injected, which is what makes the
simulation possible.
"""
from __future__ import annotations
import re
from dataclasses import dataclass
from typing import Callable, Iterator, Sequence
# ── tunables ──────────────────────────────────────────────────────────────
@dataclass
class ChunkConfig:
"""Scheduler tunables. Priors are conservative; live measurement corrects."""
# Realtime factor prior. 3090 ~3.4×, A6000 ~3.8× (plan §2). Start
# conservative — under-estimating RTF makes chunks smaller and safer.
rtf_prior: float = 3.4
# Safety fraction of the buffer to spend on the next chunk's generation.
# margin < 1 leaves headroom so estimate error doesn't starve the stream.
margin: float = 0.8
# Tighter margin on the FIRST transition (chunk 1 → chunk 2): the buffer is
# smallest there, so starvation risk is highest (plan §1.5: ~0.6-0.7).
margin_first: float = 0.65
# Seconds of audio per character of text. ~0.060 ≈ 16.6 chars/sec speech.
# Calibrated live per request.
sec_per_char_prior: float = 0.060
# EMA weight on the newest measurement when updating rtf / sec_per_char.
ema_alpha: float = 0.4
# If the first sentence's estimated audio exceeds this, clause-split it to
# protect first-audio latency (the ONLY place we split below sentence
# granularity — plan §1.1).
max_first_sec: float = 2.0
# ── result record ─────────────────────────────────────────────────────────
@dataclass
class ChunkResult:
"""Telemetry + payload for one emitted chunk."""
index: int
text: str
audio: object # opaque payload from generate() (wav tensor, fake, …)
audio_sec: float
gen_time: float
est_gen: float # gen-time the scheduler predicted before generating
buffer_before: float # unplayed audio (s) when this chunk's gen started
buffer_after: float # unplayed audio (s) once this chunk is emitted
drained: float # seconds the buffer ran dry during gen (>0 ⇒ starvation)
rtf: float # measured RTF after this chunk
sec_per_char: float # measured sec/char after this chunk
@property
def starved(self) -> bool:
return self.drained > 1e-9
# ── text splitting ────────────────────────────────────────────────────────
# Split after sentence-final punctuation when the next non-space looks like a
# new sentence start: a capital/digit (optionally behind an open quote) or an
# inline tag like "[laugh]". Pragmatic, not perfect: abbreviations ("Mr.",
# "e.g.") can over-split — a Phase-2 watch-out, harmless to quality (just an
# extra natural-pause join).
_SENTENCE_END = re.compile(r'(?<=[.!?])["\')\]]*\s+(?=["\'(]*(?:[A-Z0-9]|\[))')
# Clause boundaries for the first-sentence latency fallback only.
_CLAUSE_END = re.compile(r'(?<=[,;:])\s+')
def split_sentences(text: str) -> list[str]:
"""Split into sentence units, preserving punctuation. Whitespace-collapsed."""
text = " ".join(text.split())
if not text:
return []
return [s for s in (p.strip() for p in _SENTENCE_END.split(text)) if s]
def _split_clauses(sentence: str) -> list[str]:
return [c for c in (p.strip() for p in _CLAUSE_END.split(sentence)) if c]
def protect_first_audio(units: list[str], cfg: ChunkConfig) -> list[str]:
"""Clause-split the first unit if it's too long to hit the first-audio target.
Only the *leading* unit is split; the remainder is untouched. If the first
sentence has no clause boundary we accept the latency rather than split
mid-clause (quality > latency once we're past the budget).
"""
if not units:
return units
if len(units[0]) * cfg.sec_per_char_prior <= cfg.max_first_sec:
return units
pieces = _split_clauses(units[0])
if len(pieces) <= 1:
return units # nothing to split on; keep the long first sentence whole
return pieces + units[1:]
# ── chunk planning ────────────────────────────────────────────────────────
def _est_gen_time(text: str, *, rtf: float, sec_per_char: float) -> float:
return (len(text) * sec_per_char) / rtf
def plan_chunk(
remaining: Sequence[str],
buffer_remaining: float,
*,
margin: float,
rtf: float,
sec_per_char: float,
) -> tuple[str, list[str]]:
"""Greedily accumulate whole units until the next would blow the budget.
Always returns at least one unit (never empty, never splits a unit). With
``buffer_remaining == 0`` (the first chunk) the budget is 0, so exactly the
first unit is taken — which is the latency-critical chunk-1 rule.
"""
budget = margin * buffer_remaining
chunk = [remaining[0]]
i = 1
while i < len(remaining):
candidate = " ".join(chunk + [remaining[i]])
if _est_gen_time(candidate, rtf=rtf, sec_per_char=sec_per_char) > budget:
break
chunk.append(remaining[i])
i += 1
return " ".join(chunk), list(remaining[i:])
def relieve_leader(
remaining: list[str],
buffer_remaining: float,
*,
rtf: float,
sec_per_char: float,
) -> list[str]:
"""Expose a too-big leading sentence's clause boundaries to avoid starvation.
If generating ``remaining[0]`` alone would overrun the buffer (an audible
gap), clause-split it in place so ``plan_chunk`` can pack clause-pieces up to
the buffer — joins then land on commas (natural pauses) instead of producing
a gap. If the sentence has no clause boundary it's returned unchanged: we do
NOT split mid-clause (design rule), and the caller accepts + flags the rare
starvation in telemetry.
"""
leader = remaining[0]
if _est_gen_time(leader, rtf=rtf, sec_per_char=sec_per_char) <= buffer_remaining:
return remaining
pieces = _split_clauses(leader)
if len(pieces) <= 1:
return remaining
return pieces + remaining[1:]
def _ema(old: float, new: float, alpha: float) -> float:
return (1 - alpha) * old + alpha * new
# ── the online loop ───────────────────────────────────────────────────────
# generate(text) -> (audio_payload, audio_seconds)
GenerateFn = Callable[[str], "tuple[object, float]"]
ClockFn = Callable[[], float]
def stream_chunks(
text: str,
*,
generate: GenerateFn,
clock: ClockFn,
cfg: ChunkConfig | None = None,
) -> Iterator[ChunkResult]:
"""Run the adaptive buffer-ratchet schedule, yielding one ChunkResult per chunk.
``generate`` does the actual synthesis (real GPU or a simulation) and
returns its opaque audio payload plus the audio's duration in seconds.
``clock`` returns monotonically increasing seconds; ``gen_time`` is measured
as the clock delta around ``generate``.
Buffer model (all in audio-seconds):
* Playback begins when chunk 1 arrives, so no drain occurs during chunk 1.
* For chunk N≥2, the client plays ``gen_time`` seconds of buffered audio
while we generate it, then we add this chunk's audio. If ``gen_time``
exceeds the buffer the stream starved (``drained`` > 0).
"""
cfg = cfg or ChunkConfig()
units = protect_first_audio(split_sentences(text), cfg)
rtf = cfg.rtf_prior
sec_per_char = cfg.sec_per_char_prior
buffer_remaining = 0.0
remaining: list[str] = units
index = 0
while remaining:
first = index == 0
if not first:
# If the next indivisible sentence would starve the buffer, expose
# its clause boundaries so plan_chunk can pack to the buffer.
remaining = relieve_leader(
remaining, buffer_remaining, rtf=rtf, sec_per_char=sec_per_char
)
margin = cfg.margin_first if first else cfg.margin
chunk_text, remaining = plan_chunk(
remaining, buffer_remaining, margin=margin, rtf=rtf, sec_per_char=sec_per_char
)
est_gen = _est_gen_time(chunk_text, rtf=rtf, sec_per_char=sec_per_char)
t0 = clock()
audio, audio_sec = generate(chunk_text)
gen_time = clock() - t0
# Starvation: did the buffer run dry while we generated this chunk?
drained = 0.0 if first else max(0.0, gen_time - buffer_remaining)
# Live self-correction.
if gen_time > 0:
rtf = _ema(rtf, audio_sec / gen_time, cfg.ema_alpha)
if chunk_text:
sec_per_char = _ema(sec_per_char, audio_sec / len(chunk_text), cfg.ema_alpha)
buffer_before = buffer_remaining
if first:
buffer_remaining = audio_sec
else:
buffer_remaining = max(0.0, buffer_remaining - gen_time) + audio_sec
yield ChunkResult(
index=index,
text=chunk_text,
audio=audio,
audio_sec=audio_sec,
gen_time=gen_time,
est_gen=est_gen,
buffer_before=buffer_before,
buffer_after=buffer_remaining,
drained=drained,
rtf=rtf,
sec_per_char=sec_per_char,
)
index += 1