diff options
Diffstat (limited to 'app/converter/clients')
| -rw-r--r-- | app/converter/clients/sglomni.py | 127 |
1 files changed, 112 insertions, 15 deletions
diff --git a/app/converter/clients/sglomni.py b/app/converter/clients/sglomni.py index d4b756f..14c9990 100644 --- a/app/converter/clients/sglomni.py +++ b/app/converter/clients/sglomni.py @@ -26,6 +26,7 @@ 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 @@ -64,6 +65,21 @@ _MIME_BY_SUFFIX = { _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 @@ -82,6 +98,20 @@ def _data_url(path: Path) -> str: 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.""" @@ -94,6 +124,7 @@ class SgOmniTTSClient(BaseTTSClient): 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 @@ -114,6 +145,13 @@ class SgOmniTTSClient(BaseTTSClient): 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 @@ -278,8 +316,13 @@ class SgOmniTTSClient(BaseTTSClient): # 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 + # 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 @@ -291,22 +334,42 @@ class SgOmniTTSClient(BaseTTSClient): 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" - payload = json.dumps(self._request_payload(text)).encode("utf-8") request = urllib.request.Request( - url, data=payload, + 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 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.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 @@ -314,6 +377,37 @@ class SgOmniTTSClient(BaseTTSClient): 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. @@ -347,14 +441,17 @@ class SgOmniTTSClient(BaseTTSClient): 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 + 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 + 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) + 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") |