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:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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.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("<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)
|
||||
|
||||
@@ -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"!\[[^\]]*\]\([^)]*\)") #  → 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())
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user