"""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, MOSS-TTS, 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", } # 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 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.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 # The ref_audio request value, computed on first use (see # _ref_audio_value); None = not computed yet. self._ref_audio_cached = 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). NOTE(unverified # upstream): whether the other pipelines accept a seed too — see # the catalog's supports_seed note. 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 entry.capability == "design" and self.ref_audio: self._report(f"[WARNING] --clone is ignored with {entry.label}: " "it designs the voice from instructions.") 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}") presets = self._preset_voices() if presets and self.voice and self.voice not in presets: self._report( f"[WARNING] Voice {self.voice!r} is not one of " f"{entry.label}'s presets ({', '.join(presets)}); " "the server will reject it if it does not know the name.") def _preset_voices(self) -> List[str]: """The preset voice names ENTRY can speak with — the same list the hub's voice menu offers (catalog-declared speakers, or the checkpoint's own voice_embedding presets, e.g. Voxtral's).""" from backends.sglomni.models import preset_voices return preset_voices(self.entry) 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 urllib.error.HTTPError as exc: # A booting server answers 503 with an "unhealthy" body — # urlopen turns that into an HTTPError before the healthy # check below can see it. Tell the user to wait for the # server that is already starting, not to start another. detail = _http_error_detail(exc) if exc.code == 503: raise RuntimeError( f"The SGLang-Omni server at {url} is not healthy yet " f"(HTTP 503: {detail[:300] or 'no body'}). Wait for it " "to finish booting and retry.") from exc raise RuntimeError( f"The SGLang-Omni server at {url} answered HTTP {exc.code} " f"on /health ({detail[:300] or 'no body'}). Is this an " "sgl-omni server? 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 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 self.skip_transcription: self._report("[INFO] Skipping reference audio transcription " "(--no-transcription).") elif not self.ref_text: 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. Computed once per run and cached: the clip is validated at connect and cannot change mid-run, and re-encoding its bytes for every sub-request would ship the same payload over and over.""" cached = self._ref_audio_cached if cached is None: 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): cached = str(path.resolve()) else: cached = _data_url(path) self._ref_audio_cached = cached return cached def _request_payload(self, text: str) -> dict: """The /v1/audio/speech JSON body for one sub-chunk.""" entry = self.entry payload = { "model": entry.repo, # NOTE(unverified upstream): "voice" is sent even when nothing # was picked (the "default" sentinel) and to design runs, # which have no voice — audio.cpp omits the field there. # Verify the server tolerates it for every pipeline. "voice": self.voice or DEFAULT_VOICE, "input": text, "response_format": RESPONSE_FORMAT, # NOTE(unverified upstream): Qwen-style display names # ("English", "Auto") go to every model; audio.cpp maps per # family. Verify each pipeline accepts them (or wants ISO # codes / the field omitted). "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 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 len(wav) < 12 or wav[:4] != b"RIFF" or wav[8:12] != b"WAVE": # A JSON error body handed back with HTTP 200 would otherwise # be written as chunk bytes and fail later, confusingly, in # the concat step. raise RuntimeError( "SGLang-Omni server returned audio that is not a WAV file") 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}" try: envelope = json.loads(detail) error = envelope.get("error") if isinstance(error, dict): message = str(error.get("message") or message) except ValueError: pass # Every 4xx envelope is deterministic — the identical request # fails identically on every attempt (this is a single-user local # server: it queues work rather than answering 429-style limits). if 400 <= status < 500: 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") output_path = self._chunk_path(chunk_num, ".wav") with tempfile.TemporaryDirectory( prefix="sglomni_parts_") as parts_dir: # One part per sub-request, spooled to disk as it arrives # (like the other clients) instead of buffering every # response in memory until the chunk is complete. part_paths: List[Path] = [] with self._chunk_heartbeat(chunk_num): for index, sub_text in enumerate(sub_chunks, 1): part = Path(parts_dir) / f"part_{index:02d}.wav" part.write_bytes(self._request_wav(sub_text)) part_paths.append(part) if len(part_paths) == 1: output_path.write_bytes(part_paths[0].read_bytes()) else: # Several sub-request WAVs: concatenate through the # shared ffmpeg path (each part is a complete file # with headers). concat_audio_files(part_paths, output_path) logger.debug("Chunk %d generated (%d sub-request(s))", chunk_num, len(part_paths)) return str(output_path) except ConversionCancelled: raise except NonRetryableTTSError: # Propagate past the generic handler so the retry loop skips # its remaining attempts for deterministic server errors. raise except Exception as exc: logger.error("SGLang-Omni chunk processing failed for chunk " "%d: %s", chunk_num, exc) return None