diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 47790286a..7c520e889 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -1025,19 +1025,31 @@ class TestAsyncOmniSupportedTasks: @pytest.mark.asyncio async def test_tts_only_no_generate_task(self): """TTS-only models (audio output, no text) should not include 'generate'.""" + from unittest.mock import MagicMock + from vllm_omni.entrypoints.async_omni import AsyncOmni omni = AsyncOmni.__new__(AsyncOmni) omni.output_modalities = [None, "audio"] + stage = MagicMock() + stage.is_comprehension = False + omni.stage_list = [stage] tasks = await omni.get_supported_tasks() assert "generate" not in tasks + assert "speech" in tasks @pytest.mark.asyncio async def test_omni_model_includes_generate(self): """Models with text output (e.g. Qwen3-Omni) should include 'generate'.""" + from unittest.mock import MagicMock + from vllm_omni.entrypoints.async_omni import AsyncOmni omni = AsyncOmni.__new__(AsyncOmni) omni.output_modalities = ["text", None, "audio"] + stage = MagicMock() + stage.is_comprehension = True + omni.stage_list = [stage] tasks = await omni.get_supported_tasks() assert "generate" in tasks + assert "speech" in tasks diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 0761ae6ef..1ba95124a 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -141,10 +141,12 @@ class AsyncOmni(OmniBase): ) async def get_supported_tasks(self) -> set[str]: - """Return supported tasks based on the configured stage output modalities.""" + """Return supported tasks based on stage output modalities and capabilities.""" tasks: set[str] = set() - if "text" in self.output_modalities: + if "text" in self.output_modalities or any(stage.is_comprehension for stage in self.stage_list): tasks.add("generate") + if "audio" in self.output_modalities: + tasks.add("speech") return tasks def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]: @@ -838,16 +840,6 @@ class AsyncOmni(OmniBase): stage.submit(abort_task) return None - async def get_supported_tasks(self) -> set[str]: - tasks: set[str] = set() - has_comprehension = any(stage.is_comprehension for stage in self.stage_list) - if has_comprehension: - tasks.add("generate") - for stage in self.stage_list: - if getattr(stage, "final_output_type", None) == "audio": - tasks.add("speech") - return tasks if tasks else {"generate"} - async def get_vllm_config(self) -> VllmConfig: for stage in self.stage_list: if stage.is_comprehension: