From c0b714a8aced43bbbafe54d6f44a06e6d006c195 Mon Sep 17 00:00:00 2001
From: Jack Lee <695697442@qq.com>
Date: Fri, 6 Mar 2026 02:39:09 +0800
Subject: [PATCH] [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                              | 3 +++
 vllm_omni/engine/arg_utils.py                          | 8 ++++++++
 vllm_omni/entrypoints/cli/serve.py                     | 8 ++++++++
 vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py | 4 +++-
 4 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py
index f13a90bb7..a1d9e7765 100644
--- a/vllm_omni/config/model.py
+++ b/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):
diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index e3cfecc36..cbb42829f 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/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
 
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index 416b86f8c..a580f9536 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/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,
diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py
index 028d4a53d..b70f2f3b3 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py
+++ b/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