[Feature] Support output modalities control per request (#298)

Signed-off-by: Gaohan123 <hgaoaf@connect.ust.hk>

docs/.nav.yml

@@ -7,10 +7,10 @@ nav:
   - Examples:
     - examples/README.md
     - Offline Inference:
+      - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
       - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
-      - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
     - Online Serving:
       - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md

docs/api/README.md

@@ -10,6 +10,7 @@ Main entry points for vLLM-Omni inference and serving.
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][]
 - [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][]
 - [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
+- [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
 - [vllm_omni.entrypoints.log_utils.OrchestratorMetrics][]
 - [vllm_omni.entrypoints.omni.Omni][]
 - [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]

docs/user_guide/examples/offline_inference/qwen2_5_omni.md

@@ -38,6 +38,14 @@ Then run the command below.
 bash run_single_prompt.sh
 ```
 
+### Modality control
+If you want to control output modalities, e.g. only output text, you can run the command below:
+```bash
+python end2end.py --output-wav output_audio \
+                  --query-type mixed_modalities \
+                  --modalities text
+```
+
 #### Using Local Media Files
 The `end2end.py` script supports local media files (audio, video, image) via CLI arguments:
 

docs/user_guide/examples/offline_inference/qwen3_omni.md

@@ -33,6 +33,14 @@ If you have not enough memory, you can set thinker with tensor parallel. Just ru
 bash run_single_prompt_tp.sh
 ```
 
+### Modality control
+If you want to control output modalities, e.g. only output text, you can run the command below:
+```bash
+python end2end.py --output-wav output_audio \
+                  --query-type use_audio \
+                  --modalities text
+```
+
 #### Using Local Media Files
 The `end2end.py` script supports local media files (audio, video, image) via command-line arguments:
 

docs/user_guide/examples/online_serving/qwen2_5_omni.md

@@ -27,7 +27,7 @@ Get into the example folder
 cd examples/online_serving/qwen2_5_omni
 ```
 
-####  Send request via python
+#### Send request via python
 
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
@@ -73,6 +73,14 @@ python openai_chat_completion_client_for_multimodal_generation.py \
 bash run_curl_multimodal_generation.sh mixed_modalities
 ```
 
+## Modality control
+If you want to control output modalities, e.g. only output text, you can run the command below:
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+    --query-type mixed_modalities \
+    --modalities text
+```
+
 ## Run Local Web UI Demo
 
 This Web UI demo allows users to interact with the model through a web browser.

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -81,13 +81,21 @@ sudo apt update
 sudo apt install ffmpeg
 ```
 
+## Modality control
+If you want to control output modalities, e.g. only output text, you can run the command below:
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+    --query-type use_image \
+    --modalities text
+```
+
 ## Run Local Web UI Demo
 
 This Web UI demo allows users to interact with the model through a web browser.
 
 ### Running Gradio Demo
 
-Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by:
 
 ```bash
 python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861

examples/offline_inference/qwen2_5_omni/README.md

@@ -35,6 +35,14 @@ Then run the command below.
 bash run_single_prompt.sh
 ```
 
+### Modality control
+If you want to control output modalities, e.g. only output text, you can run the command below:
+```bash
+python end2end.py --output-wav output_audio \
+                  --query-type mixed_modalities \
+                  --modalities text
+```
+
 #### Using Local Media Files
 The `end2end.py` script supports local media files (audio, video, image) via CLI arguments:
 

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -373,6 +373,11 @@ def main(args):
             prompts = [get_text_query(ln).inputs for ln in lines if ln != ""]
             print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
 
+    if args.modalities is not None:
+        output_modalities = args.modalities.split(",")
+        for i, prompt in enumerate(prompts):
+            prompt["modalities"] = output_modalities
+
     omni_outputs = omni_llm.generate(prompts, sampling_params_list)
 
     # Determine output directory: prefer --output-dir; fallback to --output-wav
@@ -381,11 +386,11 @@ def main(args):
     for stage_outputs in omni_outputs:
         if stage_outputs.final_output_type == "text":
             for output in stage_outputs.request_output:
-                request_id = int(output.request_id)
+                request_id = output.request_id
                 text_output = output.outputs[0].text
                 # Save aligned text file per request
-                prompt_text = prompts[request_id]["prompt"]
-                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
+                prompt_text = output.prompt
+                out_txt = os.path.join(output_dir, f"{request_id}.txt")
                 lines = []
                 lines.append("Prompt:\n")
                 lines.append(str(prompt_text) + "\n")
@@ -399,9 +404,9 @@ def main(args):
                 print(f"Request ID: {request_id}, Text saved to {out_txt}")
         elif stage_outputs.final_output_type == "audio":
             for output in stage_outputs.request_output:
-                request_id = int(output.request_id)
+                request_id = output.request_id
                 audio_tensor = output.multimodal_output["audio"]
-                output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
+                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
                 sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
                 print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
@@ -505,6 +510,12 @@ def parse_args():
         default=None,
         help="Address of the Ray cluster.",
     )
+    parser.add_argument(
+        "--modalities",
+        type=str,
+        default=None,
+        help="Modalities to use for the prompts.",
+    )
     return parser.parse_args()
 
 

examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh

@@ -1,3 +1,3 @@
 python end2end.py --output-wav output_audio \
                   --query-type text \
-                  --txt-prompts top10.txt
+                  --txt-prompts top100.txt

examples/offline_inference/qwen3_omni/README.md

@@ -30,6 +30,14 @@ If you have not enough memory, you can set thinker with tensor parallel. Just ru
 bash run_single_prompt_tp.sh
 ```
 
+### Modality control
+If you want to control output modalities, e.g. only output text, you can run the command below:
+```bash
+python end2end.py --output-wav output_audio \
+                  --query-type use_audio \
+                  --modalities text
+```
+
 #### Using Local Media Files
 The `end2end.py` script supports local media files (audio, video, image) via command-line arguments:
 

examples/offline_inference/qwen3_omni/end2end.py

@@ -221,6 +221,11 @@ def main(args):
             prompts = [get_text_query(ln).inputs for ln in lines if ln != ""]
             print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
 
+    if args.modalities is not None:
+        output_modalities = args.modalities.split(",")
+        for i, prompt in enumerate(prompts):
+            prompt["modalities"] = output_modalities
+
     omni_outputs = omni_llm.generate(prompts, sampling_params_list)
     # Determine output directory: prefer --output-dir; fallback to --output-wav
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
@@ -229,11 +234,11 @@ def main(args):
     for stage_outputs in omni_outputs:
         if stage_outputs.final_output_type == "text":
             for output in stage_outputs.request_output:
-                request_id = int(output.request_id)
+                request_id = output.request_id
                 text_output = output.outputs[0].text
                 # Save aligned text file per request
-                prompt_text = prompts[request_id]["prompt"]
-                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
+                prompt_text = output.prompt
+                out_txt = os.path.join(output_dir, f"{request_id}.txt")
                 lines = []
                 lines.append("Prompt:\n")
                 lines.append(str(prompt_text) + "\n")
@@ -247,9 +252,9 @@ def main(args):
                 print(f"Request ID: {request_id}, Text saved to {out_txt}")
         elif stage_outputs.final_output_type == "audio":
             for output in stage_outputs.request_output:
-                request_id = int(output.request_id)
+                request_id = output.request_id
                 audio_tensor = output.multimodal_output["audio"]
-                output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
+                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
 
                 # Convert to numpy array and ensure correct format
                 audio_numpy = audio_tensor.float().detach().cpu().numpy()
@@ -359,6 +364,12 @@ def parse_args():
         default=16000,
         help="Sampling rate for audio loading (default: 16000).",
     )
+    parser.add_argument(
+        "--modalities",
+        type=str,
+        default=None,
+        help="Output modalities to use for the prompts.",
+    )
 
     return parser.parse_args()
 

examples/offline_inference/qwen3_omni/run_multiple_prompts.sh

@@ -1,3 +1,3 @@
 python end2end.py --output-wav output_audio \
                   --query-type text \
-                  --txt-prompts top10.txt
+                  --txt-prompts top100.txt

examples/online_serving/qwen2_5_omni/README.md

@@ -24,7 +24,7 @@ Get into the example folder
 cd examples/online_serving/qwen2_5_omni
 ```
 
-####  Send request via python
+#### Send request via python
 
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
@@ -70,6 +70,14 @@ python openai_chat_completion_client_for_multimodal_generation.py \
 bash run_curl_multimodal_generation.sh mixed_modalities
 ```
 
+## Modality control
+If you want to control output modalities, e.g. only output text, you can run the command below:
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+    --query-type mixed_modalities \
+    --modalities text
+```
+
 ## Run Local Web UI Demo
 
 This Web UI demo allows users to interact with the model through a web browser.

examples/online_serving/qwen2_5_omni/gradio_demo.py

@@ -126,6 +126,15 @@ def parse_args():
         default=ASYNC_INIT_TIMEOUT,
         help="Timeout (seconds) for initializing all stages.",
     )
+    parser.add_argument(
+        "--worker-backend", type=str, default="multi_process", choices=["multi_process", "ray"], help="backend"
+    )
+    parser.add_argument(
+        "--ray-address",
+        type=str,
+        default=None,
+        help="Address of the Ray cluster.",
+    )
     return parser.parse_args()
 
 
@@ -140,6 +149,8 @@ def build_async_omni_cli_args(base_args: argparse.Namespace) -> argparse.Namespa
         shm_threshold_bytes=int(getattr(base_args, "shm_threshold_bytes", 65536)),
         batch_timeout=int(getattr(base_args, "batch_timeout", 10)),
         init_timeout=int(getattr(base_args, "init_timeout", ASYNC_INIT_TIMEOUT)),
+        worker_backend=getattr(base_args, "worker_backend", "multi_process"),
+        ray_address=getattr(base_args, "ray_address", None),
     )
 
 
@@ -294,11 +305,15 @@ async def run_inference_async_omni(
     image_file: Image.Image | None = None,
     video_file: str | None = None,
     use_audio_in_video: bool = False,
+    output_modalities: str | None = None,
 ):
     """Run inference using AsyncOmni directly with multimodal support."""
     if not user_prompt.strip() and not audio_file and not image_file and not video_file:
         return "Please provide at least a text prompt or multimodal input.", None
 
+    if output_modalities is not None:
+        output_modalities = output_modalities.split(",")
+
     try:
         # Build prompt with multimodal data
         prompt_args = SimpleNamespace(**prompt_args_template.__dict__)
@@ -381,6 +396,7 @@ async def run_inference_async_omni(
             prompt=omni_prompt,
             request_id=request_id,
             sampling_params_list=sampling_params,
+            output_modalities=output_modalities,
         ):
             # stage_outputs.request_output is a RequestOutput object, not a list
             request_output = stage_outputs.request_output
@@ -424,6 +440,7 @@ def build_interface(
         image_file: Image.Image | None,
         video_file: str | None,
         use_audio_in_video: bool,
+        output_modalities: str | None = None,
     ):
         return await run_inference_async_omni(
             omni,
@@ -434,6 +451,7 @@ def build_interface(
             image_file,
             video_file,
             use_audio_in_video,
+            output_modalities,
         )
 
     css = """
@@ -501,6 +519,14 @@ def build_interface(
                     elem_classes="audio-input",
                 )
 
+        with gr.Row():
+            output_modalities = gr.Textbox(
+                label="Output Modalities",
+                placeholder="For example: text, image, video. Use comma to separate multiple modalities.",
+                lines=1,
+                scale=1,
+            )
+
         with gr.Row():
             generate_btn = gr.Button(
                 "Generate",
@@ -515,7 +541,7 @@ def build_interface(
 
         generate_btn.click(
             fn=run_inference,
-            inputs=[input_box, audio_input, image_input, video_input, use_audio_in_video_checkbox],
+            inputs=[input_box, audio_input, image_input, video_input, use_audio_in_video_checkbox, output_modalities],
             outputs=[text_output, audio_output],
         )
         demo.queue()

examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -341,12 +341,18 @@ def run_multimodal_generation(args) -> None:
     if args.query_type == "use_audio_in_video":
         extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}
 
+    if args.modalities is not None:
+        output_modalities = args.modalities.split(",")
+    else:
+        output_modalities = None
+
     chat_completion = client.chat.completions.create(
         messages=[
             get_system_prompt(),
             prompt,
         ],
         model=model_name,
+        modalities=output_modalities,
         extra_body=extra_body,
     )
 
@@ -401,6 +407,12 @@ def parse_args():
         default=None,
         help="Custom text prompt/question to use instead of the default prompt for the selected query type.",
     )
+    parser.add_argument(
+        "--modalities",
+        type=str,
+        default=None,
+        help="Output modalities to use for the prompts.",
+    )
 
     return parser.parse_args()
 

examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh

@@ -4,14 +4,18 @@ set -euo pipefail
 # Default query type
 QUERY_TYPE="${1:-mixed_modalities}"
 
+# Default modalities argument
+MODALITIES="${2:-null}"
+
 # Validate query type
 if [[ ! "$QUERY_TYPE" =~ ^(mixed_modalities|use_audio_in_video|multi_audios|text)$ ]]; then
     echo "Error: Invalid query type '$QUERY_TYPE'"
-    echo "Usage: $0 [mixed_modalities|use_audio_in_video|multi_audios|text]"
+    echo "Usage: $0 [mixed_modalities|use_audio_in_video|multi_audios|text] [modalities]"
     echo "  mixed_modalities: Audio + Image + Video + Text query"
     echo "  use_audio_in_video: Video + Text query (with audio extraction from video)"
     echo "  multi_audios: Two audio clips + Text query"
     echo "  text: Text query"
+    echo "  modalities: Modalities parameter (default: null)"
     exit 1
 fi
 
@@ -164,6 +168,7 @@ output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
   "model": "Qwen/Qwen2.5-Omni-7B",
   "sampling_params_list": $sampling_params_list,
   "mm_processor_kwargs": $mm_processor_kwargs,
+  "modalities": $MODALITIES,
   "messages": [
     {
       "role": "system",

examples/online_serving/qwen3_omni/README.md

@@ -78,13 +78,21 @@ sudo apt update
 sudo apt install ffmpeg
 ```
 
+## Modality control
+If you want to control output modalities, e.g. only output text, you can run the command below:
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+    --query-type use_image \
+    --modalities text
+```
+
 ## Run Local Web UI Demo
 
 This Web UI demo allows users to interact with the model through a web browser.
 
 ### Running Gradio Demo
 
-Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by:
 
 ```bash
 python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861

examples/online_serving/qwen3_omni/gradio_demo.py

@@ -129,6 +129,15 @@ def parse_args():
         default=ASYNC_INIT_TIMEOUT,
         help="Timeout (seconds) for initializing all stages.",
     )
+    parser.add_argument(
+        "--worker-backend", type=str, default="multi_process", choices=["multi_process", "ray"], help="backend"
+    )
+    parser.add_argument(
+        "--ray-address",
+        type=str,
+        default=None,
+        help="Address of the Ray cluster.",
+    )
     return parser.parse_args()
 
 
@@ -143,6 +152,8 @@ def build_async_omni_cli_args(base_args: argparse.Namespace) -> argparse.Namespa
         shm_threshold_bytes=int(getattr(base_args, "shm_threshold_bytes", 65536)),
         batch_timeout=int(getattr(base_args, "batch_timeout", 10)),
         init_timeout=int(getattr(base_args, "init_timeout", ASYNC_INIT_TIMEOUT)),
+        worker_backend=getattr(base_args, "worker_backend", "multi_process"),
+        ray_address=getattr(base_args, "ray_address", None),
     )
 
 
@@ -297,11 +308,15 @@ async def run_inference_async_omni(
     image_file: Image.Image | None = None,
     video_file: str | None = None,
     use_audio_in_video: bool = False,
+    output_modalities: str | None = None,
 ):
     """Run inference using AsyncOmni directly with multimodal support."""
     if not user_prompt.strip() and not audio_file and not image_file and not video_file:
         return "Please provide at least a text prompt or multimodal input.", None
 
+    if output_modalities is not None:
+        output_modalities = output_modalities.split(",")
+
     try:
         # Build prompt with multimodal data
         prompt_args = SimpleNamespace(**prompt_args_template.__dict__)
@@ -386,6 +401,7 @@ async def run_inference_async_omni(
             prompt=omni_prompt,
             request_id=request_id,
             sampling_params_list=sampling_params,
+            output_modalities=output_modalities,
         ):
             # stage_outputs.request_output is a RequestOutput object, not a list
             request_output = stage_outputs.request_output
@@ -430,6 +446,7 @@ def build_interface(
         image_file: Image.Image | None,
         video_file: str | None,
         use_audio_in_video: bool,
+        output_modalities: str | None = None,
     ):
         return await run_inference_async_omni(
             omni,
@@ -440,6 +457,7 @@ def build_interface(
             image_file,
             video_file,
             use_audio_in_video,
+            output_modalities,
         )
 
     css = """
@@ -507,6 +525,14 @@ def build_interface(
                     elem_classes="audio-input",
                 )
 
+        with gr.Row():
+            output_modalities = gr.Textbox(
+                label="Output Modalities",
+                placeholder="For example: text, image, video. Use comma to separate multiple modalities.",
+                lines=1,
+                scale=1,
+            )
+
         with gr.Row():
             generate_btn = gr.Button(
                 "Generate",
@@ -521,7 +547,7 @@ def build_interface(
 
         generate_btn.click(
             fn=run_inference,
-            inputs=[input_box, audio_input, image_input, video_input, use_audio_in_video_checkbox],
+            inputs=[input_box, audio_input, image_input, video_input, use_audio_in_video_checkbox, output_modalities],
             outputs=[text_output, audio_output],
         )
         demo.queue()

examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -329,12 +329,18 @@ def run_multimodal_generation(args) -> None:
     if args.query_type == "use_audio_in_video":
         extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}
 
+    if args.modalities is not None:
+        output_modalities = args.modalities.split(",")
+    else:
+        output_modalities = None
+
     chat_completion = client.chat.completions.create(
         messages=[
             get_system_prompt(),
             prompt,
         ],
         model=model_name,
+        modalities=output_modalities,
         extra_body=extra_body,
     )
 
@@ -396,6 +402,12 @@ def parse_args():
         default=None,
         help="Custom text prompt/question to use instead of the default prompt for the selected query type.",
     )
+    parser.add_argument(
+        "--modalities",
+        type=str,
+        default=None,
+        help="Output modalities to use for the prompts.",
+    )
 
     return parser.parse_args()
 

examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh

@@ -4,14 +4,18 @@ set -euo pipefail
 # Default query type
 QUERY_TYPE="${1:-use_video}"
 
+# Default modalities argument
+MODALITIES="${2:-null}"
+
 # Validate query type
 if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video)$ ]]; then
     echo "Error: Invalid query type '$QUERY_TYPE'"
-    echo "Usage: $0 [text|use_audio|use_image|use_video]"
+    echo "Usage: $0 [text|use_audio|use_image|use_video] [modalities]"
     echo "  text: Text query"
     echo "  use_audio: Audio + Text query"
     echo "  use_image: Image + Text query"
     echo "  use_video: Video + Text query"
+    echo "  modalities: Modalities parameter (default: null)"
     exit 1
 fi
 
@@ -142,6 +146,7 @@ output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
   "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
   "sampling_params_list": $sampling_params_list,
   "mm_processor_kwargs": $mm_processor_kwargs,
+  "modalities": $MODALITIES,
   "messages": [
     {
       "role": "system",

tests/e2e/offline_inference/conftest.py

@@ -79,6 +79,7 @@ class OmniRunner:
         images: PromptImageInput = None,
         videos: PromptVideoInput = None,
         mm_processor_kwargs: dict[str, Any] | None = None,
+        modalities: list[str] | None = None,
     ) -> list[dict[str, Any]]:
         """
         Construct Omni input format from prompts and multimodal data.
@@ -175,6 +176,8 @@ class OmniRunner:
             input_dict: dict[str, Any] = {"prompt": full_prompt}
             if multi_modal_data:
                 input_dict["multi_modal_data"] = multi_modal_data
+            if modalities:
+                input_dict["modalities"] = modalities
             if mm_processor_kwargs:
                 input_dict["mm_processor_kwargs"] = mm_processor_kwargs
 
@@ -213,6 +216,7 @@ class OmniRunner:
         images: PromptImageInput = None,
         videos: PromptVideoInput = None,
         mm_processor_kwargs: dict[str, Any] | None = None,
+        modalities: list[str] | None = None,
     ) -> list[Any]:
         """
         Convenience method to generate with multimodal inputs.
@@ -236,6 +240,7 @@ class OmniRunner:
             images=images,
             videos=videos,
             mm_processor_kwargs=mm_processor_kwargs,
+            modalities=modalities,
         )
         return self.generate(omni_inputs, sampling_params_list)
 

tests/e2e/offline_inference/test_qwen2_5_omni.py

@@ -85,3 +85,46 @@ def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: t
         audio_tensor = audio_output.request_output[0].multimodal_output["audio"]
         assert audio_tensor is not None
         assert audio_tensor.numel() > 0
+
+
+@pytest.mark.core_model
+@pytest.mark.parametrize("test_config", test_params)
+@create_new_process_for_each_test()
+def test_mixed_modalities_to_text_only(omni_runner: type[OmniRunner], test_config: tuple[str, str]) -> None:
+    """Test processing audio, image, and video together, generating audio output."""
+    model, stage_config_path = test_config
+    with omni_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
+        # Prepare multimodal inputs
+        question = "What is recited in the audio? What is in this image? Describe the video briefly."
+        audio = AudioAsset("mary_had_lamb").audio_and_sample_rate
+        audio = (audio[0][: 16000 * 5], audio[1])  # Trim to first 5 seconds
+        image = convert_image_mode(ImageAsset("cherry_blossom").pil_image.resize((128, 128)), "RGB")
+        video = VideoAsset(name="baby_reading", num_frames=4).np_ndarrays
+        modalities = ["text"]
+
+        outputs = runner.generate_multimodal(
+            prompts=question,
+            audios=audio,
+            images=image,
+            videos=video,
+            modalities=modalities,
+        )
+
+        # Verify we got outputs from multiple stages
+        assert len(outputs) > 0
+
+        for stage_output in outputs:
+            assert stage_output.final_output_type != "audio"
+
+        # Find and verify text output (thinker stage)
+        text_output = None
+        for stage_output in outputs:
+            if stage_output.final_output_type == "text":
+                text_output = stage_output
+                break
+
+        assert text_output is not None
+        assert len(text_output.request_output) > 0
+        text_content = text_output.request_output[0].outputs[0].text
+        assert text_content is not None
+        assert len(text_content.strip()) > 0

tests/entrypoints/test_omni_llm.py

@@ -1,3 +1,4 @@
+import uuid
 import warnings
 from queue import Empty, Queue
 from typing import Any
@@ -639,17 +640,25 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
     monkeypatch.setattr(omni_llm_module, "load_stage_configs_from_model", _fake_loader)
     monkeypatch.setattr(omni_llm_module, "OmniStage", lambda cfg: _FakeStage(cfg))
 
+    # Mock uuid.uuid4() to return a predictable value for request ID generation
+    test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
+    monkeypatch.setattr(uuid, "uuid4", lambda: test_uuid)
+    monkeypatch.setattr(omni_llm_module, "uuid", uuid)
+
     from vllm_omni.entrypoints.omni_llm import OmniLLM
 
     llm = OmniLLM(model="any", init_timeout=1)
 
+    # Generate the expected request ID format: "0_<uuid>"
+    expected_request_id = f"0_{test_uuid}"
+
     # Simulate worker behavior: manually put results into output queues
     # Note: We put results before calling generate, which simulates worker processes
     # that have already completed. The polling loop will collect them in stage order.
     # Stage 0 output (will be collected first)
     llm.stage_list[0]._out_q.put_nowait(
         {
-            "request_id": 0,
+            "request_id": expected_request_id,
             "engine_outputs": [{"stage": 0, "text": "s0"}],
             "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
         }
@@ -661,7 +670,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
     # then stage 1 result will be collected.
     llm.stage_list[1]._out_q.put_nowait(
         {
-            "request_id": 0,
+            "request_id": expected_request_id,
             "engine_outputs": [{"stage": 1, "text": "s1"}],
             "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
         }
@@ -734,21 +743,29 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
     monkeypatch.setattr(omni_llm_module, "load_stage_configs_from_model", _fake_loader)
     monkeypatch.setattr(omni_llm_module, "OmniStage", lambda cfg: _FakeStage(cfg))
 
+    # Mock uuid.uuid4() to return a predictable value for request ID generation
+    test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
+    monkeypatch.setattr(uuid, "uuid4", lambda: test_uuid)
+    monkeypatch.setattr(omni_llm_module, "uuid", uuid)
+
     from vllm_omni.entrypoints.omni_llm import OmniLLM
 
     llm = OmniLLM(model="any", init_timeout=1)
 
+    # Generate the expected request ID format: "0_<uuid>"
+    expected_request_id = f"0_{test_uuid}"
+
     # Simulate worker behavior: put results into output queues
     llm.stage_list[0]._out_q.put_nowait(
         {
-            "request_id": 0,
+            "request_id": expected_request_id,
             "engine_outputs": [{"stage": 0}],
             "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
         }
     )
     llm.stage_list[1]._out_q.put_nowait(
         {
-            "request_id": 0,
+            "request_id": expected_request_id,
             "engine_outputs": [{"stage": 1}],
             "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
         }
@@ -923,14 +940,22 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
     monkeypatch.setattr(omni_llm_module, "load_stage_configs_from_model", _fake_loader)
     monkeypatch.setattr(omni_llm_module, "OmniStage", lambda cfg: _FakeStage(cfg))
 
+    # Mock uuid.uuid4() to return a predictable value for request ID generation
+    test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
+    monkeypatch.setattr(uuid, "uuid4", lambda: test_uuid)
+    monkeypatch.setattr(omni_llm_module, "uuid", uuid)
+
     from vllm_omni.entrypoints.omni_llm import OmniLLM
 
     llm = OmniLLM(model="any", init_timeout=1)
 
+    # Generate the expected request ID format: "0_<uuid>"
+    expected_request_id = f"0_{test_uuid}"
+
     # Put error message in output queue
     llm.stage_list[0]._out_q.put_nowait(
         {
-            "request_id": 0,
+            "request_id": expected_request_id,
             "error": "test error",
         }
     )
@@ -938,7 +963,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
     # (error handling continues the loop, so we need a valid result to finish)
     llm.stage_list[0]._out_q.put_nowait(
         {
-            "request_id": 0,
+            "request_id": expected_request_id,
             "engine_outputs": [{"stage": 0, "text": "result"}],
             "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
         }

vllm_omni/entrypoints/async_omni.py

@@ -56,6 +56,7 @@ from vllm_omni.entrypoints.log_utils import (
 from vllm_omni.entrypoints.omni_stage import OmniStage
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
+    get_final_stage_id_for_e2e,
     load_stage_configs_from_model,
     load_stage_configs_from_yaml,
     resolve_model_config_path,
@@ -164,6 +165,7 @@ class AsyncOmni(EngineClient):
         results.sort(key=lambda x: x[0])
         self.stage_list = [st for _, st in results]
         self.default_sampling_params_list = [st.default_sampling_params for st in self.stage_list]
+        self.output_modalities = [st.final_output_type for st in self.stage_list]
         logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
 
         if self.worker_backend == "ray":
@@ -269,6 +271,7 @@ class AsyncOmni(EngineClient):
         prompt: PromptType,
         request_id: str,
         sampling_params_list: SamplingParams | Sequence[SamplingParams] | None = None,
+        output_modalities: list[str] | None = None,
         lora_request: LoRARequest | None = None,
         trace_headers: Mapping[str, str] | None = None,
         priority: int = 0,
@@ -323,23 +326,8 @@ class AsyncOmni(EngineClient):
 
         # Determine the final stage for E2E stats (highest stage_id with
         # final_output=True; fallback to last stage)
-        final_stage_id_for_e2e = -1
-        last_stage_id = len(self.stage_list) - 1
-        try:
-            for _sid in range(last_stage_id, -1, -1):
-                if getattr(self.stage_list[_sid], "final_output", False):
-                    final_stage_id_for_e2e = _sid
-                    break
-            if final_stage_id_for_e2e < 0:
-                final_stage_id_for_e2e = last_stage_id
-        except Exception as e:
-            logger.debug(
-                "[Orchestrator] Failed to determine final stage for E2E; \
-                    falling back to last: %s",
-                e,
-                exc_info=True,
-            )
-            final_stage_id_for_e2e = last_stage_id
+        final_stage_id_for_e2e = get_final_stage_id_for_e2e(output_modalities, self.output_modalities, self.stage_list)
+
         # Metrics/aggregation helper
         metrics = OrchestratorMetrics(
             num_stages,
@@ -367,7 +355,7 @@ class AsyncOmni(EngineClient):
         logger.debug("[Orchestrator] Enqueued request %s to stage-0", request_id)
 
         logger.debug("[Orchestrator] Entering scheduling loop: stages=%d", num_stages)
-        for stage_id, stage in enumerate(self.stage_list):
+        for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]):
             result = await req_state.queue.get()
             assert stage_id == req_state.stage_id
 
@@ -439,8 +427,9 @@ class AsyncOmni(EngineClient):
                     request_output=engine_outputs,
                 )
 
+            # Forward to next stage if there is one
             next_stage_id = stage_id + 1
-            if next_stage_id < num_stages:
+            if next_stage_id <= final_stage_id_for_e2e:
                 next_stage: OmniStage = self.stage_list[next_stage_id]
                 next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
                 sp_next: SamplingParams = sampling_params_list[next_stage_id]

vllm_omni/entrypoints/log_utils.py

@@ -638,7 +638,7 @@ class OrchestratorMetrics:
             e2e_tpt = (e2e_ms / num_tokens) if num_tokens > 0 else 0.0
             log_orchestrator_e2e(self.stats_file, req_id, stage_id, e2e_ms, int(num_tokens), e2e_tpt)
 
-    def build_and_log_summary(self, final_stage_id_for_e2e: int) -> dict[str, Any]:
+    def build_and_log_summary(self, final_stage_id_to_prompt: dict[str, int]) -> dict[str, Any]:
         # Compute stage summary using wall time between first input and last output per stage
         stage_summary: list[dict[str, Any]] = []
         for sid in range(self.num_stages):
@@ -673,7 +673,7 @@ class OrchestratorMetrics:
             "e2e_avg_time_per_request_ms": e2e_avg_req,
             "e2e_avg_tokens_per_s": e2e_avg_tok,
             "wall_time_ms": wall_time_ms,
-            "final_stage_id": final_stage_id_for_e2e,
+            "final_stage_id": final_stage_id_to_prompt,
             "stages": stage_summary,
             "transfers": transfer_summary,
         }

vllm_omni/entrypoints/omni_llm.py

@@ -44,6 +44,7 @@ from vllm_omni.entrypoints.log_utils import (
 from vllm_omni.entrypoints.omni_stage import OmniStage
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
+    get_final_stage_id_for_e2e,
     load_stage_configs_from_model,
     load_stage_configs_from_yaml,
     resolve_model_config_path,
@@ -147,6 +148,7 @@ class OmniLLM:
                 results.append(fut.result())
         results.sort(key=lambda x: x[0])
         self.stage_list = [st for _, st in results]
+        self.output_modalities = [st.final_output_type for st in self.stage_list]
         logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
 
         if self.worker_backend == "ray":
@@ -291,22 +293,23 @@ class OmniLLM:
         _wall_start_ts: float = time.time()
 
         # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
-        final_stage_id_for_e2e = -1
-        last_stage_id = len(self.stage_list) - 1
+        final_stage_id_to_prompt = {}
         try:
-            for _sid in range(last_stage_id, -1, -1):
-                if getattr(self.stage_list[_sid], "final_output", False):
-                    final_stage_id_for_e2e = _sid
-                    break
-            if final_stage_id_for_e2e < 0:
-                final_stage_id_for_e2e = last_stage_id
+            for rid, prompt in request_id_to_prompt.items():
+                final_stage_id_for_e2e = get_final_stage_id_for_e2e(
+                    prompt.get("modalities", None), self.output_modalities, self.stage_list
+                )
+                final_stage_id_to_prompt[rid] = final_stage_id_for_e2e
+
         except Exception as e:
             logger.debug(
                 "[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s",
                 e,
                 exc_info=True,
             )
-            final_stage_id_for_e2e = last_stage_id
+            for rid, prompt in request_id_to_prompt.items():
+                final_stage_id_to_prompt[rid] = len(self.stage_list) - 1
+
         # Metrics/aggregation helper
         metrics = OrchestratorMetrics(
             num_stages,
@@ -406,7 +409,7 @@ class OmniLLM:
                     # (only once per request at the designated final stage)
                     try:
                         rid_key = str(req_id)
-                        if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done:
+                        if stage_id == final_stage_id_to_prompt[req_id] and rid_key not in metrics.e2e_done:
                             metrics.on_finalize_request(
                                 stage_id,
                                 req_id,
@@ -422,7 +425,7 @@ class OmniLLM:
                         )
 
                 next_stage_id = stage_id + 1
-                if next_stage_id < num_stages:
+                if next_stage_id <= final_stage_id_to_prompt[req_id]:
                     next_stage: OmniStage = self.stage_list[next_stage_id]
                     try:
                         next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]])
@@ -479,7 +482,7 @@ class OmniLLM:
 
         # Summarize and print stats
         try:
-            summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
+            summary = metrics.build_and_log_summary(final_stage_id_to_prompt)
             logger.info("[Summary] %s", summary)
         except Exception as e:
             logger.exception("[Orchestrator] Failed to build/log summary: %s", e)

vllm_omni/entrypoints/omni_stage.py

@@ -626,6 +626,7 @@ def _stage_worker(
                 idx = 0
                 for ro in unmapped:
                     target_rid = batch_request_ids[idx % len(batch_request_ids)]
+                    ro.request_id = target_rid
                     req_to_outputs[target_rid].append(ro)
                     idx += 1
 

vllm_omni/entrypoints/openai/serving_chat.py

@@ -172,6 +172,8 @@ class OmniOpenAIServingChat(OpenAIServingChat):
         if raw_request:
             raw_request.state.request_metadata = request_metadata
 
+        output_modalities = getattr(request, "modalities", self.engine_client.output_modalities)
+
         # Schedule the request and get the result generator.
         generators: list[AsyncGenerator[RequestOutput, None]] = []
         try:
@@ -194,6 +196,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
                     prompt=engine_prompt,
                     request_id=request_id,
                     sampling_params_list=sampling_params_list,
+                    output_modalities=output_modalities,
                     lora_request=lora_request,
                     trace_headers=trace_headers,
                     priority=request.priority,

vllm_omni/entrypoints/utils.py

@@ -5,6 +5,7 @@ from pathlib import Path
 from typing import Any
 
 from omegaconf import OmegaConf
+from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_config
 
 from vllm_omni.utils import detect_device_type
@@ -12,6 +13,8 @@ from vllm_omni.utils import detect_device_type
 # Get the project root directory (2 levels up from this file)
 PROJECT_ROOT = Path(__file__).parent.parent.parent
 
+logger = init_logger(__name__)
+
 
 def _convert_dataclasses_to_dict(obj: Any) -> Any:
     """Recursively convert non-serializable objects to OmegaConf-compatible types.
