diff options
Diffstat (limited to 'app/tests/test_tts_sglomni.py')
| -rw-r--r-- | app/tests/test_tts_sglomni.py | 148 |
1 files changed, 146 insertions, 2 deletions
diff --git a/app/tests/test_tts_sglomni.py b/app/tests/test_tts_sglomni.py index 85327e1..0e3a7de 100644 --- a/app/tests/test_tts_sglomni.py +++ b/app/tests/test_tts_sglomni.py @@ -191,6 +191,7 @@ class PayloadTests(unittest.TestCase): client.instructions = kwargs.get("instructions", "") client.language = "English" client._seed = None + client._kv_fit = None return client def test_speaker_payload_sends_the_preset_name(self): @@ -256,10 +257,19 @@ class PayloadTests(unittest.TestCase): def test_higgs_payload_raises_the_generation_cap(self): """Higgs's 2048-frame engine default caps a request at ~27 s - (75 fps), below a full 250-word sub-chunk.""" + (75 fps); the catalog raises it to the most its admission window + allows (~40 s after the prompt tokens).""" client = self._make_client("higgs_audio_v3_tts") payload = client._request_payload("Hello.") - self.assertEqual(payload["max_new_tokens"], 12288) + self.assertEqual(payload["max_new_tokens"], 3000) + + def test_payload_keeps_a_learned_kv_fit(self): + """A capacity learned from an admission rejection caps later + requests below the catalog value.""" + client = self._make_client("higgs_audio_v3_tts") + client._kv_fit = 2500 + self.assertEqual(client._request_payload("Hello.")["max_new_tokens"], + 2500) def test_models_without_a_cap_send_no_max_new_tokens(self): client = self._make_client("moss_tts") @@ -293,6 +303,126 @@ class RequestErrorTests(unittest.TestCase): self.assertIsInstance(error, NonRetryableTTSError) +_KV_REJECTION_BODY = json.dumps({"error": { + "message": "Request requires more tokens than the thinker KV cache " + "can hold (input_tokens=684, max_new_tokens=12288, " + "required_tokens=12972, kv_capacity=4095). Current " + "mem_fraction_static is 0.800; try setting " + "--thinker-mem-fraction-static higher.", + "type": "InternalServerError", "code": 500}}) + + +def _http_error(code: int, body: str) -> urllib.error.HTTPError: + return urllib.error.HTTPError( + "http://127.0.0.1:8100/v1/audio/speech", code, "error", + hdrs=None, fp=io.BytesIO(body.encode("utf-8"))) + + +def _speech_response(): + response = MagicMock() + response.__enter__.return_value = response + response.read.return_value = _WAV_BYTES + return response + + +class KvAdmissionTests(unittest.TestCase): + """The KV-window admission rejection refits max_new_tokens once.""" + + def _client(self): + client = SgOmniTTSClient.__new__(SgOmniTTSClient) + from backends.sglomni.catalog import entry_by_key + client.entry = entry_by_key("higgs_audio_v3_tts") + client.api_url = "http://127.0.0.1:8100" + client.voice = None + client.ref_audio = None + client.ref_text = "" + client.instructions = "" + client.language = "English" + client._seed = None + client.chunk_size = None + client._kv_fit = None + return client + + def test_fit_is_parsed_from_the_server_message(self): + fit = self._client()._kv_admission_fit(_KV_REJECTION_BODY) + # kv_capacity 4095 - input 684 - the 64-frame margin. + self.assertEqual(fit, 3347) + + def test_fit_is_cached_for_later_requests(self): + client = self._client() + client._kv_admission_fit(_KV_REJECTION_BODY) + client._kv_admission_fit(_KV_REJECTION_BODY) + self.assertEqual(client._kv_fit, 3347) + + def test_unrelated_errors_do_not_fit(self): + client = self._client() + self.assertIsNone(client._kv_admission_fit("CUDA out of memory")) + self.assertIsNone(client._kv_fit) + + def test_a_window_below_the_floor_raises_with_guidance(self): + body = json.dumps({"error": {"message": + "Request requires more tokens than the thinker KV cache can " + "hold (input_tokens=4000, max_new_tokens=12288, " + "required_tokens=16288, kv_capacity=4095).", "code": 500}}) + with self.assertRaises(NonRetryableTTSError) as ctx: + self._client()._kv_admission_fit(body) + self.assertIn("shorter reference clip", str(ctx.exception)) + + def test_request_wav_refits_and_resends_once(self): + client = self._client() + with patch( + "converter.clients.sglomni.urllib.request.urlopen", + side_effect=[_http_error(500, _KV_REJECTION_BODY), + _speech_response()]) as mock_open: + wav = client._request_wav("Hello.") + self.assertEqual(wav, _WAV_BYTES) + self.assertEqual(mock_open.call_count, 2) + refit = json.loads(mock_open.call_args[0][0].data) + self.assertEqual(refit["max_new_tokens"], 3347) + + def test_request_wav_surfaces_a_refit_that_fails_again(self): + client = self._client() + with patch( + "converter.clients.sglomni.urllib.request.urlopen", + side_effect=[_http_error(500, _KV_REJECTION_BODY), + _http_error(500, _KV_REJECTION_BODY)]): + with self.assertRaises(RuntimeError) as ctx: + client._request_wav("Hello.") + message = str(ctx.exception) + self.assertIn("HTTP 500", message) + self.assertIn("thinker KV cache", message) + self.assertNotIsInstance(ctx.exception, NonRetryableTTSError) + + def test_request_wav_does_not_refit_other_errors(self): + client = self._client() + with patch( + "converter.clients.sglomni.urllib.request.urlopen", + side_effect=[_http_error(500, "CUDA out of memory")]): + with self.assertRaises(RuntimeError) as ctx: + client._request_wav("Hello.") + self.assertIn("CUDA out of memory", str(ctx.exception)) + self.assertIsNone(client._kv_fit) + + def test_request_wav_keeps_the_refit_across_sub_requests(self): + # A tight window (a long reference clip): the fit binds below the + # 3000-frame catalog cap, and every later request carries it. + tight_body = json.dumps({"error": {"message": + "Request requires more tokens than the thinker KV cache can " + "hold (input_tokens=1500, max_new_tokens=3000, " + "required_tokens=4500, kv_capacity=4095).", "code": 500}}) + client = self._client() + with patch( + "converter.clients.sglomni.urllib.request.urlopen", + side_effect=[_http_error(500, tight_body), + _speech_response(), + _speech_response()]) as mock_open: + client._request_wav("Hello.") + client._request_wav("Hello again.") + self.assertEqual(mock_open.call_count, 3) + second = json.loads(mock_open.call_args[0][0].data) + self.assertEqual(second["max_new_tokens"], 2531) + + class GenerateChunkTests(unittest.TestCase): """Chunk generation: WAV output, sub-chunking, bookkeeping.""" @@ -315,6 +445,8 @@ class GenerateChunkTests(unittest.TestCase): client.instructions = "" client.language = "English" client._seed = None + client.chunk_size = None + client._kv_fit = None return client def _read_wav(self, path): @@ -347,6 +479,18 @@ class GenerateChunkTests(unittest.TestCase): self.assertEqual(len(args[0]), 3) self.assertEqual(args[1], Path(result)) + def test_run_chunk_size_caps_the_sub_requests(self): + """The pre-flight clamp (a chunk_words-capped model's popup + answer) overrides CHUNK_SIZE for this run.""" + client = self._make_client() + client.chunk_size = 10 + text = " ".join(f"word{i}" for i in range(24)) + with patch.object(client, "_request_wav", + return_value=_WAV_BYTES) as mock_wav, \ + patch("converter.clients.sglomni.concat_audio_files"): + client.generate_chunk(text, 1) + self.assertEqual(mock_wav.call_count, 3) + def test_single_subchunk_skips_concatenation(self): client = self._make_client() with patch.object(client, "_request_wav", return_value=_WAV_BYTES), \ |