"""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 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") 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}" 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, 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() # 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. payload["max_new_tokens"] = entry.max_new_tokens 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.""" url = f"{self.api_url}/v1/audio/speech" payload = json.dumps(self._request_payload(text)).encode("utf-8") request = urllib.request.Request( url, data=payload, 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 as exc: detail = "" try: detail = exc.read().decode("utf-8", errors="replace") except Exception: pass 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 if not wav: raise RuntimeError("SGLang-Omni server returned empty audio") return wav 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 ``config.CHUNK_SIZE`` words each (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=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