[Migrate] Move sampling parameter overrides from legacy dispatches to TTS adapters (#5272)

Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>

tests/entrypoints/openai_api/test_audex_serving_guards.py

@@ -26,8 +26,13 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def _serving(model_type: str = "audex") -> OmniOpenAIServingSpeech:
+    from vllm_omni.entrypoints.openai.tts_adapters import resolve_adapter
+    from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext
+
     serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
     serving._tts_model_type = model_type
+    adapter_cls = resolve_adapter(model_type)
+    serving._adapter = adapter_cls(SpeechServingContext(server=serving)) if adapter_cls is not None else None
     serving._mark_ref_audio_artifact_ready_for_request = lambda request_id: None
     serving._discard_ref_audio_artifact_warmup = lambda request_id: None
     return serving
@@ -154,7 +159,7 @@ class TestAudexCfgValidation:
         assert _adapter().validate(_speech_request(cfg_scale=None)) is None
         serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
         params = SimpleNamespace(extra_args={"cfg_scale": None})
-        out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": "p"}, params)
+        out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": "p"}, params)
         assert out.extra_args == {}
 
     @pytest.mark.parametrize("key", ["cfg_role", "cfg_pair_id", "cfg_null_prompt"])
@@ -177,21 +182,21 @@ class TestAudexCfgInjection:
     def test_no_extra_args_is_untouched(self):
         serving = self._serving_with_fake_tokenizer()
         params = SimpleNamespace(extra_args=None)
-        out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
+        out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
         assert out.extra_args is None
         assert params.extra_args is None
 
     def test_scale_one_leaves_params_byte_identical(self):
         serving = self._serving_with_fake_tokenizer()
         params = SimpleNamespace(extra_args={"cfg_scale": 1.0})
-        out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
+        out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
         assert out.extra_args == {}
 
     def test_guided_request_gains_pair_contract(self):
         serving = self._serving_with_fake_tokenizer()
         params = SimpleNamespace(extra_args={"cfg_scale": 1.5})
         cond = self._cond_prompt("Some transcription text here.")
-        out = serving._inject_audex_cfg_pair_args("req-42", {"prompt": cond}, params)
+        out = serving._adapter._inject_audex_cfg_pair_args("req-42", {"prompt": cond}, params)
 
         extra = out.extra_args
         assert extra["cfg_scale"] == 1.5
@@ -207,14 +212,14 @@ class TestAudexCfgInjection:
         serving = self._serving_with_fake_tokenizer()
         params = SimpleNamespace(extra_args={"cfg_scale": 1.5})
         with pytest.raises(ValueError, match="adapter-built text prompt"):
-            serving._inject_audex_cfg_pair_args("req-1", {"prompt_token_ids": [1]}, params)
+            serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt_token_ids": [1]}, params)
 
     def test_injection_never_mutates_input_params(self):
         """The input may be the engine's SHARED default params (review P2):
         pair state must land only on the returned clone."""
         serving = self._serving_with_fake_tokenizer()
         shared = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
-        out = serving._inject_audex_cfg_pair_args("req-7", {"prompt": self._cond_prompt()}, shared)
+        out = serving._adapter._inject_audex_cfg_pair_args("req-7", {"prompt": self._cond_prompt()}, shared)
         assert out is not shared
         assert shared.extra_args == {"cfg_scale": 1.5}
         assert shared.temperature == 0.1
@@ -252,11 +257,11 @@ def test_guided_injection_sets_guided_temperature():
     serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
     cond = TestAudexCfgInjection()._cond_prompt()
     params = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
-    out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": cond}, params)
+    out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": cond}, params)
     assert out.temperature == 0.05
 
     unguided = SimpleNamespace(extra_args={"cfg_scale": 1.0}, temperature=0.1)
-    out = serving._inject_audex_cfg_pair_args("req-2", {"prompt": cond}, unguided)
+    out = serving._adapter._inject_audex_cfg_pair_args("req-2", {"prompt": cond}, unguided)
     assert out.temperature == 0.1
 
 
@@ -340,7 +345,7 @@ class TestAudexTTAInjection:
     def test_default_injection_applies_rvq_and_cfg3(self):
         serving = self._serving()
         params = SimpleNamespace(extra_args=None)
-        out = serving._inject_audex_tta_args("req-9", self._prompt(), params)
+        out = serving._adapter._inject_audex_tta_args("req-9", self._prompt(), params)
 
         extra = out.extra_args
         rvq = extra["tta_rvq"]
@@ -355,15 +360,15 @@ class TestAudexTTAInjection:
     def test_explicit_scale_one_disables_cfg_but_keeps_rvq(self):
         serving = self._serving()
         params = SimpleNamespace(extra_args={"cfg_scale": 1.0})
-        out = serving._inject_audex_tta_args("req-1", self._prompt(), params)
+        out = serving._adapter._inject_audex_tta_args("req-1", self._prompt(), params)
         assert "tta_rvq" in out.extra_args
         assert "cfg_scale" not in out.extra_args
         assert "cfg_role" not in out.extra_args
 
     def test_rvq_contract_cached_per_process(self):
         serving = self._serving()
-        out_a = serving._inject_audex_tta_args("a", self._prompt(), SimpleNamespace(extra_args=None))
-        out_b = serving._inject_audex_tta_args("b", self._prompt(), SimpleNamespace(extra_args=None))
+        out_a = serving._adapter._inject_audex_tta_args("a", self._prompt(), SimpleNamespace(extra_args=None))
+        out_b = serving._adapter._inject_audex_tta_args("b", self._prompt(), SimpleNamespace(extra_args=None))
         assert out_a.extra_args["tta_rvq"] is out_b.extra_args["tta_rvq"]
 
     def test_injection_never_mutates_input_params(self):
@@ -372,8 +377,8 @@ class TestAudexTTAInjection:
         per-request pair state must land only on the returned clone."""
         serving = self._serving()
         shared = SimpleNamespace(extra_args=None)
-        out_1 = serving._inject_audex_tta_args("req-1", self._prompt(), shared)
-        out_2 = serving._inject_audex_tta_args("req-2", self._prompt(), shared)
+        out_1 = serving._adapter._inject_audex_tta_args("req-1", self._prompt(), shared)
+        out_2 = serving._adapter._inject_audex_tta_args("req-2", self._prompt(), shared)
         assert shared.extra_args is None
         assert out_1 is not shared and out_2 is not shared
         assert out_1.extra_args["cfg_pair_id"] == "req-1"

tests/entrypoints/openai_api/test_serving_speech.py

@@ -52,6 +52,7 @@ from vllm_omni.model_executor.models.ming_tts.constants import (
     KEY_SPEAKER_WAVEFORM,
     KEY_SPEAKER_WAVEFORM_LENGTHS,
     SPEAKER_EMBEDDING_DIM,
+    TEXT_EOS_TOKEN_ID,
 )
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -3618,7 +3619,21 @@ class TestCosyVoice3Serving:
         assert error is not None
         assert "max_new_tokens" in error
 
-    def test_prepare_speech_generation_cosyvoice3(self, cosyvoice3_server, mocker: MockerFixture):
+    @pytest.mark.parametrize(
+        ("max_new_tokens", "expected_min_tokens", "expected_max_tokens"),
+        [
+            (None, 10, 2000),
+            (5, 5, 5),
+        ],
+    )
+    def test_prepare_speech_generation_cosyvoice3(
+        self,
+        cosyvoice3_server,
+        mocker: MockerFixture,
+        max_new_tokens: int | None,
+        expected_min_tokens: int,
+        expected_max_tokens: int,
+    ):
         cosyvoice3_server._build_cosyvoice3_prompt = mocker.AsyncMock(
             return_value={
                 "prompt": "Hello",
@@ -3626,15 +3641,23 @@ class TestCosyVoice3Serving:
                 "mm_processor_kwargs": {"prompt_text": "ref text", "sample_rate": 24000},
             }
         )
-        cosyvoice3_server._apply_cosyvoice3_dynamic_tokens = mocker.MagicMock(side_effect=lambda spl, req: spl)
+        cosyvoice3_server.model_config.hf_config = SimpleNamespace(
+            min_token_text_ratio=1, max_token_text_ratio=200, allowed_special=True
+        )
+        cosyvoice3_server._cosyvoice3_tokenizer = mocker.MagicMock()
+        mocker.patch("vllm_omni.model_executor.models.cosyvoice3.utils.extract_text_token", return_value=[None, 10])
 
         request = OpenAICreateSpeechRequest(
             input="Hello",
             ref_audio="data:audio/wav;base64,abc",
             ref_text="Reference text",
+            max_new_tokens=max_new_tokens,
         )
         request_id, generator, tts_params = asyncio.run(cosyvoice3_server._prepare_speech_generation(request))
+        sampling_params = cosyvoice3_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
 
+        assert sampling_params.min_tokens == expected_min_tokens
+        assert sampling_params.max_tokens == expected_max_tokens
         assert request_id.startswith("speech-")
         assert generator == "generator"
         assert tts_params == {}
@@ -3718,6 +3741,90 @@ class TestGLMTTSServing:
         assert text_token_len == 3
         load_tokenizer.assert_called_once()
 
+    def test_prepare_speech_generation_glm_tts(self, glm_tts_server, mocker: MockerFixture):
+        glm_tts_server._build_glm_tts_prompt = mocker.AsyncMock(
+            return_value={
+                "prompt": "Hello",
+                "multi_modal_data": {"audio": (np.zeros(24000), 24000)},
+                "mm_processor_kwargs": {"prompt_text": "ref text", "sample_rate": 24000},
+            }
+        )
+        request = OpenAICreateSpeechRequest(
+            input="Hello",
+            ref_audio="data:audio/wav;base64,abc",
+            ref_text="Reference text",
+        )
+        glm_tts_server.model_config.hf_config = SimpleNamespace(min_token_text_ratio=1, max_token_text_ratio=200)
+        glm_tts_server._estimate_glm_tts_text_token_len = mocker.MagicMock(return_value=10)
+        request_id, generator, tts_params = asyncio.run(glm_tts_server._prepare_speech_generation(request))
+        sampling_params = glm_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
+
+        assert sampling_params.min_tokens == 10
+        assert sampling_params.max_tokens == 2000
+        assert request_id.startswith("speech-")
+        assert generator == "generator"
+        assert tts_params == {}
+        glm_tts_server._build_glm_tts_prompt.assert_awaited_once()
+
+
+@pytest.fixture
+def ming_tts_server(mocker: MockerFixture):
+    mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
+    mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+
+    mock_engine_client = mocker.MagicMock()
+    mock_engine_client.errored = False
+    mock_engine_client.model_config = mocker.MagicMock(
+        model="inclusionAI/Ming-omni-tts-0.5B",
+        hf_config=SimpleNamespace(min_token_text_ratio=2, max_token_text_ratio=20),
+    )
+    mock_engine_client.default_sampling_params_list = [
+        SimpleNamespace(max_tokens=2048, min_tokens=None, extra_args=None)
+    ]
+    mock_engine_client.tts_batch_max_items = 32
+    mock_engine_client.generate = mocker.MagicMock(return_value="generator")
+    mock_engine_client.stage_configs = [
+        SimpleNamespace(
+            engine_args=SimpleNamespace(
+                model_stage="llm", model_arch="MingTTSForConditionalGeneration", worker_type="ar"
+            ),
+            tts_args={},
+        )
+    ]
+
+    mock_models = mocker.MagicMock()
+    mock_models.is_base_model.return_value = True
+    return OmniOpenAIServingSpeech(
+        engine_client=mock_engine_client,
+        models=mock_models,
+        request_logger=mocker.MagicMock(),
+    )
+
+
+class TestMingTTSServing:
+    def test_prepare_speech_generation_ming(self, ming_tts_server, mocker: MockerFixture):
+        ming_tts_server._adapter.build = mocker.AsyncMock(
+            return_value=PreparedRequest(
+                prompt={"prompt": "hello"},
+                tts_params={},
+                model_type="ming_tts",
+            )
+        )
+
+        request = OpenAICreateSpeechRequest(
+            input="Hello",
+            ref_audio="data:audio/wav;base64,abc",
+            ref_text="Reference text",
+            max_new_tokens=7,
+        )
+        asyncio.run(ming_tts_server._prepare_speech_generation(request))
+
+        assert ming_tts_server._tts_model_type == "ming_tts"
+        ming_tts_server._adapter.build.assert_awaited_once()
+        sampling_params = ming_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
+        assert sampling_params.stop_token_ids == [TEXT_EOS_TOKEN_ID]
+        assert sampling_params.max_tokens == 8
+
 
 @pytest.fixture
 def ming_flash_omni_tts_server(mocker: MockerFixture):
@@ -3901,6 +4008,9 @@ class TestTTSAsyncOffloading:
                     model_type=adapter_model_type,
                 )
 
+            def apply_sampling_overrides(self, sampling_params_list, request, prompt, request_id=None):
+                return sampling_params_list
+
         voxtral_server._tts_model_type = legacy_tts_model_type
         mocker.patch.object(voxtral_server, "_get_tts_adapter", return_value=FakeAdapter())
         log_info = mocker.patch("vllm_omni.entrypoints.openai.serving_speech.logger.info")

tests/entrypoints/openai_api/test_tts_detection.py

@@ -14,7 +14,6 @@ import itertools
 import pytest
 
 from vllm_omni.entrypoints.openai.tts_adapters import (
-    LEGACY_TTS_DETECTORS,
     TTS_ADAPTER_REGISTRY,
     all_tts_stage_keys,
     detect_tts_model_type,
@@ -282,27 +281,17 @@ def test_entry_stage_archs_is_ming_only():
     assert tts_entry_stage_archs() == frozenset({"MingTTSForConditionalGeneration"})
 
 
-def test_every_detected_type_is_adapter_backed_or_declared_legacy():
-    """Detection may only name a model that has an adapter or an explicit
-    :data:`LEGACY_TTS_DETECTORS` entry — never an undeclared string."""
-    legacy_names = {d.name for d in LEGACY_TTS_DETECTORS}
+def test_every_detected_type_is_adapter_backed():
+    """Detection may only name a model that has a registered adapter."""
     for stage, arch in itertools.product(_STAGES, _ARCHS):  # full cross product on purpose
         detected = detect_tts_model_type(stage, arch)
         if detected is None:
             continue
-        assert detected in TTS_ADAPTER_REGISTRY or detected in legacy_names, (
+        assert detected in TTS_ADAPTER_REGISTRY, (
             f"detection produced undeclared model type {detected!r} for stage={stage!r} arch={arch!r}"
         )
 
 
-def test_legacy_detectors_have_no_adapter():
-    """A legacy entry that gained an adapter must be deleted from the list."""
-    for detector in LEGACY_TTS_DETECTORS:
-        assert resolve_adapter(detector.name) is None, (
-            f"{detector.name!r} now has an adapter; remove it from LEGACY_TTS_DETECTORS"
-        )
-
-
 def test_no_ambiguous_detectors_at_equal_priority():
     """Two detectors that can match the same deployment must be ordered.