Fix duplicate get_supported_tasks definition in async_omni.py (#1825)

Signed-off-by: linyueqian <linyueqian@outlook.com>

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

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: