"""Client for the SGLang-Omni OpenAI-compatible TTS server. SGLang-Omni (``sgl-omni serve --model-path ``) hosts one TTS model per process behind the OpenAI-style ``/v1/audio/speech`` endpoint. This client speaks that endpoint for every catalog model, resolving the request shape from the model's voice capability (``backends.sglomni. catalog``): speaker the voice names a preset shipped with the model (Qwen3-TTS CustomVoice speakers; Voxtral preset voices) clone the voice comes from a reference clip sent per request as ``ref_audio`` + ``ref_text``. The reference is transcribed with a local Whisper backend when no transcript is given (the qwen backend's flow). On a loopback server the clip travels as a local path the server reads directly; anywhere else it is inlined as a base64 data URL, so ``--api-url`` remote servers work without any server-side file setup. design the voice is described by instructions (``task_type= "VoiceDesign"`` + ``instructions``, Qwen3-TTS VoiceDesign). Clone-capable models without a reference synthesize their built-in default voice ("default") unless the catalog marks a reference as mandatory (Qwen3-TTS Base, dots.tts, ZONOS2 — those refuse at connect). """ import base64 import json import logging import re import tempfile import urllib.error import urllib.parse import urllib.request from pathlib import Path from typing import List, Optional from .. import config from ..audio import concat_audio_files from ..chunking import split_into_chunks from .base import (BaseTTSClient, ConversionCancelled, NonRetryableTTSError, resolve_request_seed) from .languages import normalize_language logger = logging.getLogger(__name__) # Response formats the endpoint offers; complete WAV files need no # sample-rate handling client-side (the header carries it, and models # differ: 24 kHz Voxtral/Higgs, 44.1 kHz ZONOS2, 48 kHz MOSS Local). RESPONSE_FORMAT = "wav" # The voice name the server synthesizes with when the request does not # pick a preset or clone a reference. DEFAULT_VOICE = "default" # Mimetypes for inlined reference audio (data URLs), by file suffix. _MIME_BY_SUFFIX = { ".wav": "audio/wav", ".mp3": "audio/mpeg", ".flac": "audio/flac", ".ogg": "audio/ogg", ".aac": "audio/aac", ".m4a": "audio/mp4", ".webm": "audio/webm", ".mp4": "audio/mp4", } # Error-envelope types the server returns for deterministic request # problems (bad voice, missing reference, unknown model): the identical # request fails on every retry, so the chunk loop gives up immediately. _NON_RETRYABLE_TYPES = ("BadRequestError", "InvalidRequestError", "NotFoundError", "PermissionDeniedError") # The scheduler's KV-window admission error ("Request requires more tokens # than the thinker KV cache can hold (input_tokens=684, max_new_tokens= # 12288, required_tokens=12972, kv_capacity=4095)..."): the server names # the numbers a refit needs, and upstream classifies the message as a # deterministic bad request — the identical request fails on every retry, # so the only useful response is to send a smaller one. _KV_ADMISSION_MARKER = "thinker KV cache can hold" # Frames kept below the capacity the server reported, and the smallest # refitted cap worth generating with (~14 s of speech at 75 fps): below # the floor the request would truncate almost immediately, so the run # surfaces guidance instead of near-empty audio. _KV_FIT_MARGIN = 64 _KV_FIT_FLOOR = 1024 def _is_loopback(url: str) -> bool: """True when URL's host is this machine (the server can read local reference files by path).""" try: host = urllib.parse.urlsplit(url).hostname or "127.0.0.1" except ValueError: return False return host in ("127.0.0.1", "localhost", "::1") def _data_url(path: Path) -> str: """PATH's audio bytes as a base64 data URL (for remote servers).""" mime = _MIME_BY_SUFFIX.get(path.suffix.lower(), "audio/wav") encoded = base64.b64encode(path.read_bytes()).decode("ascii") return f"data:{mime};base64,{encoded}" def _http_error_detail(exc: urllib.error.HTTPError) -> str: """The error response body as text (empty when it cannot be read).""" try: return exc.read().decode("utf-8", errors="replace") except Exception: return "" def _kv_error_number(detail: str, name: str) -> Optional[int]: """The integer NAME=... reports in a KV-window admission message.""" match = re.search(rf"\b{name}=(\d+)", detail) return int(match.group(1)) if match else None class SgOmniTTSClient(BaseTTSClient): """Generates audio chunks through an SGLang-Omni server.""" def __init__(self, chunks_dir: Path, model: Optional[str] = None, voice: Optional[str] = None, ref_audio: Optional[str] = None, ref_text: Optional[str] = None, skip_transcription: bool = False, instructions: Optional[str] = None, language: Optional[str] = None, api_url: Optional[str] = None, chunk_size: Optional[int] = None, quiet: bool = False, cancel=None): super().__init__(chunks_dir, quiet=quiet, cancel=cancel) # The catalog entry this run targets (the backend package validates # the key; only its repo id and capability are client business). from backends.sglomni.catalog import entry_by_key from backends.sglomni.constants import DEFAULT_PORT, SERVER_NAME from backends.common import port_of self.entry = entry_by_key((model or "").strip()) if self.entry is None: raise RuntimeError( f"Unknown SGLang-Omni model {model!r} — pick a catalog key " "(see Configure Backends → SGLang-Omni or the backend docs).") self.api_url = ((api_url or config.SGLOMNI_API_URL).strip() .rstrip("/")) self.port = port_of(self.api_url, DEFAULT_PORT) self.voice = (voice or "").strip() or None self.ref_audio = (ref_audio or "").strip() or None self.ref_text = (ref_text or "").strip() self.skip_transcription = skip_transcription self.instructions = (instructions or "").strip() # Per-run sub-request word cap (the pre-flight chunk popup's "set # chunk" answer); None follows config.CHUNK_SIZE. self.chunk_size = chunk_size # A KV-window capacity the server taught us via an admission # rejection (None = none learned): later requests keep their # max_new_tokens under it. See _kv_admission_fit. self._kv_fit = None # Seed sent with every request: config.SEED as-is, or (with # CONSTANT_SEED and SEED < 0) one random value drawn per run and # reused for every chunk so the voice stays consistent across # chunk boundaries. Only sent to models that accept a # request-scoped seed (Voxtral rejects it outright), and only # when a concrete seed is in play (a negative one means "re-sample # every generation", so there is nothing to send). seed = resolve_request_seed() if self.entry.supports_seed else None self._seed = seed if (seed is not None and seed >= 0) else None if language is None: language = config.LANGUAGE self.language = normalize_language(language) self._check_connect_inputs() self._connect() # ------------------------------------------------------------------ # Connection # ------------------------------------------------------------------ def _check_connect_inputs(self) -> None: """Validate the voice inputs against the model's capability.""" entry = self.entry if entry.capability == "design" and not self.instructions: raise RuntimeError( f"{entry.label} designs the voice from an instruction: " 'pass --instructions "..." describing the voice.') if entry.capability == "clone" and entry.requires_reference \ and not self.ref_audio: raise RuntimeError( f"{entry.label} requires reference audio to narrate: " "pass --clone PATH (a .wav reference clip), or pick a " "model that synthesizes without one.") if entry.capability == "speaker" and self.ref_audio: self._report(f"[WARNING] --clone is ignored with {entry.label}: " "it voices text with its built-in presets.") self.ref_audio = None elif self.ref_audio and not Path(self.ref_audio).is_file(): raise RuntimeError( f"Reference audio not found: {self.ref_audio}") if entry.speakers and self.voice \ and self.voice not in entry.speakers: self._report( f"[WARNING] Voice {self.voice!r} is not one of " f"{entry.label}'s presets ({', '.join(entry.speakers)}); " "the server will reject it if it does not know the name.") def _connect(self) -> None: """Verify the server is up, healthy, and hosting the expected model. The managed-server lifecycle (managed.ensure_running / the run view's autostart) normally boots exactly the selected model; a foreign server hosting something else — or a remote one the form could not classify — fails here with both model names instead of producing per-chunk failures later. """ entry, url = self.entry, self.api_url try: payload = self._fetch_json("/health", timeout=10) except Exception as exc: raise RuntimeError( f"SGLang-Omni server not reachable at {url}: {exc}. Start " "the sgl-omni server first (the CLI and the hub start the " "managed instance automatically when the backend is " "installed), or point --api-url at a running server." ) from exc if not isinstance(payload, dict) \ or payload.get("status") != "healthy": raise RuntimeError( f"The SGLang-Omni server at {url} is not healthy yet " f"(health: {payload}). Wait for it to finish booting and " "retry.") served = self._served_model() if served is not None and served != entry.repo: raise RuntimeError( f"The SGLang-Omni server at {url} hosts {served}, but " f"this run selected {entry.repo}. Restart it with that " "model (the managed server restarts automatically), or " "pick the hosted model for this run.") self._resolve_reference_text() mode = {"speaker": "built-in presets", "clone": "voice cloning", "design": "voice design"}[entry.capability] self._report(f"[OK] Connected to SGLang-Omni at {url} " f"({entry.label}, {mode})") def _served_model(self) -> Optional[str]: """The repo id the server hosts (None when it cannot be read).""" try: payload = self._get_json("/v1/models", timeout=10) except Exception: return None entries = (payload or {}).get("data") if isinstance(entries, list) and entries \ and isinstance(entries[0], dict): return entries[0].get("id") return None def _resolve_reference_text(self) -> None: """Resolve the clone reference transcript: explicit text, then a local Whisper transcription.""" if self.entry.capability != "clone" or not self.ref_audio: return if not self.ref_text and not self.skip_transcription: self._report("[INFO] Transcribing reference audio for voice " "cloning...") from .transcribe import transcribe_reference_audio self.ref_text = transcribe_reference_audio(self.ref_audio) or "" if self.ref_text: self._report(f"[OK] Reference text: {self.ref_text}") else: self._report("[WARNING] No reference transcript: cloning runs " "without ref_text, which lowers quality for " "models that use it. Pass --transcription \"...\" " "for best results.") # ------------------------------------------------------------------ # HTTP requests # ------------------------------------------------------------------ def _fetch_json(self, path: str, timeout: int = 10) -> dict: """GET PATH and parse the JSON body, raising on connection errors. Unlike _get_json this surfaces unreachable servers to the caller — the connect flow needs to tell "nothing is listening" (start the server) apart from "listening but still booting" (wait). """ url = f"{self.api_url}{path}" with urllib.request.urlopen(url, timeout=timeout) as response: payload = json.loads(response.read().decode("utf-8")) return payload if isinstance(payload, dict) else {} def _get_json(self, path: str, timeout: int = 10) -> Optional[dict]: """GET PATH and parse a JSON object, or None on any error.""" try: return self._fetch_json(path, timeout=timeout) except (OSError, ValueError): return None def _ref_audio_value(self) -> str: """The ref_audio request value: a local path on a loopback server (the server reads the file directly), else a base64 data URL.""" path = Path(self.ref_audio) if not path.is_file(): raise RuntimeError( f"Reference audio not found: {self.ref_audio}") if _is_loopback(self.api_url): return str(path.resolve()) return _data_url(path) def _request_payload(self, text: str) -> dict: """The /v1/audio/speech JSON body for one sub-chunk.""" entry = self.entry payload = { "model": entry.repo, "voice": self.voice or DEFAULT_VOICE, "input": text, "response_format": RESPONSE_FORMAT, "language": self.language, } if self._seed is not None: payload["seed"] = self._seed if entry.max_new_tokens is not None: # Models whose engine caps a request below what a full # sub-chunk can narrate (Zonos2's 1024-frame default is ~12 s): # raise the ceiling per request. Generation still stops at # natural EOS, so an unused margin costs nothing. A capacity # learned from an admission rejection (Higgs pins the window) # keeps later requests under it too. cap = entry.max_new_tokens if self._kv_fit is not None: cap = min(cap, self._kv_fit) payload["max_new_tokens"] = cap if entry.capability == "design": payload["task_type"] = "VoiceDesign" payload["instructions"] = self.instructions elif entry.capability == "clone" and self.ref_audio: payload["ref_audio"] = self._ref_audio_value() if self.ref_text: payload["ref_text"] = self.ref_text return payload def _request_wav(self, text: str) -> bytes: """POST one sub-chunk and return the complete WAV bytes.""" payload = self._request_payload(text) try: return self._post_speech(payload) except urllib.error.HTTPError as exc: # A KV-window rejection is deterministic (upstream maps it to a # bad request): refit the generation cap to the capacity the # server reported and resend once before surfacing anything. detail = _http_error_detail(exc) fitted = (self._kv_admission_fit(detail) if "max_new_tokens" in payload else None) if fitted is not None: try: return self._post_speech( dict(payload, max_new_tokens=fitted)) except urllib.error.HTTPError as retry_exc: exc = retry_exc detail = _http_error_detail(exc) raise self._request_error(exc.code, detail) from exc except urllib.error.URLError as exc: raise RuntimeError( f"SGLang-Omni request failed: {exc.reason}") from exc def _post_speech(self, payload: dict) -> bytes: """POST PAYLOAD to /v1/audio/speech; HTTPErrors propagate raw.""" url = f"{self.api_url}/v1/audio/speech" request = urllib.request.Request( url, data=json.dumps(payload).encode("utf-8"), headers={"Content-Type": "application/json"}, method="POST") try: with urllib.request.urlopen(request, timeout=config.API_TIMEOUT) as response: wav = response.read() except urllib.error.HTTPError: # Re-raise raw (HTTPError subclasses URLError): the caller maps # it — and refits KV-window rejections — from the status code. raise except urllib.error.URLError as exc: raise RuntimeError( f"SGLang-Omni request failed: {exc.reason}") from exc if not wav: raise RuntimeError("SGLang-Omni server returned empty audio") return wav def _kv_admission_fit(self, detail: str) -> Optional[int]: """A refitted max_new_tokens for a KV-window rejection, or None. DETAIL is the error response body. The server's message names the request's prompt length and the KV window it must fit; the refit keeps a small margin below the window, is remembered for this client's remaining sub-requests, and the refitted request carries it. When the window leaves less than a useful minimum after the prompt (a very long reference clip), the run fails with guidance instead of near-empty audio. """ if _KV_ADMISSION_MARKER not in detail: return None input_tokens = _kv_error_number(detail, "input_tokens") kv_capacity = _kv_error_number(detail, "kv_capacity") if input_tokens is None or kv_capacity is None: return None fitted = kv_capacity - input_tokens - _KV_FIT_MARGIN if fitted < _KV_FIT_FLOOR: raise NonRetryableTTSError( f"SGLang-Omni rejected the request: the model's KV window " f"({kv_capacity} tokens) leaves {fitted} frames after this " f"request's prompt ({input_tokens} tokens) — too little to " "narrate anything useful. Use a shorter reference clip or " "a smaller Chunk Size setting; the server caps prompt plus " "generation at that window for every request.") if self._kv_fit is not None: fitted = min(fitted, self._kv_fit) self._kv_fit = fitted return fitted def _request_error(self, status: int, detail: str) -> Exception: """Map the OpenAI-style error envelope to the retry decision. A 4xx envelope (BadRequestError et al.) is deterministic — the identical request fails identically on every attempt — so it surfaces as NonRetryableTTSError and the chunk loop aborts with the server's message; anything else stays retryable. """ message = detail[:500] or f"HTTP {status}" kind = None try: envelope = json.loads(detail) error = envelope.get("error") if isinstance(error, dict): message = str(error.get("message") or message) kind = error.get("type") except ValueError: pass if 400 <= status < 500 and (kind is None or kind in _NON_RETRYABLE_TYPES): return NonRetryableTTSError( f"SGLang-Omni rejected the request (HTTP {status}): " f"{message}") return RuntimeError( f"SGLang-Omni server returned HTTP {status}: {message}") # ------------------------------------------------------------------ # Chunk generation # ------------------------------------------------------------------ def generate_chunk(self, text: str, chunk_num: int) -> Optional[str]: """Generate one audio chunk; returns its path in the chunks folder. The text is split into sub-requests of at most ``chunk_size`` words each — the per-run cap the pre-flight chunk popup sets for models whose engine cannot narrate a full CHUNK_SIZE sub-chunk (Higgs), else ``config.CHUNK_SIZE`` (the book-level chunker normally guarantees this already; the split is defense in depth against pathological input such as a punctuation-free run of text) — and the returned WAV files are concatenated into one chunk file. """ try: sub_chunks = split_into_chunks( text, max_words=self.chunk_size or config.CHUNK_SIZE) if not sub_chunks: raise RuntimeError("No text to synthesize") with self._chunk_heartbeat(chunk_num): wav_parts: List[bytes] = [ self._request_wav(sub_text) for sub_text in sub_chunks] output_path = self._chunk_path(chunk_num, ".wav") if len(wav_parts) == 1: output_path.write_bytes(wav_parts[0]) else: # Several sub-request WAVs: concatenate through the shared # ffmpeg path (each part is a complete file with headers). with tempfile.TemporaryDirectory( prefix="sglomni_parts_") as parts_dir: part_paths: List[Path] = [] for index, wav in enumerate(wav_parts, 1): part = Path(parts_dir) / f"part_{index:02d}.wav" part.write_bytes(wav) part_paths.append(part) concat_audio_files(part_paths, output_path) logger.debug("Chunk %d generated (%d sub-request(s))", chunk_num, len(wav_parts)) return str(output_path) except ConversionCancelled: raise except Exception as exc: logger.error("SGLang-Omni chunk processing failed for chunk " "%d: %s", chunk_num, exc) return None