@@ -143,3 +146,46 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None
             base_engine_args_tmp = OmegaConf.merge(base_engine_args_tmp, stage_arg.engine_args)
         stage_arg.engine_args = base_engine_args_tmp
     return stage_args
+
+
+def get_final_stage_id_for_e2e(output_modalities: list[str], default_modalities: list[str], stage_list: list) -> int:
+    """Get the final stage id for e2e.
+
+    Args:
+        stage_list: List of stage configurations
+
+    Returns:
+        Final stage id for e2e
+    """
+    last_stage_id = len(stage_list) - 1
+    if output_modalities is not None:
+        prompt_modalities = []
+        for modality in output_modalities:
+            if modality not in default_modalities:
+                logger.warning(f"Invalid output modality: {modality}, ignoring it")
+                continue
+            prompt_modalities.append(modality)
+        output_modalities = prompt_modalities
+    else:
+        output_modalities = default_modalities
+
+    try:
+        for _sid in range(last_stage_id, -1, -1):
+            if (
+                getattr(stage_list[_sid], "final_output", False)
+                and stage_list[_sid].final_output_type in output_modalities
+            ):
+                final_stage_id_for_e2e = _sid
+                break
+        if final_stage_id_for_e2e < 0:
+            final_stage_id_for_e2e = last_stage_id
+    except Exception as e:
+        logger.debug(
+            "[Orchestrator] Failed to determine final stage for E2E; \
+                falling back to last: %s",
+            e,
+            exc_info=True,
+        )
+        final_stage_id_for_e2e = last_stage_id
+
+    return final_stage_id_for_e2e

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -61,7 +61,7 @@ stage_args:
   - stage_id: 2
     runtime:
       process: true
-      devices: "2"            # Example: use a different GPU than the previous stage; use "0" if single GPU
+      devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
       max_batch_size: 1
     engine_args:
       model_stage: code2wav