fix: warn before using a likely too-big chunk size for sglang-omni models

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")