fix: ZONOS2 only generating 1024 tokens max

4 files changed, 29 insertions, 1 deletions

diff --git a/app/backends/sglomni/catalog.py b/app/backends/sglomni/catalog.py
index c89e865..5ff0c81 100644
--- a/app/backends/sglomni/catalog.py
+++ b/app/backends/sglomni/catalog.py
@@ -60,6 +60,16 @@ class ModelEntry:
     fp8_moe: bool = False
     fp8_min_compute_capability: Optional[Tuple[int, int]] = None
     bf16_config: Optional[str] = None
+    # max_new_tokens sent with every /v1/audio/speech request (None = send
+    # nothing and leave the server default in charge). Zonos2's AR engine
+    # defaults to 1024 audio frames (44100 Hz / 512 hop = 86.13 fps), which
+    # caps one request at ~12 s of speech and silently truncates longer
+    # text. 12288 frames (~143 s) covers a full 250-word CHUNK_SIZE
+    # sub-chunk (~8600 frames) and stays under the KV-pool admission check
+    # on every GPU that can host the model (all are >= 24 GB with
+    # >= 28083-token pools; the scheduler rejects prompt + max_new_tokens
+    # above it, and zonos2 requests are not auto-clamped).
+    max_new_tokens: Optional[int] = None
 
 
 # The Qwen3-TTS CustomVoice speaker table — the same built-in speakers the
@@ -197,6 +207,7 @@ ENTRIES: Tuple[ModelEntry, ...] = (
         fp8_moe=True,
         fp8_min_compute_capability=(8, 9),
         bf16_config="zonos2_bf16.yaml",
+        max_new_tokens=12288,
     ),
 )
 
diff --git a/app/converter/clients/sglomni.py b/app/converter/clients/sglomni.py
index 8484aba..d4b756f 100644
--- a/app/converter/clients/sglomni.py
+++ b/app/converter/clients/sglomni.py
@@ -274,6 +274,12 @@ class SgOmniTTSClient(BaseTTSClient):
         }
         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
diff --git a/app/docs/backend-sglomni.md b/app/docs/backend-sglomni.md
index 13c8876..2aac93e 100644
--- a/app/docs/backend-sglomni.md
+++ b/app/docs/backend-sglomni.md
@@ -66,7 +66,7 @@ inside the start timeout.
 | `voxtral_tts` | Voxtral TTS 4B | preset named voices | e.g. `default`, `casual_male` |
 | `dots_tts_mf` | dots.tts (MeanFlow) | clone (reference required) | |
 | `fish_s2_pro` | Fish Speech S2-Pro | default voice or clone | needs ~24 GB VRAM (known OOM on a single RTX 3090, upstream issue #359) |
-| `zonos2` | ZONOS2 | clone (reference required) | 44.1 kHz; FP8 pipeline falls back to bf16 on GPUs below compute capability 8.9 |
+| `zonos2` | ZONOS2 | clone (reference required) | 44.1 kHz; FP8 pipeline falls back to bf16 on GPUs below compute capability 8.9; requests carry `max_new_tokens=12288` because the engine's 1024-frame default caps one request at ~12 s of speech |
 
 In the hub's **Generate Audiobooks** form the Model picker reads as a table,
 like the audio.cpp one: each entry's label is padded to the widest one and its
diff --git a/app/tests/test_tts_sglomni.py b/app/tests/test_tts_sglomni.py
index 2dee364..58ca57c 100644
--- a/app/tests/test_tts_sglomni.py
+++ b/app/tests/test_tts_sglomni.py
@@ -248,6 +248,17 @@ class PayloadTests(unittest.TestCase):
         client._seed = 7
         self.assertEqual(client._request_payload("Hello.")["seed"], 7)
 
+    def test_zonos2_payload_raises_the_generation_cap(self):
+        """Zonos2's 1024-frame engine default caps a request at ~12 s."""
+        client = self._make_client("zonos2")
+        payload = client._request_payload("Hello.")
+        self.assertEqual(payload["max_new_tokens"], 12288)
+
+    def test_models_without_a_cap_send_no_max_new_tokens(self):
+        client = self._make_client("higgs_audio_v3_tts")
+        self.assertNotIn("max_new_tokens",
+                         client._request_payload("Hello."))
+
 
 class RequestErrorTests(unittest.TestCase):
     """OpenAI-style error envelopes decide retryability."""