[Feature] Support flexible task_type configuration for Qwen3-TTS models (#1197)

Signed-off-by: Jack Lee <695697442@qq.com>
Signed-off-by: jackleehal <695697442@qq.com>
Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
Co-authored-by: JackLeeHal <15664755+JackLeeHal@users.noreply.github.com>

vllm_omni/config/model.py

@@ -34,6 +34,8 @@ class OmniModelConfig(ModelConfig):
             "audio", "latents"). If None, output type is inferred.
         stage_connector_config: Stage connector configuration dictionary.
             Contains "name" (connector name), "extra" (extra connector config).
+        task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base).
+            If not specified, will be inferred from model path.
 
     Example:
         >>> config = OmniModelConfig(
@@ -59,6 +61,7 @@ class OmniModelConfig(ModelConfig):
     )
     omni_kv_config: dict | None = None
     codec_frame_rate_hz: float | None = None
+    task_type: str | None = None
 
     @property
     def registry(self):

vllm_omni/engine/arg_utils.py

@@ -65,6 +65,8 @@ class OmniEngineArgs(EngineArgs):
         stage_connector_spec: Extra configuration for stage connector
         async_chunk: If set to True, perform async chunk
         worker_type: Model Type, e.g., "ar" or "generation"
+        task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base).
+            If not specified, will be inferred from model path.
     """
 
     stage_id: int = 0
@@ -78,6 +80,7 @@ class OmniEngineArgs(EngineArgs):
     omni_kv_config: dict | None = None
     quantization_config: Any | None = None
     worker_type: str | None = None
+    task_type: str | None = None
 
     def __post_init__(self) -> None:
         load_omni_general_plugins()
@@ -199,6 +202,7 @@ class OmniEngineArgs(EngineArgs):
             custom_process_next_stage_input_func=self.custom_process_next_stage_input_func,
             stage_connector_config=stage_connector_config,
             omni_kv_config=self.omni_kv_config,
+            task_type=self.task_type,
         )
         omni_config.hf_config.architectures = omni_config.architectures
 
@@ -221,6 +225,8 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
             "audio", "latents"). If None, output type is inferred.
         stage_connector_spec: Extra configuration for stage connector
         worker_type: Model Type, e.g., "ar" or "generation"
+        task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base).
+            If not specified, will be inferred from model path.
     """
 
     stage_id: int = 0
@@ -234,6 +240,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
     omni_kv_config: dict | None = None
     quantization_config: Any | None = None
     worker_type: str | None = None
+    task_type: str | None = None
 
     def __post_init__(self) -> None:
         load_omni_general_plugins()
@@ -355,6 +362,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
             custom_process_next_stage_input_func=self.custom_process_next_stage_input_func,
             stage_connector_config=stage_connector_config,
             omni_kv_config=self.omni_kv_config,
+            task_type=self.task_type,
         )
         omni_config.hf_config.architectures = omni_config.architectures
 

vllm_omni/entrypoints/cli/serve.py

@@ -113,6 +113,14 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             help="Enable vLLM-Omni mode for multi-modal and diffusion models",
         )
+        omni_config_group.add_argument(
+            "--task-type",
+            type=str,
+            default=None,
+            choices=["CustomVoice", "VoiceDesign", "Base"],
+            help="Default task type for TTS models (CustomVoice, VoiceDesign, or Base). "
+            "If not specified, will be inferred from model path.",
+        )
         omni_config_group.add_argument(
             "--stage-configs-path",
             type=str,

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py

@@ -94,7 +94,9 @@ class Qwen3TTSModelForGeneration(nn.Module):
             torch_dtype=torch.bfloat16,
             **attn_kwargs,
         )
-        self.task_type = _normalize_task_type(model_path.split("-")[-1].split("/")[0])
+        self.task_type = getattr(vllm_config.model_config, "task_type", None) or _normalize_task_type(
+            model_path.split("-")[-1].split("/")[0]
+        )
         # Mark that this model produces multimodal outputs
         self.have_multimodal_outputs = True