mARTin-B78/tts-dgx-spark-faster-qwen3-tts
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
119b92ea15
tts-dgx-spark-faster-qwen3-tts/patches/openai_server.patch
mARTin-B78 ccccdf820a feat: graph-decoded streaming, prebuffer_ms and cancel-on-disconnect (v6.13) 
- Fixed-window CUDA-graph codec decode for streaming (TTS_GRAPH_DECODE=0 disables)
- Per-request chunk_size (2-24) on /v1/audio/speech
- prebuffer_ms (0-10000) to hold back audio before sending
- Stop generation on client disconnect
- Recommend --max-seq-len 2048 for the streaming container

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-10-04 00:11:12 +02:00

571 lines
24 KiB
Diff
Raw Blame History

diff --git a/examples/openai_server.py b/examples/openai_server.py
index 2199e14..3f3a4e6 100644
--- a/examples/openai_server.py
+++ b/examples/openai_server.py
@@ -36,11 +36,13 @@ API usage:
"""
import argparse
import asyncio
+import hashlib
import io
import json
import logging
import os
import queue
+import re
import struct
import sys
import threading
@@ -66,9 +68,67 @@ app = FastAPI(title="faster-qwen3-tts OpenAI-compatible API")
tts_model = None
voices: dict = {}
+voices_file_path: Optional[str] = None
+last_voices_mtime: float = 0.0
default_voice: Optional[str] = None
SAMPLE_RATE = 24000 # updated once the model loads
+MAX_SEQ_LEN = 4096 # CUDA graph static cache size, updated from --max-seq-len
_model_lock = threading.Lock() # prevent concurrent GPU inference
+aligner_model = None
+
+def _get_aligner():
+ global aligner_model
+ if aligner_model is None:
+ try:
+ from qwen_asr import Qwen3ForcedAligner
+ import torch
+ except ImportError:
+ raise HTTPException(status_code=500, detail="qwen-asr is not installed. Run: pip install qwen-asr")
+ logger.info("Loading Qwen3-ForcedAligner-0.6B...")
+ aligner_model = Qwen3ForcedAligner.from_pretrained(
+ "Qwen/Qwen3-ForcedAligner-0.6B",
+ dtype=torch.bfloat16,
+ device_map="cuda"
+ )
+ logger.info("Aligner loaded.")
+ return aligner_model
+
+
+def _voice_seed(voice_name: str) -> int:
+ """Return a stable per-voice seed derived from the voice name.
+
+ Used as the default when no explicit 'seed' is set in voices.json.
+ MD5 is used only for its stable byte output — not for security.
+ """
+ return int(hashlib.md5(voice_name.encode()).hexdigest(), 16) % (2 ** 31)
+
+
+def _seed_rng(seed: int) -> None:
+ """Seed PyTorch CPU and CUDA RNGs for reproducible sampling."""
+ torch.manual_seed(seed)
+ if torch.cuda.is_available():
+ torch.cuda.manual_seed_all(seed)
+
+
+def _merge_instruct(voice_cfg: dict, req_instruct: Optional[str]) -> Optional[str]:
+ """Combine a voice's identity instruct with a per-request delivery direction.
+
+ A voice's instruct in voices.json describes WHO is speaking; an instruct
+ sent with the request describes HOW this one line is delivered. Replacing
+ the former with the latter would re-roll the character from a few words of
+ direction, so combine them, identity first — the same rule the VoiceDesign
+ server applies.
+
+ Instructions are only followed reliably in ICL mode. The .pt prompts written
+ by extract_embeddings.py are full ICL prompts (ref codec tokens + ref_text),
+ not x-vector-only, so this path qualifies.
+ """
+ base = str(voice_cfg.get("instruct") or "").strip()
+ line = str(req_instruct or "").strip()
+ if base and line and line != base:
+ return f"{base} {line}"
+ return line or base or None
+
# ---------------------------------------------------------------------------
# Request / response models
@@ -79,8 +139,68 @@ class SpeechRequest(BaseModel):
model: str = "tts-1"
input: str
voice: str = "alloy"
- response_format: str = "wav" # wav | pcm | mp3
- speed: float = 1.0 # accepted but not yet applied
+ response_format: str = "wav" # wav | pcm | mp3 | zip
+ speed: float = 1.0 # scales audio tempo
+ instruct: Optional[str] = None # per-line delivery direction, see _merge_instruct
+ # Streaming only: codec frames per streamed chunk (12 = 1 s). Small values
+ # cut time-to-first-audio but cost throughput; clamped to 2..24.
+ # None keeps the voice's voices.json chunk_size.
+ chunk_size: Optional[int] = None
+ # Streaming only: hold back the first N ms of audio before sending, so a
+ # client that plays on arrival rides out a slower-than-realtime start
+ # (shared GPU) instead of stuttering. 0/None = send immediately. Max 10 s.
+ prebuffer_ms: Optional[int] = None
+
+
+# ---------------------------------------------------------------------------
+# Text chunking for long input
+#
+# A single generate call is capped at MAX_SEQ_LEN codec tokens (the CUDA
+# graph's static cache size). Rather than raising that cache to cover the
+# longest text anyone might send — which burns VRAM permanently for every
+# request — long input is split into paragraph/sentence-sized segments that
+# each fit comfortably under the cap, generated back-to-back, and stitched
+# into one continuous stream/file.
+# ---------------------------------------------------------------------------
+
+_CODEC_HZ = 12 # talker codec frame rate (generated audio tokens/sec)
+_CHARS_PER_SEC = 15.0 # conservative estimate of spoken characters/sec
+_SEQ_LEN_SAFETY = 0.6 # stay well under the hard cap to leave room for slower speech
+_SENTENCE_SPLIT_RE = re.compile(r"(?<=[.!?])\s+")
+
+
+def _max_chars_per_segment() -> int:
+ budget_seconds = (MAX_SEQ_LEN * _SEQ_LEN_SAFETY) / _CODEC_HZ
+ return max(200, int(budget_seconds * _CHARS_PER_SEC))
+
+
+def _split_text_for_generation(text: str, max_chars: int) -> list:
+ """Split text into segments that each fit under max_chars.
+
+ Splits on paragraph breaks first (kept as separate segments so natural
+ pauses land there), then falls back to sentence boundaries for any
+ paragraph that alone exceeds the budget.
+ """
+ paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
+ if not paragraphs:
+ return [text]
+
+ segments = []
+ for para in paragraphs:
+ if len(para) <= max_chars:
+ segments.append(para)
+ continue
+ buf = ""
+ for sent in _SENTENCE_SPLIT_RE.split(para):
+ candidate = f"{buf} {sent}".strip() if buf else sent
+ if len(candidate) > max_chars and buf:
+ segments.append(buf)
+ buf = sent
+ else:
+ buf = candidate
+ if buf:
+ segments.append(buf)
+ return segments or [text]
# ---------------------------------------------------------------------------
@@ -144,9 +264,39 @@ def _to_mp3_bytes(pcm: np.ndarray, sample_rate: int) -> bytes:
def resolve_voice(voice_name: str) -> dict:
- """Return voice config dict or fall back to default, else raise 400."""
+ """Return voice config dict, accepting a bare language code as an alias."""
+ global voices, last_voices_mtime
+ voice_name = voice_name.strip()
+
+ # Hot-reload voices.json if it was modified
+ if voices_file_path and os.path.exists(voices_file_path):
+ try:
+ current_mtime = os.path.getmtime(voices_file_path)
+ if current_mtime > last_voices_mtime:
+ with open(voices_file_path, "r", encoding="utf-8") as f:
+ voices = json.load(f)
+ last_voices_mtime = current_mtime
+ logger.info("Hot-reloaded %d voices from %s", len(voices), voices_file_path)
+ except Exception as e:
+ logger.warning("Failed to hot-reload voices.json: %s", e)
+
if voice_name in voices:
return voices[voice_name]
+ # OpenAI-compatible clients sometimes send a route alias (e.g. ``FR``,
+ # ``FR_F`` or ``FR_M``) in the voice field. Resolve it to a configured
+ # speaker instead of silently switching to the default English voice.
+ alias = re.fullmatch(r"([A-Z]{2})(?:_([MF]))?", voice_name, re.IGNORECASE)
+ if alias:
+ prefix = alias.group(1).upper() + "_"
+ if alias.group(2):
+ prefix += alias.group(2).upper() + "_"
+ candidates = [name for name in voices if name.upper().startswith(prefix)]
+ if candidates:
+ # For a language-only alias prefer a female voice, then registry order.
+ selected = next((name for name in candidates
+ if not alias.group(2) and name.upper().startswith(prefix + "F_")), candidates[0])
+ logger.warning("Voice %r is a language alias; using configured voice %r", voice_name, selected)
+ return voices[selected]
if default_voice and default_voice in voices:
logger.warning(
"Voice %r not configured; falling back to default voice %r",
@@ -168,42 +318,214 @@ def resolve_voice(voice_name: str) -> dict:
# ---------------------------------------------------------------------------
-async def _stream_chunks(voice_cfg: dict, text: str) -> AsyncGenerator[bytes, None]:
+def _load_voice_clone_prompt(voice_cfg: dict, voice_name: str, tts_model):
+ spk_emb_path = voice_cfg.get("speaker_embeddings") or voice_cfg.get("speaker embeddings")
+ if spk_emb_path and os.path.isfile(spk_emb_path):
+ try:
+ return torch.load(spk_emb_path, map_location="cpu", weights_only=False)
+ except Exception as e:
+ logger.error("Failed to load speaker embeddings from %s: %s", spk_emb_path, e)
+
+ # Auto-generate if missing
+ ref_audio = voice_cfg.get("ref_audio")
+ if not ref_audio or not os.path.isfile(ref_audio):
+ return None
+
+ logger.info("Precomputing and saving speaker embedding for voice %r...", voice_name)
+ try:
+ ref_text = voice_cfg.get("ref_text", "")
+ # generate prompt using the model's built-in helper
+ prompt_items = tts_model.model.create_voice_clone_prompt(ref_audio, [ref_text])
+ vcp = tts_model.model._prompt_items_to_voice_clone_prompt(prompt_items)
+
+ # save it to the speakers directory
+ # If the file path is already in voice_cfg but doesn't exist, use that, otherwise generate a path
+ if spk_emb_path and not os.path.exists(spk_emb_path) and spk_emb_path.endswith('.pt'):
+ pt_path = spk_emb_path
+ else:
+ pt_path = f"/config/speakers/{voice_name}.pt"
+
+ os.makedirs(os.path.dirname(pt_path), exist_ok=True)
+ torch.save(vcp, pt_path)
+ logger.info("Saved speaker embedding to %s", pt_path)
+
+ # update in memory so future requests skip generating
+ voice_cfg["speaker_embeddings"] = pt_path
+
+ return vcp
+ except Exception as e:
+ logger.error("Failed to precompute speaker embedding: %s", e)
+ return None
+
+
+def _clamp_chunk_size(n: int) -> int:
+ return max(2, min(24, n))
+
+
+# Fixed-window codec decode. Upstream re-decodes the whole reference prompt plus
+# every generated frame for the first chunks (~250 frames on long references)
+# and decodes eagerly (hundreds of small kernel launches). The 12 Hz codec
+# decoder is causal, so decoding the last (context + chunk) frames and keeping
+# the tail is exact; with the reference codes as left context the window is
+# always full, so one CUDA graph per window size covers every chunk.
+_DECODE_CONTEXT = 25 # frames, same as upstream chunked_decode / streaming
+_GRAPH_DECODE = os.environ.get("TTS_GRAPH_DECODE", "1") != "0"
+_decoder_graphs: dict = {} # window frames -> (graph, static_in, static_out)
+
+
+def _decode_window(decoder, window: torch.Tensor) -> torch.Tensor:
+ """Decode [frames, Q] codes to a waveform tensor, graph-replayed when possible."""
+ codes = window.T.unsqueeze(0)
+ entry = _decoder_graphs.get(codes.shape[-1])
+ if entry is None:
+ static_in = codes.clone()
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ for _ in range(2):
+ decoder(static_in)
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph, capture_error_mode="thread_local"):
+ static_out = decoder(static_in)
+ entry = _decoder_graphs[codes.shape[-1]] = (graph, static_in, static_out)
+ graph, static_in, static_out = entry
+ static_in.copy_(codes)
+ graph.replay()
+ return static_out
+
+
+@torch.inference_mode()
+def _voice_clone_stream(chunk_size: int, **kwargs):
+ """generate_voice_clone_streaming with fixed-window (graphed) codec decode.
+
+ Yields float32 numpy chunks. Token generation is upstream's
+ fast_generate_streaming unchanged.
+ """
+ from faster_qwen3_tts.streaming import fast_generate_streaming
+
+ gen_keys = ("temperature", "top_k", "top_p")
+ gen = {k: kwargs.pop(k) for k in gen_keys if k in kwargs}
+ m, talker, config, tie, tam, tth, tpe, ref_codes = tts_model._prepare_generation(**kwargs)
+ decoder = m.speech_tokenizer.model.decoder
+ upsample = decoder.total_upsample
+ window = _DECODE_CONTEXT + chunk_size
+ hist = ref_codes.to(tie.device)[-window:] if ref_codes is not None else None
+ for codes, _timing in fast_generate_streaming(
+ talker=talker, talker_input_embeds=tie, attention_mask=tam,
+ trailing_text_hiddens=tth, tts_pad_embed=tpe, config=config,
+ predictor_graph=tts_model.predictor_graph, talker_graph=tts_model.talker_graph,
+ chunk_size=chunk_size, **gen,
+ ):
+ hist = codes if hist is None else torch.cat([hist, codes])[-window:]
+ if hist.shape[0] == window:
+ wav = _decode_window(decoder, hist)
+ else: # x-vector voice, first chunks: short eager decode
+ wav = decoder(hist.T.unsqueeze(0))
+ yield wav[..., -codes.shape[0] * upsample:].flatten().float().cpu().numpy()
+
+
+async def _stream_chunks(voice_cfg: dict, text: str, voice_name: str, speed: float, instruct: Optional[str] = None, chunk_size: Optional[int] = None, prebuffer_ms: Optional[int] = None) -> AsyncGenerator[bytes, None]:
"""
Run generate_voice_clone_streaming in a background thread and yield
raw PCM bytes for each chunk as they arrive.
"""
q: queue.Queue = queue.Queue()
_DONE = object()
+ cancelled = threading.Event()
def producer():
+ process = None
+ if speed != 1.0:
+ import subprocess
+ cmd = [
+ "ffmpeg", "-y", "-loglevel", "error",
+ "-f", "s16le", "-ar", str(SAMPLE_RATE), "-ac", "1", "-i", "pipe:0",
+ "-filter:a", f"atempo={speed}",
+ "-f", "s16le", "-ar", str(SAMPLE_RATE), "-ac", "1", "pipe:1"
+ ]
+ process = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
+
+ def ffmpeg_reader():
+ try:
+ while True:
+ out = process.stdout.read(4096)
+ if not out:
+ break
+ q.put(out)
+ except Exception as e:
+ q.put(e)
+ finally:
+ q.put(_DONE)
+
+ threading.Thread(target=ffmpeg_reader, daemon=True).start()
+
try:
with _model_lock:
- for chunk, _sr, _timing in tts_model.generate_voice_clone_streaming(
- text=text,
- language=voice_cfg.get("language", "Auto"),
- ref_audio=voice_cfg["ref_audio"],
- ref_text=voice_cfg.get("ref_text", ""),
- chunk_size=voice_cfg.get("chunk_size", 12),
- non_streaming_mode=False,
- ):
- q.put(chunk)
+ _seed_rng(voice_cfg.get("seed", _voice_seed(voice_name)))
+ segments = _split_text_for_generation(text, _max_chars_per_segment())
+ for seg in segments:
+ kwargs = dict(
+ text=seg,
+ language=voice_cfg.get("language", "Auto"),
+ ref_audio=voice_cfg.get("ref_audio"),
+ ref_text=voice_cfg.get("ref_text", ""),
+ instruct=_merge_instruct(voice_cfg, instruct),
+ voice_clone_prompt=_load_voice_clone_prompt(voice_cfg, voice_name, tts_model),
+ non_streaming_mode=True,
+ temperature=voice_cfg.get("temperature", 0.8),
+ top_k=voice_cfg.get("top_k", 50),
+ top_p=voice_cfg.get("top_p", 0.9),
+ )
+ cs = _clamp_chunk_size(chunk_size) if chunk_size is not None else voice_cfg.get("chunk_size", 12)
+ if _GRAPH_DECODE:
+ chunks = _voice_clone_stream(cs, **kwargs)
+ else:
+ chunks = (c for c, _sr, _t in tts_model.generate_voice_clone_streaming(chunk_size=cs, **kwargs))
+ for chunk in chunks:
+ if cancelled.is_set(): # client gone: free the GPU for the next request
+ chunks.close()
+ return
+ raw = _to_pcm16(chunk)
+ if process:
+ process.stdin.write(raw)
+ process.stdin.flush()
+ else:
+ q.put(raw)
except Exception as exc:
q.put(exc)
finally:
- q.put(_DONE)
+ if process:
+ try:
+ process.stdin.close()
+ except Exception:
+ pass
+ else:
+ q.put(_DONE)
thread = threading.Thread(target=producer, daemon=True)
thread.start()
loop = asyncio.get_event_loop()
- while True:
- item = await loop.run_in_executor(None, q.get)
- if item is _DONE:
- break
- if isinstance(item, Exception):
- raise item
- yield _to_pcm16(item)
+ hold = max(0, min(10000, prebuffer_ms or 0)) * SAMPLE_RATE * 2 // 1000 # PCM16 bytes
+ pending = b""
+ try:
+ while True:
+ item = await loop.run_in_executor(None, q.get)
+ if item is _DONE:
+ break
+ if isinstance(item, Exception):
+ raise item
+ if hold:
+ pending += item
+ if len(pending) < hold:
+ continue
+ item, pending, hold = pending, b"", 0
+ yield item
+ if pending:
+ yield pending
+ finally:
+ cancelled.set()
# ---------------------------------------------------------------------------
@@ -220,7 +542,8 @@ async def health():
async def create_speech(req: SpeechRequest):
if tts_model is None:
raise HTTPException(status_code=503, detail="Model not loaded")
- if not req.input.strip():
+ req.input = req.input.strip()
+ if not req.input:
raise HTTPException(status_code=400, detail="'input' text is empty")
voice_cfg = resolve_voice(req.voice)
@@ -230,36 +553,84 @@ async def create_speech(req: SpeechRequest):
"wav": "audio/wav",
"pcm": "audio/pcm",
"mp3": "audio/mpeg",
+ "zip": "application/zip",
}
if fmt not in _CONTENT_TYPES:
raise HTTPException(
status_code=400,
- detail=f"response_format {fmt!r} not supported. Use: wav, pcm, mp3",
+ detail=f"response_format {fmt!r} not supported. Use: wav, pcm, mp3, zip",
)
content_type = _CONTENT_TYPES[fmt]
- # --- MP3: generate all audio, then encode (non-streaming) ---
- if fmt == "mp3":
+ # --- MP3 / ZIP: generate all audio, then encode (non-streaming) ---
+ if fmt in ("mp3", "zip"):
loop = asyncio.get_event_loop()
def _generate():
with _model_lock:
- return tts_model.generate_voice_clone(
- text=req.input,
- language=voice_cfg.get("language", "Auto"),
- ref_audio=voice_cfg["ref_audio"],
- ref_text=voice_cfg.get("ref_text", ""),
- )
-
- audio_arrays, sr = await loop.run_in_executor(None, _generate)
- audio = audio_arrays[0] if audio_arrays else np.zeros(1, dtype=np.float32)
+ _seed_rng(voice_cfg.get("seed", _voice_seed(req.voice)))
+ segments = _split_text_for_generation(req.input, _max_chars_per_segment())
+ pieces = []
+ sr = SAMPLE_RATE
+ for seg in segments:
+ audio_arrays, sr = tts_model.generate_voice_clone(
+ text=seg,
+ language=voice_cfg.get("language", "Auto"),
+ ref_audio=voice_cfg.get("ref_audio"),
+ ref_text=voice_cfg.get("ref_text", ""),
+ instruct=_merge_instruct(voice_cfg, req.instruct),
+ voice_clone_prompt=_load_voice_clone_prompt(voice_cfg, req.voice, tts_model),
+ non_streaming_mode=True,
+ temperature=voice_cfg.get("temperature", 0.8),
+ top_k=voice_cfg.get("top_k", 50),
+ top_p=voice_cfg.get("top_p", 0.9),
+ )
+ if audio_arrays:
+ pieces.append(audio_arrays[0])
+ return pieces, sr
+
+ pieces, sr = await loop.run_in_executor(None, _generate)
+ audio = np.concatenate(pieces) if pieces else np.zeros(1, dtype=np.float32)
+
+ if req.speed != 1.0:
+ import subprocess
+ cmd = [
+ "ffmpeg", "-y", "-loglevel", "error",
+ "-f", "f32le", "-ar", str(sr), "-ac", "1", "-i", "pipe:0",
+ "-filter:a", f"atempo={req.speed}",
+ "-f", "f32le", "-ar", str(sr), "-ac", "1", "pipe:1"
+ ]
+ process = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
+ process.stdin.write(audio.tobytes())
+ process.stdin.close()
+ out = process.stdout.read()
+ audio = np.frombuffer(out, dtype=np.float32)
+
+ if fmt == "zip":
+ def _align():
+ aligner = _get_aligner()
+ res = aligner.align(audio=(audio, sr), text=req.input, language=voice_cfg.get("language", "Auto"))
+ import dataclasses
+ return [dataclasses.asdict(x) for x in res]
+
+ align_data = await loop.run_in_executor(None, _align)
+
+ import zipfile
+ mp3_bytes = _to_mp3_bytes(audio, sr)
+ zip_buf = io.BytesIO()
+ with zipfile.ZipFile(zip_buf, "w", zipfile.ZIP_DEFLATED) as zf:
+ zf.writestr("audio.mp3", mp3_bytes)
+ zf.writestr("timer.json", json.dumps(align_data, ensure_ascii=False))
+
+ return Response(content=zip_buf.getvalue(), media_type=content_type)
+
return Response(content=_to_mp3_bytes(audio, sr), media_type=content_type)
# --- WAV / PCM: stream chunks as they are generated ---
async def audio_stream():
if fmt == "wav":
yield _wav_header(SAMPLE_RATE) # stream with unknown data length
- async for raw_chunk in _stream_chunks(voice_cfg, req.input):
+ async for raw_chunk in _stream_chunks(voice_cfg, req.input, req.voice, req.speed, req.instruct, req.chunk_size, req.prebuffer_ms):
yield raw_chunk
return StreamingResponse(audio_stream(), media_type=content_type)
@@ -306,16 +677,21 @@ def _parse_args():
p.add_argument("--host", default="0.0.0.0", help="Bind host (default: 0.0.0.0)")
p.add_argument("--port", type=int, default=8000, help="Bind port (default: 8000)")
p.add_argument("--device", default="cuda", help="Torch device (default: cuda)")
+ p.add_argument("--max-seq-len", type=int, default=4096, help="Max sequence length for CUDA graph static cache (default: 4096)")
return p.parse_args()
def main():
- global tts_model, voices, default_voice, SAMPLE_RATE
+ global tts_model, voices, voices_file_path, last_voices_mtime, default_voice, SAMPLE_RATE, MAX_SEQ_LEN
args = _parse_args()
+ MAX_SEQ_LEN = args.max_seq_len
# Build voice registry
if args.voices:
+ voices_file_path = args.voices
+ if os.path.exists(args.voices):
+ last_voices_mtime = os.path.getmtime(args.voices)
with open(args.voices) as f:
voices = json.load(f)
default_voice = next(iter(voices))
@@ -344,6 +720,7 @@ def main():
args.model,
device=args.device,
dtype=torch.bfloat16,
+ max_seq_len=args.max_seq_len,
)
SAMPLE_RATE = tts_model.sample_rate
logger.info("Model ready. Sample rate: %d Hz", SAMPLE_RATE)
Reference in New Issue View Git Blame Copy Permalink