diff --git a/.gitignore b/.gitignore index bf16260..7e108f0 100644 --- a/.gitignore +++ b/.gitignore @@ -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 diff --git a/playbooks/deploy-omnivoice.yaml b/playbooks/deploy-omnivoice.yaml index 8e1aa38..bcb2348 100644 --- a/playbooks/deploy-omnivoice.yaml +++ b/playbooks/deploy-omnivoice.yaml @@ -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 diff --git a/stacks/omnivoice/Dockerfile b/stacks/omnivoice/Dockerfile index 5e9a7b7..3d1df6c 100644 --- a/stacks/omnivoice/Dockerfile +++ b/stacks/omnivoice/Dockerfile @@ -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 diff --git a/stacks/omnivoice/README.md b/stacks/omnivoice/README.md index 7a7b5ac..48bc69b 100644 --- a/stacks/omnivoice/README.md +++ b/stacks/omnivoice/README.md @@ -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 (`` 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 diff --git a/stacks/omnivoice/app.py b/stacks/omnivoice/app.py index 2152af8..a1e7bc4 100644 --- a/stacks/omnivoice/app.py +++ b/stacks/omnivoice/app.py @@ -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=) and/or voice-DESIGN (instruct=) + - clone (voice=) and/or voice-DESIGN (instruct=) - 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": [, ...]} - GET /v1/audio/languages -> {"languages": ["Auto", , ...]} - POST /v1/audio/speech -> audio/wav + GET /healthz -> readiness (200 once model + >=1 voice are loaded) + GET /v1/audio/voices -> {"voices": [, ...]} + GET /v1/audio/languages -> {"languages": ["Auto", , ...]} + GET /v1/audio/instruct-items -> {"instruct_items": [, ...]} + 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.5–1.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(" 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(" 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) diff --git a/stacks/omnivoice/sanitize.py b/stacks/omnivoice/sanitize.py new file mode 100644 index 0000000..f12b82d --- /dev/null +++ b/stacks/omnivoice/sanitize.py @@ -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: …, …, etc. +# Tag + content both removed (the content is private chain-of-thought, not speech). +_THINK_RE = re.compile( + r"<(think|thinking|reasoning|reflection|scratchpad)>.*?", + 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"|" # + r"|<(?:pad|unk|mask|bos|eos|sep|cls)>" # common HF special tokens + r"|\[/?INST\]" # [INST] [/INST] + r"|<>", # <> <> + 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()) diff --git a/stacks/omnivoice/scheduler.py b/stacks/omnivoice/scheduler.py new file mode 100644 index 0000000..9ce0b46 --- /dev/null +++ b/stacks/omnivoice/scheduler.py @@ -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