[Feature]Basic version of supporting streaming output (#367)

Signed-off-by: Gaohan123 <hgaoaf@connect.ust.hk>
This commit is contained in:
Gao Han authored and GitHub committed 2026-01-01 03:02:04 +13:00
1 parent c94ce4d85f
commit 9c7de5373a
41 files changed
+1567 -243

No files matched your search

+8 -3
View File
@@ -4,9 +4,8 @@
Main entry points for vLLM-Omni inference and serving.
- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
- [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalContentParser][]
- [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][]
@@ -14,11 +13,14 @@ Main entry points for vLLM-Omni inference and serving.
- [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
- [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
- [vllm_omni.entrypoints.log_utils.OrchestratorMetrics][]
- [vllm_omni.entrypoints.log_utils.StageRequestMetrics][]
- [vllm_omni.entrypoints.log_utils.StageStats][]
- [vllm_omni.entrypoints.omni.Omni][]
- [vllm_omni.entrypoints.omni.OmniBase][]
- [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]
- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
- [vllm_omni.entrypoints.omni_llm.OmniStageLLM][]
- [vllm_omni.entrypoints.omni_stage.OmniStage][]
- [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
## Inputs
@@ -34,6 +36,7 @@ Input data structures for multi-modal inputs.
Engine classes for offline and online inference.
- [vllm_omni.diffusion.diffusion_engine.BackgroundResources][]
- [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
- [vllm_omni.engine.AdditionalInformationEntry][]
- [vllm_omni.engine.AdditionalInformationPayload][]
@@ -59,6 +62,7 @@ Core scheduling and caching components.
Model execution components.
- [vllm_omni.model_executor.custom_process_mixin.CustomProcessMixin][]
- [vllm_omni.model_executor.models.output_templates.OmniOutput][]
- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni.Qwen2_5OmniForConditionalGeneration][]
- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_talker.Qwen2_5OmniTalkerForConditionalGeneration][]
@@ -108,6 +112,7 @@ Worker classes and model runners for distributed inference.
- [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
- [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
- [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
- [vllm_omni.worker.npu.npu_ar_model_runner.ExecuteModelState][]
- [vllm_omni.worker.npu.npu_ar_model_runner.NPUARModelRunner][]
- [vllm_omni.worker.npu.npu_ar_worker.NPUARWorker][]
- [vllm_omni.worker.npu.npu_generation_model_runner.NPUGenerationModelRunner][]
@@ -19,11 +19,11 @@ Then run:
```bash
python image_edit.py \
--image qwen_bear.png \
--image qwen-bear.png \
--prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
--output output_image_edit.png \
--num_inference_steps 50 \
--cfg_scale 4.0 \
--cfg_scale 4.0
```
### Multiple Image Editing (Qwen-Image-Edit-2509)
@@ -70,6 +70,13 @@ If media file paths are not provided, the script will use default assets. Suppor
- `use_audio_in_video`: Extract audio from video
- `text`: Text-only query
### FAQ
If you encounter error about backend of librosa, try to install ffmpeg with command below.
```
sudo apt update
sudo apt install ffmpeg
```
## Example materials
@@ -49,6 +49,12 @@ python end2end.py --query-type use_image --image-path /path/to/image.jpg
# Use local audio file
python end2end.py --query-type use_audio --audio-path /path/to/audio.wav
# Combine multiple local media files
python end2end.py --query-type mixed_modalities \
--video-path /path/to/video.mp4 \
--image-path /path/to/image.jpg \
--audio-path /path/to/audio.wav
```
If media file paths are not provided, the script will use default assets. Supported query types:
@@ -56,7 +62,16 @@ If media file paths are not provided, the script will use default assets. Suppor
- `use_image`: Image input
- `use_audio`: Audio input
- `text`: Text-only query
- `multi_audios`: Multiple audio inputs
- `mixed_modalities`: Combination of video, image, and audio inputs
### FAQ
If you encounter error about backend of librosa, try to install ffmpeg with command below.
```
sudo apt update
sudo apt install ffmpeg
```
## Example materials
@@ -19,7 +19,7 @@ if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
prompt = "a cup of coffee on the table"
outputs = omni.generate(prompt)
images = outputs[0].request_output[0]["images"]
images = list(outputs)[0].request_output[0].images
images[0].save("coffee.png")
```
@@ -37,7 +37,7 @@ if __name__ == "__main__":
]
outputs = omni.generate(prompts)
for i, output in enumerate(outputs):
image = output.request_output[0]["images"][0].save(f"{i}.jpg")
image = output.request_output[0].images[0].save(f"{i}.jpg")
```
## Local CLI Usage
@@ -37,10 +37,11 @@ The Python client supports the following command-line arguments:
- `--query-type` (or `-q`): Query type (default: `mixed_modalities`). Options: `mixed_modalities`, `use_audio_in_video`, `multi_audios`, `text`
- `--video-path` (or `-v`): Path to local video file or URL. If not provided and query-type uses video, uses default video URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs. Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
- `--image-path` (or `-i`): Path to local image file or URL. If not provided and query-type uses image, uses default image URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common image formats: JPEG, PNG, GIF. Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
- `--image-path` (or `-i`): Path to local image file or URL. If not provided and query-type uses image, uses default image URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common image formats: JPEG, PNG, GIF, WebP. Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
- `--audio-path` (or `-a`): Path to local audio file or URL. If not provided and query-type uses audio, uses default audio URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common audio formats: MP3, WAV, OGG, FLAC, M4A. Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
- `--prompt` (or `-p`): Custom text prompt/question. If not provided, uses default prompt for the selected query type. Example: `--prompt "What are the main activities shown in this video?"`
For example, to use mixed modalities with all local files:
```bash
@@ -59,7 +60,6 @@ bash run_curl_multimodal_generation.sh mixed_modalities
```
## Modality control
You can control output modalities to specify which types of output the model should generate. This is useful when you only need text output and want to skip audio generation stages for better performance.
### Supported modalities
@@ -137,6 +137,13 @@ response = client.chat.completions.create(
# Response contains two choices: one with text, one with audio
print(response.choices[0].message.content) # Text response
print(response.choices[1].message.audio) # Audio response
## Streaming Output
If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.
```bash
python openai_chat_completion_client_for_multimodal_generation.py \
--query-type mixed_modalities \
--stream
```
## Run Local Web UI Demo
@@ -38,10 +38,11 @@ The Python client supports the following command-line arguments:
- `--query-type` (or `-q`): Query type (default: `use_video`). Options: `text`, `use_audio`, `use_image`, `use_video`
- `--model` (or `-m`): Model name/path (default: `Qwen/Qwen3-Omni-30B-A3B-Instruct`)
- `--video-path` (or `-v`): Path to local video file or URL. If not provided and query-type is `use_video`, uses default video URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs. Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
- `--image-path` (or `-i`): Path to local image file or URL. If not provided and query-type is `use_image`, uses default image URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common image formats: JPEG, PNG, GIF. Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`.
- `--image-path` (or `-i`): Path to local image file or URL. If not provided and query-type is `use_image`, uses default image URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common image formats: JPEG, PNG, GIF, WebP. Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
- `--audio-path` (or `-a`): Path to local audio file or URL. If not provided and query-type is `use_audio`, uses default audio URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common audio formats: MP3, WAV, OGG, FLAC, M4A. Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
- `--prompt` (or `-p`): Custom text prompt/question. If not provided, uses default prompt for the selected query type. Example: `--prompt "What are the main activities shown in this video?"`
For example, to use a local video file with custom prompt:
```bash
@@ -67,7 +68,6 @@ sudo apt install ffmpeg
```
## Modality control
You can control output modalities to specify which types of output the model should generate. This is useful when you only need text output and want to skip audio generation stages for better performance.
### Supported modalities
@@ -145,6 +145,13 @@ response = client.chat.completions.create(
# Response contains two choices: one with text, one with audio
print(response.choices[0].message.content) # Text response
print(response.choices[1].message.audio) # Audio response
## Streaming Output
If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.
```bash
python openai_chat_completion_client_for_multimodal_generation.py \
--query-type use_image \
--stream
```
## Run Local Web UI Demo
@@ -63,6 +63,7 @@ from PIL import Image
from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.platform_utils import detect_device_type, is_npu
@@ -340,21 +341,23 @@ def main():
# Print profiling results
print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
# omni.generate() returns Generator[OmniRequestOutput, None, None], convert to list
outputs = list(outputs)
if not outputs:
raise ValueError("No output generated from omni.generate()")
logger.info("Outputs: %s", outputs)
# Extract images from OmniRequestOutput
# omni.generate() returns list[OmniRequestOutput], extract images from request_output[0]['images']
# omni.generate() returns Generator[OmniRequestOutput, None, None], extract images from request_output[0].images
first_output = outputs[0]
if not hasattr(first_output, "request_output") or not first_output.request_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output.request_output[0]
if not isinstance(req_out, dict) or "images" not in req_out:
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
images = req_out["images"]
images = req_out.images
if not images:
raise ValueError("No images found in request_output")
@@ -16,11 +16,11 @@ Then run:
```bash
python image_edit.py \
--image qwen_bear.png \
--image qwen-bear.png \
--prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
--output output_image_edit.png \
--num_inference_steps 50 \
--cfg_scale 4.0 \
--cfg_scale 4.0
```
### Multiple Image Editing (Qwen-Image-Edit-2509)
@@ -26,6 +26,7 @@ import PIL.Image
import torch
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.platform_utils import detect_device_type, is_npu
@@ -102,6 +103,7 @@ def main():
flow_shift=args.flow_shift,
)
# omni.generate() returns Generator[OmniRequestOutput, None, None]
frames = omni.generate(
args.prompt,
negative_prompt=args.negative_prompt,
@@ -115,6 +117,32 @@ def main():
num_frames=args.num_frames,
)
# Extract video frames from OmniRequestOutput
frames = list(frames)
if isinstance(frames, list) and len(frames) > 0:
first_item = frames[0]
# Check if it's an OmniRequestOutput
if hasattr(first_item, "final_output_type"):
if first_item.final_output_type != "image":
raise ValueError(
f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
)
# Pipeline mode: extract from nested request_output
if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
inner_output = first_item.request_output[0]
if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
frames = inner_output.images[0] if inner_output.images else None
if frames is None:
raise ValueError("No video frames found in output.")
# Diffusion mode: use direct images field
elif hasattr(first_item, "images") and first_item.images:
frames = first_item.images
else:
raise ValueError("No video frames found in OmniRequestOutput.")
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
@@ -376,12 +376,12 @@ def main(args):
for i, prompt in enumerate(prompts):
prompt["modalities"] = output_modalities
omni_outputs = omni_llm.generate(prompts, sampling_params_list)
omni_generator = 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
os.makedirs(output_dir, exist_ok=True)
for stage_outputs in omni_outputs:
for stage_outputs in omni_generator:
if stage_outputs.final_output_type == "text":
for output in stage_outputs.request_output:
request_id = output.request_id
@@ -319,12 +319,12 @@ def main(args):
for i, prompt in enumerate(prompts):
prompt["modalities"] = output_modalities
omni_outputs = omni_llm.generate(prompts, sampling_params_list)
omni_generator = 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
os.makedirs(output_dir, exist_ok=True)
for stage_outputs in omni_outputs:
for stage_outputs in omni_generator:
if stage_outputs.final_output_type == "text":
for output in stage_outputs.request_output:
request_id = output.request_id
@@ -16,7 +16,7 @@ if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
prompt = "a cup of coffee on the table"
outputs = omni.generate(prompt)
images = outputs[0].request_output[0]["images"]
images = list(outputs)[0].request_output[0].images
images[0].save("coffee.png")
```
@@ -34,7 +34,7 @@ if __name__ == "__main__":
]
outputs = omni.generate(prompts)
for i, output in enumerate(outputs):
image = output.request_output[0]["images"][0].save(f"{i}.jpg")
image = output.request_output[0].images[0].save(f"{i}.jpg")
```
## Local CLI Usage
@@ -5,6 +5,7 @@ import gradio as gr
import torch
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.platform_utils import detect_device_type, is_npu
ASPECT_RATIOS: dict[str, tuple[int, int]] = {
@@ -98,7 +99,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
if num_images not in {1, 2, 3, 4}:
raise gr.Error("Number of images must be 1, 2, 3, or 4.")
generator = torch.Generator(device=device).manual_seed(seed)
images = omni.generate(
outputs = omni.generate(
prompt.strip(),
height=height,
width=width,
@@ -107,7 +108,20 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
num_inference_steps=num_steps,
num_outputs_per_prompt=num_images,
)
return [img for img in images[:num_images]]
images_outputs = []
for output in outputs:
req_out = output.request_output[0]
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
images = req_out.images
if not images:
raise ValueError("No images found in request_output")
# Extend the list with individual images (not append the entire list)
images_outputs.extend(images)
if len(images_outputs) >= num_images:
break
# Return only the requested number of images
return images_outputs[:num_images]
with gr.Blocks(
title="vLLM-Omni Web Serving Demo",
@@ -9,6 +9,7 @@ import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.platform_utils import detect_device_type, is_npu
@@ -144,8 +145,9 @@ def main():
print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
# Extract images from OmniRequestOutput
# omni.generate() returns list[OmniRequestOutput], extract images from the first output
if not outputs or len(outputs) == 0:
# omni.generate() returns Generator[OmniRequestOutput, None, None], convert to list
outputs = list(outputs)
if not outputs:
raise ValueError("No output generated from omni.generate()")
logger.info(f"Outputs: {outputs}")
@@ -155,10 +157,10 @@ def main():
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output.request_output[0]
if not isinstance(req_out, dict) or "images" not in req_out:
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
images = req_out["images"]
images = req_out.images
if not images:
raise ValueError("No images found in request_output")
@@ -175,8 +177,6 @@ def main():
img.save(save_path)
print(f"Saved generated image to {save_path}")
omni.close()
if __name__ == "__main__":
main()
@@ -8,6 +8,7 @@ import numpy as np
import torch
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.platform_utils import detect_device_type, is_npu
@@ -53,6 +54,7 @@ def main():
flow_shift=args.flow_shift,
)
# omni.generate() returns Generator[OmniRequestOutput, None, None]
frames = omni.generate(
args.prompt,
negative_prompt=args.negative_prompt,
@@ -65,6 +67,32 @@ def main():
num_frames=args.num_frames,
)
# Extract video frames from OmniRequestOutput
frames = list(frames)
if isinstance(frames, list) and len(frames) > 0:
first_item = frames[0]
# Check if it's an OmniRequestOutput
if hasattr(first_item, "final_output_type"):
if first_item.final_output_type != "image":
raise ValueError(
f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
)
# Pipeline mode: extract from nested request_output
if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
inner_output = first_item.request_output[0]
if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
frames = inner_output.images[0] if inner_output.images else None
if frames is None:
raise ValueError("No video frames found in output.")
# Diffusion mode: use direct images field
elif hasattr(first_item, "images") and first_item.images:
frames = first_item.images
else:
raise ValueError("No video frames found in OmniRequestOutput.")
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
try:
+80 -1
View File
@@ -57,13 +57,92 @@ 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:
You can control output modalities to specify which types of output the model should generate. This is useful when you only need text output and want to skip audio generation stages for better performance.
### Supported modalities
| Modalities | Output |
|------------|--------|
| `["text"]` | Text only |
| `["audio"]` | Text + Audio |
| `["text", "audio"]` | Text + Audio |
| Not specified | Text + Audio (default) |
### Using curl
#### Text only
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Qwen/Qwen2.5-Omni-7B",
"messages": [{"role": "user", "content": "Describe vLLM in brief."}],
"modalities": ["text"]
}'
```
#### Text + Audio
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Qwen/Qwen2.5-Omni-7B",
"messages": [{"role": "user", "content": "Describe vLLM in brief."}],
"modalities": ["audio"]
}'
```
### Using Python client
```bash
python openai_chat_completion_client_for_multimodal_generation.py \
--query-type mixed_modalities \
--modalities text
```
### Using OpenAI Python SDK
#### Text only
```python
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY")
response = client.chat.completions.create(
model="Qwen/Qwen2.5-Omni-7B",
messages=[{"role": "user", "content": "Describe vLLM in brief."}],
modalities=["text"]
)
print(response.choices[0].message.content)
```
#### Text + Audio
```python
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY")
response = client.chat.completions.create(
model="Qwen/Qwen2.5-Omni-7B",
messages=[{"role": "user", "content": "Describe vLLM in brief."}],
modalities=["audio"]
)
# Response contains two choices: one with text, one with audio
print(response.choices[0].message.content) # Text response
print(response.choices[1].message.audio) # Audio response
## Streaming Output
If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.
```bash
python openai_chat_completion_client_for_multimodal_generation.py \
--query-type mixed_modalities \
--stream
```
## Run Local Web UI Demo
This Web UI demo allows users to interact with the model through a web browser.
@@ -23,16 +23,18 @@ SUPPORTED_MODELS: dict[str, dict[str, Any]] = {
"top_p": 1.0,
"top_k": -1,
"max_tokens": 2048,
"seed": SEED,
"detokenize": True,
"repetition_penalty": 1.1,
},
"talker": {
"temperature": 0.0,
"top_p": 1.0,
"top_k": -1,
"temperature": 0.9,
"top_p": 0.8,
"top_k": 40,
"max_tokens": 2048,
"seed": SEED,
"detokenize": True,
"repetition_penalty": 1.1,
"repetition_penalty": 1.05,
"stop_token_ids": [8294],
},
"code2wav": {
@@ -40,6 +42,7 @@ SUPPORTED_MODELS: dict[str, dict[str, Any]] = {
"top_p": 1.0,
"top_k": -1,
"max_tokens": 2048,
"seed": SEED,
"detokenize": True,
"repetition_penalty": 1.1,
},
@@ -241,10 +244,12 @@ def run_inference_api(
video_file: str | None = None,
use_audio_in_video: bool = False,
output_modalities: str | None = None,
stream: bool = False,
):
"""Run inference using OpenAI API client 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
yield "Please provide at least a text prompt or multimodal input.", None
return
try:
# Build message content list
@@ -324,7 +329,7 @@ def run_inference_api(
extra_body["mm_processor_kwargs"] = mm_processor_kwargs
# Parse output modalities
if output_modalities is not None:
if output_modalities and output_modalities.strip():
output_modalities_list = [m.strip() for m in output_modalities.split(",")]
else:
output_modalities_list = None
@@ -335,29 +340,71 @@ def run_inference_api(
model=model,
modalities=output_modalities_list,
extra_body=extra_body,
stream=stream,
)
# Extract outputs
text_outputs: list[str] = []
audio_output = None
if not stream:
# Non-streaming mode: extract outputs and yield once
text_outputs: list[str] = []
audio_output = None
for choice in chat_completion.choices:
if choice.message.content:
text_outputs.append(choice.message.content)
if choice.message.audio:
# Decode base64 audio
audio_data = base64.b64decode(choice.message.audio.data)
# Load audio from bytes
audio_np, sample_rate = sf.read(io.BytesIO(audio_data))
# Convert to mono if needed
if audio_np.ndim > 1:
audio_np = audio_np[:, 0]
audio_output = (int(sample_rate), audio_np.astype(np.float32))
for choice in chat_completion.choices:
if choice.message.content:
text_outputs.append(choice.message.content)
if choice.message.audio:
# Decode base64 audio
audio_data = base64.b64decode(choice.message.audio.data)
# Load audio from bytes
audio_np, sample_rate = sf.read(io.BytesIO(audio_data))
# Convert to mono if needed
if audio_np.ndim > 1:
audio_np = audio_np[:, 0]
audio_output = (int(sample_rate), audio_np.astype(np.float32))
text_response = "\n\n".join(text_outputs) if text_outputs else "No text output."
yield text_response, audio_output
else:
# Streaming mode: yield incremental updates
text_content = ""
audio_output = None
for chunk in chat_completion:
for choice in chunk.choices:
if hasattr(choice, "delta"):
content = getattr(choice.delta, "content", None)
else:
content = None
# Handle audio modality
if getattr(chunk, "modality", None) == "audio" and content:
try:
# Decode base64 audio
audio_data = base64.b64decode(content)
# Load audio from bytes
audio_np, sample_rate = sf.read(io.BytesIO(audio_data))
# Convert to mono if needed
if audio_np.ndim > 1:
audio_np = audio_np[:, 0]
audio_output = (int(sample_rate), audio_np.astype(np.float32))
# Yield current text and audio
yield text_content if text_content else "", audio_output
except Exception: # pylint: disable=broad-except
# If audio processing fails, just yield text
yield text_content if text_content else "", None
# Handle text modality
elif getattr(chunk, "modality", None) == "text":
if content:
text_content += content
# Yield updated text content (keep existing audio if any)
yield text_content, audio_output
# Final yield with accumulated text and last audio (if any)
yield text_content if text_content else "No text output.", audio_output
text_response = "\n\n".join(text_outputs) if text_outputs else "No text output."
return text_response, audio_output
except Exception as exc: # pylint: disable=broad-except
return f"Inference failed: {exc}", None
error_msg = f"Inference failed: {exc}"
yield error_msg, None
def build_interface(
@@ -374,8 +421,10 @@ def build_interface(
video_file: str | None,
use_audio_in_video: bool,
output_modalities: str | None = None,
stream: bool = False,
):
return run_inference_api(
# Always yield from the API function to maintain consistent generator behavior
yield from run_inference_api(
client,
model,
sampling_params_dict,
@@ -385,6 +434,7 @@ def build_interface(
video_file,
use_audio_in_video,
output_modalities,
stream,
)
css = """
@@ -455,8 +505,15 @@ def build_interface(
with gr.Row():
output_modalities = gr.Textbox(
label="Output Modalities",
value=None,
placeholder="For example: text, image, video. Use comma to separate multiple modalities.",
lines=1,
scale=2,
)
stream_checkbox = gr.Checkbox(
label="Stream output",
value=False,
info="Enable streaming to see output as it's generated.",
scale=1,
)
@@ -474,7 +531,15 @@ def build_interface(
generate_btn.click(
fn=run_inference,
inputs=[input_box, audio_input, image_input, video_input, use_audio_in_video_checkbox, output_modalities],
inputs=[
input_box,
audio_input,
image_input,
video_input,
use_audio_in_video_checkbox,
output_modalities,
stream_checkbox,
],
outputs=[text_output, audio_output],
)
demo.queue()
@@ -354,19 +354,43 @@ def run_multimodal_generation(args) -> None:
model=model_name,
modalities=output_modalities,
extra_body=extra_body,
stream=args.stream,
)
count = 0
for choice in chat_completion.choices:
if choice.message.audio:
audio_data = base64.b64decode(choice.message.audio.data)
audio_file_path = f"audio_{count}.wav"
with open(audio_file_path, "wb") as f:
f.write(audio_data)
print(f"Audio saved to {audio_file_path}")
count += 1
elif choice.message.content:
print("Chat completion output from text:", choice.message.content)
if not args.stream:
for choice in chat_completion.choices:
if choice.message.audio:
audio_data = base64.b64decode(choice.message.audio.data)
audio_file_path = f"audio_{count}.wav"
with open(audio_file_path, "wb") as f:
f.write(audio_data)
print(f"Audio saved to {audio_file_path}")
count += 1
elif choice.message.content:
print("Chat completion output from text:", choice.message.content)
else:
printed_content = False
for chunk in chat_completion:
for choice in chunk.choices:
if hasattr(choice, "delta"):
content = getattr(choice.delta, "content", None)
else:
content = None
if getattr(chunk, "modality", None) == "audio" and content:
audio_data = base64.b64decode(content)
audio_file_path = f"audio_{count}.wav"
with open(audio_file_path, "wb") as f:
f.write(audio_data)
print(f"\nAudio saved to {audio_file_path}")
count += 1
elif getattr(chunk, "modality", None) == "text":
if not printed_content:
printed_content = True
print("\ncontent:", end="", flush=True)
print(content, end="", flush=True)
def parse_args():
@@ -413,6 +437,11 @@ def parse_args():
default=None,
help="Output modalities to use for the prompts.",
)
parser.add_argument(
"--stream",
action="store_true",
help="Stream the response.",
)
return parser.parse_args()
+80 -1
View File
@@ -65,13 +65,92 @@ sudo apt install ffmpeg
```
## Modality control
If you want to control output modalities, e.g. only output text, you can run the command below:
You can control output modalities to specify which types of output the model should generate. This is useful when you only need text output and want to skip audio generation stages for better performance.
### Supported modalities
| Modalities | Output |
|------------|--------|
| `["text"]` | Text only |
| `["audio"]` | Text + Audio |
| `["text", "audio"]` | Text + Audio |
| Not specified | Text + Audio (default) |
### Using curl
#### Text only
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"messages": [{"role": "user", "content": "Describe vLLM in brief."}],
"modalities": ["text"]
}'
```
#### Text + Audio
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"messages": [{"role": "user", "content": "Describe vLLM in brief."}],
"modalities": ["audio"]
}'
```
### Using Python client
```bash
python openai_chat_completion_client_for_multimodal_generation.py \
--query-type use_image \
--modalities text
```
### Using OpenAI Python SDK
#### Text only
```python
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY")
response = client.chat.completions.create(
model="Qwen/Qwen3-Omni-30B-A3B-Instruct",
messages=[{"role": "user", "content": "Describe vLLM in brief."}],
modalities=["text"]
)
print(response.choices[0].message.content)
```
#### Text + Audio
```python
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY")
response = client.chat.completions.create(
model="Qwen/Qwen3-Omni-30B-A3B-Instruct",
messages=[{"role": "user", "content": "Describe vLLM in brief."}],
modalities=["audio"]
)
# Response contains two choices: one with text, one with audio
print(response.choices[0].message.content) # Text response
print(response.choices[1].message.audio) # Audio response
## Streaming Output
If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.
```bash
python openai_chat_completion_client_for_multimodal_generation.py \
--query-type use_image \
--stream
```
## Run Local Web UI Demo
This Web UI demo allows users to interact with the model through a web browser.
@@ -244,10 +244,11 @@ def run_inference_api(
video_file: str | None = None,
use_audio_in_video: bool = False,
output_modalities: str | None = None,
stream: bool = False,
):
"""Run inference using OpenAI API client 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
yield "Please provide at least a text prompt or multimodal input.", None
try:
# Build message content list
@@ -327,7 +328,7 @@ def run_inference_api(
extra_body["mm_processor_kwargs"] = mm_processor_kwargs
# Parse output modalities
if output_modalities is not None:
if output_modalities and output_modalities.strip():
output_modalities_list = [m.strip() for m in output_modalities.split(",")]
else:
output_modalities_list = None
@@ -338,29 +339,71 @@ def run_inference_api(
model=model,
modalities=output_modalities_list,
extra_body=extra_body,
stream=stream,
)
# Extract outputs
text_outputs: list[str] = []
audio_output = None
if not stream:
# Non-streaming mode: extract outputs and yield once
text_outputs: list[str] = []
audio_output = None
for choice in chat_completion.choices:
if choice.message.content:
text_outputs.append(choice.message.content)
if choice.message.audio:
# Decode base64 audio
audio_data = base64.b64decode(choice.message.audio.data)
# Load audio from bytes
audio_np, sample_rate = sf.read(io.BytesIO(audio_data))
# Convert to mono if needed
if audio_np.ndim > 1:
audio_np = audio_np[:, 0]
audio_output = (int(sample_rate), audio_np.astype(np.float32))
for choice in chat_completion.choices:
if choice.message.content:
text_outputs.append(choice.message.content)
if choice.message.audio:
# Decode base64 audio
audio_data = base64.b64decode(choice.message.audio.data)
# Load audio from bytes
audio_np, sample_rate = sf.read(io.BytesIO(audio_data))
# Convert to mono if needed
if audio_np.ndim > 1:
audio_np = audio_np[:, 0]
audio_output = (int(sample_rate), audio_np.astype(np.float32))
text_response = "\n\n".join(text_outputs) if text_outputs else "No text output."
yield text_response, audio_output
else:
# Streaming mode: yield incremental updates
text_content = ""
audio_output = None
for chunk in chat_completion:
for choice in chunk.choices:
if hasattr(choice, "delta"):
content = getattr(choice.delta, "content", None)
else:
content = None
# Handle audio modality
if getattr(chunk, "modality", None) == "audio" and content:
try:
# Decode base64 audio
audio_data = base64.b64decode(content)
# Load audio from bytes
audio_np, sample_rate = sf.read(io.BytesIO(audio_data))
# Convert to mono if needed
if audio_np.ndim > 1:
audio_np = audio_np[:, 0]
audio_output = (int(sample_rate), audio_np.astype(np.float32))
# Yield current text and audio
yield text_content if text_content else "", audio_output
except Exception: # pylint: disable=broad-except
# If audio processing fails, just yield text
yield text_content if text_content else "", None
# Handle text modality
elif getattr(chunk, "modality", None) == "text":
if content:
text_content += content
# Yield updated text content (keep existing audio if any)
yield text_content, audio_output
# Final yield with accumulated text and last audio (if any)
yield text_content if text_content else "No text output.", audio_output
text_response = "\n\n".join(text_outputs) if text_outputs else "No text output."
return text_response, audio_output
except Exception as exc: # pylint: disable=broad-except
return f"Inference failed: {exc}", None
error_msg = f"Inference failed: {exc}"
yield error_msg, None
def build_interface(
@@ -377,8 +420,10 @@ def build_interface(
video_file: str | None,
use_audio_in_video: bool,
output_modalities: str | None = None,
stream: bool = False,
):
return run_inference_api(
# Always yield from the API function to maintain consistent generator behavior
yield from run_inference_api(
client,
model,
sampling_params_dict,
@@ -388,6 +433,7 @@ def build_interface(
video_file,
use_audio_in_video,
output_modalities,
stream,
)
css = """
@@ -458,8 +504,15 @@ def build_interface(
with gr.Row():
output_modalities = gr.Textbox(
label="Output Modalities",
value=None,
placeholder="For example: text, image, video. Use comma to separate multiple modalities.",
lines=1,
scale=2,
)
stream_checkbox = gr.Checkbox(
label="Stream output",
value=False,
info="Enable streaming to see output as it's generated.",
scale=1,
)
@@ -477,7 +530,15 @@ def build_interface(
generate_btn.click(
fn=run_inference,
inputs=[input_box, audio_input, image_input, video_input, use_audio_in_video_checkbox, output_modalities],
inputs=[
input_box,
audio_input,
image_input,
video_input,
use_audio_in_video_checkbox,
output_modalities,
stream_checkbox,
],
outputs=[text_output, audio_output],
)
demo.queue()
@@ -395,19 +395,43 @@ def run_multimodal_generation(args) -> None:
model=model_name,
modalities=output_modalities,
extra_body=extra_body,
stream=args.stream,
)
count = 0
for choice in chat_completion.choices:
if choice.message.audio:
audio_data = base64.b64decode(choice.message.audio.data)
audio_file_path = f"audio_{count}.wav"
with open(audio_file_path, "wb") as f:
f.write(audio_data)
print(f"Audio saved to {audio_file_path}")
count += 1
elif choice.message.content:
print("Chat completion output from text:", choice.message.content)
if not args.stream:
for choice in chat_completion.choices:
if choice.message.audio:
audio_data = base64.b64decode(choice.message.audio.data)
audio_file_path = f"audio_{count}.wav"
with open(audio_file_path, "wb") as f:
f.write(audio_data)
print(f"Audio saved to {audio_file_path}")
count += 1
elif choice.message.content:
print("Chat completion output from text:", choice.message.content)
else:
printed_content = False
for chunk in chat_completion:
for choice in chunk.choices:
if hasattr(choice, "delta"):
content = getattr(choice.delta, "content", None)
else:
content = None
if getattr(chunk, "modality", None) == "audio" and content:
audio_data = base64.b64decode(content)
audio_file_path = f"audio_{count}.wav"
with open(audio_file_path, "wb") as f:
f.write(audio_data)
print(f"\nAudio saved to {audio_file_path}")
count += 1
elif getattr(chunk, "modality", None) == "text":
if not printed_content:
printed_content = True
print("\ncontent:", end="", flush=True)
print(content, end="", flush=True)
def parse_args():
@@ -461,6 +485,11 @@ def parse_args():
default=None,
help="Output modalities to use for the prompts.",
)
parser.add_argument(
"--stream",
action="store_true",
help="Stream the response.",
)
return parser.parse_args()
+2 -1
View File
@@ -113,7 +113,8 @@ plugins:
- https://docs.aiohttp.org/en/stable/objects.inv
- https://pillow.readthedocs.io/en/stable/objects.inv
- https://numpy.org/doc/stable/objects.inv
- https://pytorch.org/docs/stable/objects.inv
# Temporarily disabled due to decompression errors
# - https://pytorch.org/docs/stable/objects.inv
- https://psutil.readthedocs.io/en/stable/objects.inv
markdown_extensions:
+11 -10
View File
@@ -4,6 +4,7 @@
Pytest configuration and fixtures for vllm-omni tests.
"""
from collections.abc import Generator
from typing import Any
import pytest
@@ -189,7 +190,7 @@ class OmniRunner:
self,
prompts: list[dict[str, Any]],
sampling_params_list: list[SamplingParams] | None = None,
) -> list[Any]:
) -> Generator[Any, None, None]:
"""
Generate outputs for the given prompts.
@@ -205,7 +206,7 @@ class OmniRunner:
if sampling_params_list is None:
sampling_params_list = self.get_default_sampling_params_list()
return self.omni.generate(prompts, sampling_params_list)
yield from self.omni.generate(prompts, sampling_params_list)
def generate_multimodal(
self,
@@ -217,7 +218,7 @@ class OmniRunner:
videos: PromptVideoInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
modalities: list[str] | None = None,
) -> list[Any]:
) -> Generator[Any, None, None]:
"""
Convenience method to generate with multimodal inputs.
@@ -242,7 +243,7 @@ class OmniRunner:
mm_processor_kwargs=mm_processor_kwargs,
modalities=modalities,
)
return self.generate(omni_inputs, sampling_params_list)
yield from self.generate(omni_inputs, sampling_params_list)
def generate_audio(
self,
@@ -251,7 +252,7 @@ class OmniRunner:
system_prompt: str | None = None,
audios: PromptAudioInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
) -> list[Any]:
) -> Generator[Any, None, None]:
"""
Convenience method to generate with multimodal inputs.
Args:
@@ -269,7 +270,7 @@ class OmniRunner:
audios=audios,
mm_processor_kwargs=mm_processor_kwargs,
)
return self.generate(omni_inputs, sampling_params_list)
yield from self.generate(omni_inputs, sampling_params_list)
def generate_video(
self,
@@ -278,7 +279,7 @@ class OmniRunner:
system_prompt: str | None = None,
videos: PromptVideoInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
) -> list[Any]:
) -> Generator[Any, None, None]:
"""
Convenience method to generate with multimodal inputs.
Args:
@@ -296,7 +297,7 @@ class OmniRunner:
videos=videos,
mm_processor_kwargs=mm_processor_kwargs,
)
return self.generate(omni_inputs, sampling_params_list)
yield from self.generate(omni_inputs, sampling_params_list)
def generate_image(
self,
@@ -305,7 +306,7 @@ class OmniRunner:
system_prompt: str | None = None,
images: PromptImageInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
) -> list[Any]:
) -> Generator[Any, None, None]:
"""
Convenience method to generate with multimodal inputs.
Args:
@@ -323,7 +324,7 @@ class OmniRunner:
images=images,
mm_processor_kwargs=mm_processor_kwargs,
)
return self.generate(omni_inputs, sampling_params_list)
yield from self.generate(omni_inputs, sampling_params_list)
def __enter__(self):
"""Context manager entry."""
@@ -21,6 +21,7 @@ if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from vllm_omni import Omni
from vllm_omni.outputs import OmniRequestOutput
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
@@ -61,16 +62,16 @@ def test_cache_dit(model_name: str):
num_outputs_per_prompt=1, # Single output for speed
)
# Extract images from request_output[0]['images']
first_output = outputs[0]
first_output = list(outputs)[0]
assert first_output.final_output_type == "image"
if not hasattr(first_output, "request_output") or not first_output.request_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output.request_output[0]
if not isinstance(req_out, dict) or "images" not in req_out:
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
images = req_out["images"]
images = req_out.images
# Verify generation succeeded
assert images is not None
@@ -58,15 +58,15 @@ def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: t
videos=video,
)
# Verify we got outputs from multiple stages
assert len(outputs) > 0
# Find and verify text output (thinker stage)
text_output = None
output_count = 0
for stage_output in outputs:
if stage_output.final_output_type == "text":
text_output = stage_output
output_count += 1
break
assert output_count > 0
assert text_output is not None
assert len(text_output.request_output) > 0
@@ -76,10 +76,13 @@ def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: t
# Find and verify audio output (code2wav stage)
audio_output = None
output_count = 0
for stage_output in outputs:
if stage_output.final_output_type == "audio":
audio_output = stage_output
output_count += 1
break
assert output_count > 0
assert audio_output is not None
assert len(audio_output.request_output) > 0
@@ -113,18 +116,16 @@ def test_mixed_modalities_to_text_only(omni_runner: type[OmniRunner], test_confi
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
output_count = 0
for stage_output in outputs:
assert stage_output.final_output_type != "audio"
if stage_output.final_output_type == "text":
text_output = stage_output
output_count += 1
break
assert output_count > 0
assert text_output is not None
assert len(text_output.request_output) > 0
@@ -37,16 +37,16 @@ def test_video_to_audio(omni_runner: type[OmniRunner], test_config) -> None:
videos=video,
)
# Verify we got outputs from multiple stages
assert len(outputs) > 0
# Find and verify text output (thinker stage)
text_output = None
output_count = 0
for stage_output in outputs:
if stage_output.final_output_type == "text":
text_output = stage_output
output_count += 1
break
assert output_count > 0
assert text_output is not None
assert len(text_output.request_output) > 0
text_content = text_output.request_output[0].outputs[0].text
@@ -55,11 +55,14 @@ def test_video_to_audio(omni_runner: type[OmniRunner], test_config) -> None:
# Find and verify audio output (code2wav stage)
audio_output = None
output_count = 0
for stage_output in outputs:
if stage_output.final_output_type == "audio":
audio_output = stage_output
output_count += 1
break
assert output_count > 0
assert audio_output is not None
assert len(audio_output.request_output) > 0
@@ -87,7 +87,7 @@ def test_sequence_parallel(model_name: str, ulysses_degree: int, ring_degree: in
generator=torch.Generator(get_device_name()).manual_seed(seed),
num_outputs_per_prompt=1,
)
baseline_images = outputs[0].request_output[0]["images"]
baseline_images = list(outputs)[0].request_output[0].images
finally:
baseline.close()
@@ -113,7 +113,7 @@ def test_sequence_parallel(model_name: str, ulysses_degree: int, ring_degree: in
generator=torch.Generator(get_device_name()).manual_seed(seed),
num_outputs_per_prompt=1,
)
sp_images = outputs[0].request_output[0]["images"]
sp_images = list(outputs)[0].request_output[0].images
finally:
sp.close()
@@ -5,6 +5,7 @@ from pathlib import Path
import pytest
import torch
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.platform_utils import is_npu, is_rocm
# ruff: noqa: E402
@@ -47,16 +48,16 @@ def test_diffusion_model(model_name: str):
num_outputs_per_prompt=2,
)
# Extract images from request_output[0]['images']
first_output = outputs[0]
first_output = list(outputs)[0]
assert first_output.final_output_type == "image"
if not hasattr(first_output, "request_output") or not first_output.request_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output.request_output[0]
if not isinstance(req_out, dict) or "images" not in req_out:
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
images = req_out["images"]
images = req_out.images
assert len(images) == 2
# check image size
+4 -3
View File
@@ -21,6 +21,7 @@ if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from vllm_omni import Omni
from vllm_omni.outputs import OmniRequestOutput
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
@@ -57,16 +58,16 @@ def test_teacache(model_name: str):
num_outputs_per_prompt=1, # Single output for speed
)
# Extract images from request_output[0]['images']
first_output = outputs[0]
first_output = list(outputs)[0]
assert first_output.final_output_type == "image"
if not hasattr(first_output, "request_output") or not first_output.request_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output.request_output[0]
if not isinstance(req_out, dict) or "images" not in req_out:
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
images = req_out["images"]
images = req_out.images
# Verify generation succeeded
assert images is not None
+2 -2
View File
@@ -110,10 +110,10 @@ def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, Non
):
exc_info = cloudpickle.load(f)
if (original_exception := exc_info.get("pickled_exception")) is not None:
original_exception = exc_info.get("pickled_exception")
if original_exception is not None and isinstance(original_exception, Exception):
# Re-raise the actual exception object if it was
# successfully pickled.
assert isinstance(original_exception, Exception)
raise original_exception
if (original_tb := exc_info.get("traceback")) is not None:
+38 -2
View File
@@ -4,6 +4,7 @@
E2E Online tests for Qwen3-Omni model with video input and audio output.
"""
import base64
import concurrent.futures
import ctypes
import os
@@ -146,8 +147,6 @@ def client(omni_server):
@pytest.fixture(scope="session")
def base64_encoded_video() -> str:
"""Base64 encoded video for testing."""
import base64
video = VideoAsset(name="baby_reading", num_frames=4)
with open(video.video_path, "rb") as f:
content = f.read()
@@ -240,3 +239,40 @@ def test_video_to_audio_concurrent(
if hasattr(audio_message, "audio") and audio_message.audio:
assert audio_message.audio.data is not None
assert len(audio_message.audio.data) > 0
# Test streaming completion
chat_completion = client.chat.completions.create(
model=omni_server.model,
messages=messages,
stream=True,
)
# Collect text and audio data from stream
text_content = ""
audio_data = None
for chunk in chat_completion:
for choice in chunk.choices:
if hasattr(choice, "delta"):
content = getattr(choice.delta, "content", None)
else:
content = None
modality = getattr(chunk, "modality", None)
if modality == "audio" and content:
# Audio chunk - decode base64 content
if audio_data is None:
audio_data = base64.b64decode(content)
else:
audio_data += base64.b64decode(content)
elif modality == "text" and content:
# Text chunk - accumulate text content
text_content += content if content else ""
# Verify text output
assert text_content is not None and len(text_content) >= 2
# Verify audio output
assert audio_data is not None
assert len(audio_data) > 0
+5 -5
View File
@@ -553,7 +553,7 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
omni = Omni(model="any", init_timeout=1)
with pytest.raises(ValueError):
omni.generate(prompts=["hi"], sampling_params_list=[])
list(omni.generate(prompts=["hi"], sampling_params_list=[]))
def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
@@ -640,7 +640,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
# Use dicts instead of object() for serializable sampling params
sampling_params_list = [{"temperature": 0.7}, {"temperature": 0.8}]
prompts = ["hi"]
outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
outputs = list(omni.generate(prompts=prompts, sampling_params_list=sampling_params_list))
# Both stages have final_output=True, so should aggregate two OmniRequestOutput
assert len(outputs) == 2
@@ -723,7 +723,7 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
)
# Use dicts instead of object() for serializable sampling params
outputs = omni.generate(prompts=["p"], sampling_params_list=[{"temperature": 0.7}, {"temperature": 0.8}])
outputs = list(omni.generate(prompts=["p"], sampling_params_list=[{"temperature": 0.7}, {"temperature": 0.8}]))
assert outputs == []
@@ -796,7 +796,7 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
}
)
# Use the default sampling params
omni.generate(prompts=["p"], sampling_params_list=None)
list(omni.generate(prompts=["p"], sampling_params_list=None))
def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
@@ -924,7 +924,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
# Generate should handle error gracefully (log but continue)
# Use dict instead of object() for serializable sampling params
sampling_params_list = [{"temperature": 0.7}]
outputs = omni.generate(prompts=["hi"], sampling_params_list=sampling_params_list)
outputs = list(omni.generate(prompts=["hi"], sampling_params_list=sampling_params_list))
# Should return final output (error was logged but didn't stop processing)
assert isinstance(outputs, list)
# Since final_output=True, should have one output
@@ -22,6 +22,11 @@ _REQUEST_OUTPUT_KEYS = frozenset({"request_id", "prompt", "prompt_token_ids", "o
# Keys that identify a CompletionOutput dict (for reconstruction)
_COMPLETION_OUTPUT_KEYS = frozenset({"index", "text", "token_ids", "finish_reason"})
# Keys that identify an OmniRequestOutput dict (for reconstruction)
# OmniRequestOutput has 'final_output_type' which is unique, or can be identified by
# having 'finished' and ('images' or 'final_output_type')
_OMNI_REQUEST_OUTPUT_KEYS = frozenset({"finished", "final_output_type"})
class OmniMsgpackEncoder:
"""
@@ -173,7 +178,7 @@ class OmniMsgpackDecoder:
return self._post_process(result)
def _post_process(self, obj: Any) -> Any:
"""Recursively restore tensor/ndarray/image/RequestOutput from their dict representations."""
"""Recursively restore tensor/ndarray/image/RequestOutput/OmniRequestOutput from their dict representations."""
if isinstance(obj, dict):
# Check for type markers first
if obj.get(_TENSOR_MARKER):
@@ -186,6 +191,11 @@ class OmniMsgpackDecoder:
# Process values recursively first
processed = {k: self._post_process(v) for k, v in obj.items()}
# Check if this looks like an OmniRequestOutput (check before RequestOutput
# since OmniRequestOutput may also have some RequestOutput-like fields)
if self._is_omni_request_output(processed):
return self._decode_omni_request_output(processed)
# Check if this looks like a RequestOutput
if _REQUEST_OUTPUT_KEYS.issubset(processed.keys()):
return self._decode_request_output(processed)
@@ -204,6 +214,48 @@ class OmniMsgpackDecoder:
return obj
def _is_omni_request_output(self, obj: dict[str, Any]) -> bool:
"""Check if a dict looks like an OmniRequestOutput.
OmniRequestOutput can be identified by:
- Having 'finished' and 'final_output_type' fields (unique to OmniRequestOutput)
- OR having 'finished' and 'images' fields (diffusion mode)
"""
# Must have 'finished' field
if "finished" not in obj:
return False
# Check for unique identifier: 'final_output_type'
if "final_output_type" in obj:
return True
# Alternative: check for 'images' field (diffusion mode)
if "images" in obj:
return True
return False
def _decode_omni_request_output(self, obj: dict[str, Any]) -> Any:
"""Decode dict to OmniRequestOutput.
OmniRequestOutput is a dataclass, so we can use msgspec.convert
or construct it directly.
"""
from vllm_omni.outputs import OmniRequestOutput
try:
# Use msgspec.convert for dataclass reconstruction
return msgspec.convert(obj, OmniRequestOutput)
except Exception:
try:
# Fallback: construct directly if msgspec.convert fails
# (e.g., if some fields are missing or have wrong types)
return OmniRequestOutput(**obj)
except Exception:
# If both attempts fail, return dict as-is (defensive fallback)
# This should rarely happen if _is_omni_request_output is correct
return obj
def _decode_tensor(self, obj: dict[str, Any]) -> torch.Tensor:
"""Decode dict to torch.Tensor."""
dtype_str = obj["dtype"]
+70 -59
View File
@@ -341,78 +341,89 @@ class AsyncOmni(OmniBase):
logger.debug(f"[{self._name}] Entering scheduling loop: stages={num_stages}")
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
finished = False
while not finished:
result = await req_state.queue.get()
assert stage_id == req_state.stage_id
req_id = result.get("request_id")
if "error" in result:
logger.error(
f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}",
)
raise RuntimeError(result) # Request Finished due to error
req_id = result.get("request_id")
if "error" in result:
logger.error(
f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}",
)
raise RuntimeError(result) # Request Finished due to error
req_id = result.get("request_id")
if "error" in result:
logger.error(
f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}",
)
raise RuntimeError(result) # Request Finished due to error
engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
# Mark last output time for this stage whenever we receive outputs
metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
try:
_m = asdict(result.get("metrics"))
if _m is not None:
metrics.on_stage_metrics(stage_id, req_id, _m)
except Exception as e:
logger.exception(
f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}",
)
logger.debug(
f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing",
)
stage.set_engine_outputs(engine_outputs)
if getattr(stage, "final_output", False):
logger.debug(
f"[{self._name}] Request {req_id} finalized at stage-{stage_id}",
)
# End-to-end timing and time-per-token for final output
# (only once per request at the designated final stage)
engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
# Mark last output time for this stage whenever we receive outputs
metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
try:
rid_key = str(req_id)
if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done:
metrics.on_finalize_request(
stage_id,
req_id,
engine_outputs,
_req_start_ts.get(req_id, _wall_start_ts),
)
_m = asdict(result.get("metrics"))
if _m is not None:
metrics.on_stage_metrics(stage_id, req_id, _m)
except Exception as e:
logger.exception(
f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}",
f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}",
)
logger.debug(
f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing",
)
stage.set_engine_outputs(engine_outputs)
if isinstance(engine_outputs, list):
engine_outputs = engine_outputs[0]
# Handle diffusion outputs that already contain images
if stage.final_output_type == "image":
images = []
if isinstance(engine_outputs, OmniRequestOutput) and engine_outputs.images:
images = engine_outputs.images
elif hasattr(engine_outputs, "images") and engine_outputs.images:
images = engine_outputs.images
yield OmniRequestOutput(
stage_id=stage_id,
final_output_type=stage.final_output_type,
request_output=engine_outputs,
images=images,
)
else:
yield OmniRequestOutput(
stage_id=stage_id,
final_output_type=stage.final_output_type,
request_output=engine_outputs,
finished = engine_outputs.finished
if getattr(stage, "final_output", False):
logger.debug(
f"[{self._name}] Request {req_id} finalized at stage-{stage_id}",
)
# End-to-end timing and time-per-token for final output
# (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:
metrics.on_finalize_request(
stage_id,
req_id,
[engine_outputs],
_req_start_ts.get(req_id, _wall_start_ts),
)
except Exception as e:
logger.exception(
f"[{self._name}] Finalize request handling error for req "
f"{req_id} at stage {stage_id}: {e}",
)
# Handle diffusion outputs that already contain images
if stage.final_output_type == "image":
images = []
if isinstance(engine_outputs, OmniRequestOutput) and engine_outputs.images:
images = engine_outputs.images
elif hasattr(engine_outputs, "images") and engine_outputs.images:
images = engine_outputs.images
yield OmniRequestOutput(
stage_id=stage_id,
final_output_type=stage.final_output_type,
request_output=engine_outputs,
images=images,
)
else:
yield OmniRequestOutput(
stage_id=stage_id,
final_output_type=stage.final_output_type,
request_output=engine_outputs,
)
# Forward to next stage if there is one
next_stage_id = stage_id + 1
if next_stage_id <= final_stage_id_for_e2e:
if next_stage_id <= final_stage_id_for_e2e and finished:
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]
+9 -15
View File
@@ -6,7 +6,7 @@ import os
import time
import uuid
import weakref
from collections.abc import Callable, Sequence
from collections.abc import Callable, Generator, Sequence
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import asdict
from pprint import pformat
@@ -388,7 +388,7 @@ class Omni(OmniBase):
self._ray_pg,
)
def generate(self, *args: Any, **kwargs: dict[str, Any]) -> list[OmniRequestOutput]:
def generate(self, *args: Any, **kwargs: dict[str, Any]) -> Generator[OmniRequestOutput, None, None]:
"""Generate outputs for the given prompts.
Orchestrates the multi-stage pipeline based on YAML configuration.
@@ -439,7 +439,7 @@ class Omni(OmniBase):
sampling_params_list = per_stage_params
try:
return self._run_generation(prompts, sampling_params_list)
yield from self._run_generation(prompts, sampling_params_list)
except Exception as e:
logger.exception("[Orchestrator] Failed to run generation: %s", e)
raise e
@@ -451,7 +451,7 @@ class Omni(OmniBase):
prompts: PromptType | Sequence[PromptType] | OmniDiffusionRequest | Sequence[OmniDiffusionRequest],
sampling_params_list: Any | Sequence[Any] | None = None,
use_tqdm: bool | Callable[..., tqdm] = True,
) -> list[OmniRequestOutput]:
) -> Generator[OmniRequestOutput, None, None]:
"""Run generation through all stages in the pipeline."""
logger.debug(f"[{self._name}] generate() called")
if sampling_params_list is None:
@@ -472,8 +472,6 @@ class Omni(OmniBase):
else:
request_prompts = list(prompts)
final_outputs: list[OmniRequestOutput] = []
# Orchestrator keeps stage objects for input derivation
num_stages = len(self.stage_list)
@@ -600,13 +598,6 @@ class Omni(OmniBase):
stage.set_engine_outputs(engine_outputs)
if getattr(stage, "final_output", False):
final_outputs.append(
OmniRequestOutput(
stage_id=stage_id,
final_output_type=stage.final_output_type, # type: ignore[attr-defined]
request_output=engine_outputs,
)
)
logger.debug(
f"[{self._name}] Request {req_id} finalized at stage-{stage_id}",
)
@@ -626,6 +617,11 @@ class Omni(OmniBase):
logger.exception(
f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}",
)
yield OmniRequestOutput(
stage_id=stage_id,
final_output_type=stage.final_output_type, # type: ignore[attr-defined]
request_output=engine_outputs,
)
next_stage_id = stage_id + 1
if next_stage_id <= final_stage_id_to_prompt[req_id]:
@@ -690,8 +686,6 @@ class Omni(OmniBase):
except Exception as e:
logger.exception(f"[{self._name}] Failed to build/log summary: {e}")
return final_outputs
@property
def _name(self) -> str:
return "Orchestrator"
+6 -4
View File
@@ -1149,16 +1149,18 @@ async def _stage_worker_async(
diffusion_kwargs = prepare_sampling_params(sampling_params, "diffusion")
# AsyncOmniDiffusion.generate returns a single result, not an async generator
gen_output = await stage_engine.generate(prompt=prompt, request_id=rid, **diffusion_kwargs)
_gen_t1 = _time.time()
_gen_ms = (_gen_t1 - _gen_t0) * 1000.0
await generation_out_q.put((rid, gen_output, _gen_ms))
else:
# LLM stages: ensure using SamplingParams
llm_sampling_params = prepare_sampling_params(sampling_params, "llm")
gen_output = None
async for res in stage_engine.generate(ein, llm_sampling_params, rid):
gen_output = res
_gen_t1 = _time.time()
_gen_ms = (_gen_t1 - _gen_t0) * 1000.0
await generation_out_q.put((rid, gen_output, _gen_ms))
_gen_t1 = _time.time()
_gen_ms = (_gen_t1 - _gen_t0) * 1000.0
await generation_out_q.put((rid, gen_output, _gen_ms))
except Exception as e:
logger.exception("Failed on request %s: %s", rid, e)
out_q.put(
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionStreamResponse
from vllm_omni.entrypoints.openai.protocol.images import (
ImageData,
ImageGenerationRequest,
@@ -13,4 +14,5 @@ __all__ = [
"ImageGenerationRequest",
"ImageGenerationResponse",
"ResponseFormat",
"OmniChatCompletionStreamResponse",
]
@@ -0,0 +1,5 @@
from vllm.entrypoints.openai.protocol import ChatCompletionStreamResponse
class OmniChatCompletionStreamResponse(ChatCompletionStreamResponse):
modality: str | None = "text"
+778 -17
View File
@@ -6,7 +6,7 @@ import uuid
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Sequence
from datetime import datetime, timedelta, timezone
from io import BytesIO
from typing import TYPE_CHECKING, Any, Optional
from typing import TYPE_CHECKING, Any, Final, Optional
import jinja2
from fastapi import Request
@@ -29,13 +29,18 @@ from vllm.entrypoints.chat_utils import (
make_tool_call_id,
resolve_chat_template_content_format,
)
from vllm.entrypoints.harmony_utils import parse_chat_output
from vllm.entrypoints.harmony_utils import get_streamable_parser_for_assistant, parse_chat_output
from vllm.entrypoints.openai.protocol import (
ChatCompletionNamedToolChoiceParam,
ChatCompletionRequest,
ChatCompletionResponse,
ChatCompletionResponseChoice,
ChatCompletionResponseStreamChoice,
ChatCompletionStreamResponse,
ChatMessage,
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
ErrorResponse,
FunctionCall,
FunctionDefinition,
@@ -56,6 +61,8 @@ from vllm.entrypoints.openai.serving_engine import (
)
from vllm.entrypoints.openai.tool_parsers import ToolParser
from vllm.entrypoints.openai.tool_parsers.mistral_tool_parser import MistralToolCall
from vllm.entrypoints.openai.utils import maybe_filter_parallel_tool_calls
from vllm.entrypoints.utils import should_include_usage
from vllm.inputs.data import PromptType
from vllm.logger import init_logger
from vllm.lora.request import LoRARequest
@@ -68,10 +75,12 @@ from vllm.tokenizers.mistral import (
truncate_tool_call_ids,
validate_request_params,
)
from vllm.transformers_utils.tokenizer import AnyTokenizer
from vllm.utils.collection_utils import as_list
from vllm_omni.entrypoints.chat_utils import parse_chat_messages_futures
from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse
from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio
from vllm_omni.outputs import OmniRequestOutput
@@ -226,6 +235,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
raw_request.state.request_metadata = request_metadata
output_modalities = getattr(request, "modalities", self.engine_client.output_modalities)
request.modalities = (
output_modalities if output_modalities is not None else self.engine_client.output_modalities
)
# Schedule the request and get the result generator.
generators: list[AsyncGenerator[RequestOutput, None]] = []
@@ -266,7 +278,15 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
# Streaming response
if request.stream:
raise RuntimeError("Not support streaming output now.")
return self.chat_completion_stream_generator(
request,
result_generator,
request_id,
model_name,
conversation,
tokenizer,
request_metadata,
)
try:
return await self.chat_completion_full_generator(
@@ -518,6 +538,732 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
lora_request,
)
async def chat_completion_stream_generator(
self,
request: ChatCompletionRequest,
result_generator: AsyncIterator[RequestOutput],
request_id: str,
model_name: str,
conversation: list[ConversationMessage],
tokenizer: AnyTokenizer,
request_metadata: RequestResponseMetadata,
):
created_time = int(time.time())
chunk_object_type: Final = "chat.completion.chunk"
first_iteration_dict = {}
assert hasattr(request, "modalities") and request.modalities is not None, (
"Streaming request must specify output modalities"
)
for modality in request.modalities:
first_iteration_dict[modality] = True
# Send response for each token for each request.n (index)
num_choices = 1 if request.n is None else request.n
previous_num_tokens = [0] * num_choices
finish_reason_sent = [False] * num_choices
num_prompt_tokens = 0
num_cached_tokens = None
if self.use_harmony:
harmony_parsers = [get_streamable_parser_for_assistant() for _ in range(num_choices)]
harmony_tools_streamed = [False] * num_choices
tools_streamed = [False] * num_choices
if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam):
tool_choice_function_name = request.tool_choice.function.name
else:
tool_choice_function_name = None
# Determine whether tools are in use with "auto" tool choice
tool_choice_auto = not tool_choice_function_name and self._should_stream_with_auto_tool_parsing(request)
all_previous_token_ids: list[list[int]] | None
function_name_returned = [False] * num_choices
if self.tool_call_id_type == "kimi_k2":
history_tool_call_cnt = get_history_tool_calls_cnt(conversation)
else:
history_tool_call_cnt = 0
# Always track previous_texts for comprehensive output logging
previous_texts = [""] * num_choices
# Only one of these will be used, thus previous_texts and
# all_previous_token_ids will not be used twice in the same iteration.
if tool_choice_auto or self.reasoning_parser:
# These are only required in "auto" tool choice case
all_previous_token_ids = [[]] * num_choices
# For reasoning parser and tool call all enabled
added_content_delta_arr = [False] * num_choices
reasoning_end_arr = [False] * num_choices
else:
all_previous_token_ids = None
try:
if self.reasoning_parser:
reasoning_parser = self.reasoning_parser(
tokenizer,
chat_template_kwargs=request.chat_template_kwargs, # type: ignore
)
except RuntimeError as e:
logger.exception("Error in reasoning parser creation.")
data = self.create_streaming_error_response(str(e))
yield f"data: {data}\n\n"
yield "data: [DONE]\n\n"
return
# Prepare the tool parser if it's needed
try:
if tool_choice_auto and self.tool_parser:
tool_parsers: list[ToolParser | None] = [self.tool_parser(tokenizer)] * num_choices
else:
tool_parsers = [None] * num_choices
except Exception as e:
logger.exception("Error in tool parser creation.")
data = self.create_streaming_error_response(str(e))
yield f"data: {data}\n\n"
yield "data: [DONE]\n\n"
return
stream_options = request.stream_options
include_usage, include_continuous_usage = should_include_usage(stream_options, self.enable_force_include_usage)
try:
async for omni_res in result_generator:
final_output_type = omni_res.final_output_type
res = omni_res.request_output
if final_output_type not in first_iteration_dict:
logger.warning(f"final output type: {final_output_type} is not needed by the request")
continue
if res.prompt_token_ids is not None:
num_prompt_tokens = len(res.prompt_token_ids)
if res.encoder_prompt_token_ids is not None:
num_prompt_tokens += len(res.encoder_prompt_token_ids)
# We need to do it here, because if there are exceptions in
# the result_generator, it needs to be sent as the FIRST
# response (by the try...catch).
if first_iteration_dict[final_output_type] and final_output_type == "text":
num_cached_tokens = res.num_cached_tokens
# Send first response for each request.n (index) with
# the role
role = self.get_chat_request_role(request)
# NOTE num_choices defaults to 1 so this usually executes
# once per request
for i in range(num_choices):
choice_data = ChatCompletionResponseStreamChoice(
index=i,
delta=DeltaMessage(
role=role,
content="",
),
logprobs=None,
finish_reason=None,
)
# return prompt_token_ids at the first chunk ever
chunk = OmniChatCompletionStreamResponse(
id=request_id,
object=chunk_object_type,
created=created_time,
choices=[choice_data],
model=model_name,
prompt_token_ids=(res.prompt_token_ids if request.return_token_ids else None),
modality=final_output_type,
)
# if continuous usage stats are requested, add it
if include_continuous_usage:
chunk.usage = UsageInfo(
prompt_tokens=num_prompt_tokens,
completion_tokens=0,
total_tokens=num_prompt_tokens,
)
data = chunk.model_dump_json(exclude_unset=True)
yield f"data: {data}\n\n"
# Send response to echo the input portion of the
# last message
if request.echo:
last_msg_content: str | list[dict[str, str]] = ""
if conversation and "content" in conversation[-1] and conversation[-1].get("role") == role:
last_msg_content = conversation[-1]["content"] or ""
if last_msg_content:
for i in range(num_choices):
choice_data = ChatCompletionResponseStreamChoice(
index=i,
delta=DeltaMessage(content=last_msg_content),
logprobs=None,
finish_reason=None,
)
chunk = OmniChatCompletionStreamResponse(
id=request_id,
object=chunk_object_type,
created=created_time,
choices=[choice_data],
model=model_name,
modality=final_output_type,
)
if include_continuous_usage:
chunk.usage = UsageInfo(
prompt_tokens=num_prompt_tokens,
completion_tokens=0,
total_tokens=num_prompt_tokens,
)
data = chunk.model_dump_json(exclude_unset=True)
yield f"data: {data}\n\n"
first_iteration_dict[final_output_type] = False
if final_output_type == "text":
for output in res.outputs:
i = output.index
tool_parser = tool_parsers[i]
if finish_reason_sent[i]:
continue
if request.logprobs and request.top_logprobs is not None:
assert output.logprobs is not None, "Did not output logprobs"
logprobs = self._create_chat_logprobs(
token_ids=output.token_ids,
top_logprobs=output.logprobs,
tokenizer=tokenizer,
num_output_top_logprobs=request.top_logprobs,
return_as_token_id=request.return_tokens_as_token_ids,
)
else:
logprobs = None
if self.use_harmony:
harmony_parser = harmony_parsers[i]
prev_recipient = harmony_parser.current_recipient
delta_text = ""
for token_id in output.token_ids:
harmony_parser.process(token_id)
delta_text += harmony_parser.last_content_delta or ""
cur_channel = harmony_parser.current_channel
cur_recipient = harmony_parser.current_recipient
else:
# output.text is cumulative, extract only the delta portion
previous_text = previous_texts[i] if previous_texts else ""
if output.text is not None:
delta_text = output.text[len(previous_text) :]
else:
delta_text = ""
if not delta_text and not output.token_ids and not previous_num_tokens[i]:
# Chunked prefill case, don't return empty chunks
continue
delta_message: DeltaMessage | None
# just update previous_texts and previous_token_ids
if tool_choice_auto or self.reasoning_parser:
assert previous_texts is not None
assert all_previous_token_ids is not None
previous_text = previous_texts[i]
previous_token_ids = all_previous_token_ids[i]
current_text = previous_text + delta_text
# avoid the None + list error.
if previous_token_ids:
current_token_ids = previous_token_ids + as_list(output.token_ids)
else:
current_token_ids = as_list(output.token_ids)
if self.use_harmony:
if cur_channel == "final":
delta_message = DeltaMessage(content=delta_text)
elif cur_channel == "analysis":
if request.include_reasoning:
delta_message = DeltaMessage(reasoning=delta_text)
else:
delta_message = None
elif (
cur_channel == "commentary" and cur_recipient and cur_recipient.startswith("functions.")
):
# Count completed tool calls to determine index
base_index = 0
for msg in harmony_parser.messages:
if (
msg.channel == "commentary"
and msg.recipient
and msg.recipient.startswith("functions.")
):
base_index += 1
if prev_recipient != cur_recipient:
tool_name = cur_recipient.split("functions.", 1)[1]
delta_message = DeltaMessage(
tool_calls=[
DeltaToolCall(
id=make_tool_call_id(),
type="function",
function=DeltaFunctionCall(
name=tool_name,
arguments="",
),
index=base_index,
)
]
)
elif delta_text:
delta_message = DeltaMessage(
tool_calls=[
DeltaToolCall(
index=base_index,
function=DeltaFunctionCall(arguments=delta_text),
)
]
)
else:
delta_message = None
if delta_message is not None:
harmony_tools_streamed[i] = True
else:
delta_message = None
# handle streaming deltas for tools with named tool_choice
elif tool_choice_function_name:
if (
self.reasoning_parser
and not reasoning_end_arr[i]
and not reasoning_parser.is_reasoning_end(previous_token_ids)
):
assert reasoning_parser is not None
delta_message = reasoning_parser.extract_reasoning_streaming(
previous_text,
current_text,
delta_text,
previous_token_ids,
current_token_ids,
output.token_ids,
)
# When encountering think end id in delta_token_ids
# or think end id in prompt_token_ids
# i.e {"enable_thinking": False},
# set reasoning status to end.
# Only keep 'content', remove 'reasoning'.
if reasoning_parser.is_reasoning_end(as_list(output.token_ids)) or (
res.prompt_token_ids and reasoning_parser.is_reasoning_end(res.prompt_token_ids)
):
reasoning_end_arr[i] = True
if delta_message and delta_message.content:
# This need to be added to next `delta_text`
current_text = delta_message.content
delta_message.content = None
else:
current_text = ""
else:
# Just to add remaining `content`
if self.reasoning_parser:
delta_text = previous_text + delta_text
current_text = ""
if function_name_returned[i]:
delta_tool_call = DeltaToolCall(
function=DeltaFunctionCall(arguments=delta_text),
index=i,
)
else:
delta_tool_call = DeltaToolCall(
id=make_tool_call_id(),
type="function",
function=DeltaFunctionCall(
name=tool_choice_function_name,
arguments=delta_text,
),
index=i,
)
function_name_returned[i] = True
delta_message = DeltaMessage(
tool_calls=[
delta_tool_call,
]
)
tools_streamed[i] = True
elif request.tool_choice == "required":
assert previous_texts is not None
previous_text = previous_texts[i]
current_text = previous_text + delta_text
fn_name_returned = function_name_returned[i]
output_token_ids = as_list(output.token_ids)
if (
self.reasoning_parser is not None
and not reasoning_end_arr[i]
and res.prompt_token_ids
and reasoning_parser.is_reasoning_end(res.prompt_token_ids)
):
reasoning_end_arr[i] = True
if self.reasoning_parser and not reasoning_end_arr[i]:
delta_message = reasoning_parser.extract_reasoning_streaming(
previous_text,
current_text,
delta_text,
previous_token_ids,
current_token_ids,
output_token_ids,
)
if reasoning_parser.is_reasoning_end(output_token_ids):
reasoning_end_arr[i] = True
if delta_message and delta_message.content:
current_text = delta_message.content
delta_message.content = None
else:
# reasoning ended
current_text = ""
else:
# either finished reasoning or no reasoning at all
content = current_text
delta_message, function_name_returned[i] = self.extract_tool_call_required_streaming(
previous_text=previous_text,
current_text=content,
delta_text=delta_text,
function_name_returned=fn_name_returned,
tool_call_idx=history_tool_call_cnt,
)
if (
delta_message
and delta_message.tool_calls
and delta_message.tool_calls[0].id is not None
):
history_tool_call_cnt += 1
tools_streamed[i] = True
# handle streaming deltas for tools with "auto" tool choice
# and reasoning parser
elif tool_choice_auto and self.reasoning_parser:
assert tool_parser is not None
assert reasoning_parser is not None
assert added_content_delta_arr is not None
assert reasoning_end_arr is not None
output_token_ids = as_list(output.token_ids)
if not reasoning_end_arr[i]:
delta_message = reasoning_parser.extract_reasoning_streaming(
previous_text,
current_text,
delta_text,
previous_token_ids,
current_token_ids,
output_token_ids,
)
# When encountering think end id in prompt_token_ids
# i.e {"enable_thinking": False},
# set reasoning status to end.
# Remove the text and token ids related
# to 'reasoning'.
if res.prompt_token_ids and reasoning_parser.is_reasoning_end(res.prompt_token_ids):
reasoning_end_arr[i] = True
current_token_ids = output_token_ids
if delta_message and delta_message.content:
current_text = delta_message.content
delta_message.content = None
else:
current_text = ""
# When encountering think end id in delta_token_ids,
# set reasoning status to end.
# Remove the text and token ids related
# to 'reasoning'.
if reasoning_parser.is_reasoning_end(output_token_ids):
reasoning_end_arr[i] = True
current_token_ids = reasoning_parser.extract_content_ids(output_token_ids)
if delta_message and delta_message.content:
current_text = delta_message.content
delta_message.content = None
else:
current_text = ""
# handle tool calls only after reasoning is done,
else:
delta_token_ids = output_token_ids
# First time to tool call,
# add the remaining text and token ids
# to delta from previous
if not added_content_delta_arr[i]:
added_content_delta_arr[i] = True
previous_text = ""
previous_token_ids = []
delta_text = current_text
delta_token_ids = current_token_ids
delta_message = tool_parser.extract_tool_calls_streaming(
previous_text=previous_text,
current_text=current_text,
delta_text=delta_text,
previous_token_ids=previous_token_ids,
current_token_ids=current_token_ids,
delta_token_ids=delta_token_ids,
request=request,
)
if delta_message and delta_message.tool_calls:
tools_streamed[i] = True
# when only tool calls
elif tool_choice_auto:
assert tool_parser is not None
delta_message = tool_parser.extract_tool_calls_streaming(
previous_text=previous_text,
current_text=current_text,
delta_text=delta_text,
previous_token_ids=previous_token_ids,
current_token_ids=current_token_ids,
delta_token_ids=output.token_ids,
request=request,
)
if delta_message and delta_message.tool_calls:
tools_streamed[i] = True
# when only reasoning
elif self.reasoning_parser:
delta_message = reasoning_parser.extract_reasoning_streaming(
previous_text,
current_text,
delta_text,
previous_token_ids,
current_token_ids,
output.token_ids,
)
# handle streaming just a content delta
else:
delta_message = DeltaMessage(content=delta_text)
# update the previous values for the next iteration
if (tool_choice_auto or self.reasoning_parser) and not self.use_harmony:
assert previous_texts is not None
assert all_previous_token_ids is not None
previous_texts[i] = current_text
all_previous_token_ids[i] = current_token_ids
else:
# Update for comprehensive logging even in simple case
assert previous_texts is not None
previous_texts[i] += delta_text
# set the previous values for the next iteration
previous_num_tokens[i] += len(output.token_ids)
# if the message delta is None (e.g. because it was a
# "control token" for tool calls or the parser otherwise
# wasn't ready to send a token, then
# get the next token without streaming a chunk
if delta_message is None:
if output.finish_reason is None:
continue
else:
delta_message = DeltaMessage()
# Log streaming delta if output logging is enabled
if self.enable_log_outputs and self.request_logger:
delta_content = ""
if delta_message.content:
delta_content = delta_message.content
elif delta_message.tool_calls:
delta_content = "".join(
tc.function.arguments
for tc in delta_message.tool_calls
if tc.function and tc.function.arguments
)
if delta_content:
self.request_logger.log_outputs(
request_id=request_id,
outputs=delta_content,
output_token_ids=as_list(output.token_ids),
finish_reason=output.finish_reason,
is_streaming=True,
delta=True,
)
if output.finish_reason is None:
# Send token-by-token response for each request.n
choice_data = ChatCompletionResponseStreamChoice(
index=i,
delta=delta_message,
logprobs=logprobs,
finish_reason=None,
token_ids=(as_list(output.token_ids) if request.return_token_ids else None),
)
# if the model is finished generating
else:
# check to make sure we haven't "forgotten" to stream
# any tokens that were generated but previously
# matched by partial json parsing
# only happens if we are NOT using structured outputs
auto_tools_called = False
if tool_parser:
auto_tools_called = len(tool_parser.prev_tool_call_arr) > 0
index = len(tool_parser.prev_tool_call_arr) - 1 if auto_tools_called else 0
else:
index = 0
if self._should_check_for_unstreamed_tool_arg_tokens(delta_message, output) and tool_parser:
latest_delta_len = 0
if (
isinstance(
delta_message.tool_calls[0].function,
DeltaFunctionCall,
)
) and isinstance(delta_message.tool_calls[0].function.arguments, str):
latest_delta_len = len(delta_message.tool_calls[0].function.arguments)
# get the expected call based on partial JSON
# parsing which "autocompletes" the JSON
expected_call = json.dumps(
tool_parser.prev_tool_call_arr[index].get("arguments", {}),
ensure_ascii=False,
)
# get what we've streamed so far for arguments
# for the current tool
actual_call = tool_parser.streamed_args_for_tool[index]
if latest_delta_len > 0:
actual_call = actual_call[:-latest_delta_len]
# check to see if there's anything left to stream
remaining_call = expected_call.replace(actual_call, "", 1)
# set that as a delta message
delta_message = DeltaMessage(
tool_calls=[
DeltaToolCall(
index=index,
function=DeltaFunctionCall(arguments=remaining_call).model_dump(
exclude_none=True
),
)
]
)
# Send the finish response for each request.n only once
# In OpenAI's API, when a tool is called, the
# finish_reason is:
# "tool_calls" for "auto" or "required" tool calls,
# and "stop" for named tool calls.
if (
auto_tools_called
or (tools_streamed[i] and not tool_choice_function_name)
or (self.use_harmony and harmony_tools_streamed[i])
):
finish_reason_ = "tool_calls"
else:
finish_reason_ = output.finish_reason if output.finish_reason else "stop"
choice_data = ChatCompletionResponseStreamChoice(
index=i,
delta=delta_message,
logprobs=logprobs,
finish_reason=finish_reason_,
stop_reason=output.stop_reason,
token_ids=(as_list(output.token_ids) if request.return_token_ids else None),
)
finish_reason_sent[i] = True
choice_data = maybe_filter_parallel_tool_calls(choice_data, request)
chunk = OmniChatCompletionStreamResponse(
id=request_id,
object=chunk_object_type,
created=created_time,
choices=[choice_data],
model=model_name,
modality=final_output_type,
)
# handle usage stats if requested & if continuous
if include_continuous_usage:
completion_tokens = previous_num_tokens[i]
chunk.usage = UsageInfo(
prompt_tokens=num_prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=num_prompt_tokens + completion_tokens,
)
data = chunk.model_dump_json(exclude_unset=True)
yield f"data: {data}\n\n"
elif final_output_type == "audio":
choices_data = self._create_audio_choice(omni_res, role, request, stream=True)
chunk = OmniChatCompletionStreamResponse(
id=request_id,
object=chunk_object_type,
created=created_time,
choices=choices_data,
model=model_name,
modality=final_output_type,
)
chunk.usage = UsageInfo(
prompt_tokens=num_prompt_tokens,
completion_tokens=0,
total_tokens=num_prompt_tokens,
)
data = chunk.model_dump_json(exclude_unset=True)
yield f"data: {data}\n\n"
else:
logger.warning(f"Unsupported streaming final output type: {final_output_type}")
continue
# once the final token is handled, if stream_options.include_usage
# is sent, send the usage
if include_usage:
completion_tokens = sum(previous_num_tokens)
final_usage = UsageInfo(
prompt_tokens=num_prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=num_prompt_tokens + completion_tokens,
)
if self.enable_prompt_tokens_details and num_cached_tokens:
final_usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=num_cached_tokens)
final_usage_chunk = ChatCompletionStreamResponse(
id=request_id,
object=chunk_object_type,
created=created_time,
choices=[],
model=model_name,
usage=final_usage,
)
final_usage_data = final_usage_chunk.model_dump_json(exclude_unset=True, exclude_none=True)
yield f"data: {final_usage_data}\n\n"
# report to FastAPI middleware aggregate usage across all choices
num_completion_tokens = sum(previous_num_tokens)
request_metadata.final_usage_info = UsageInfo(
prompt_tokens=num_prompt_tokens,
completion_tokens=num_completion_tokens,
total_tokens=num_prompt_tokens + num_completion_tokens,
)
# Log complete streaming response if output logging is enabled
if self.enable_log_outputs and self.request_logger:
# Log the complete response for each choice
for i in range(num_choices):
full_text = (
previous_texts[i]
if previous_texts and i < len(previous_texts)
else f"<streaming_complete: {previous_num_tokens[i]} tokens>"
)
self.request_logger.log_outputs(
request_id=request_id,
outputs=full_text,
output_token_ids=None, # Consider also logging all token IDs
finish_reason="streaming_complete",
is_streaming=True,
delta=False,
)
except Exception as e:
# TODO: Use a vllm-specific Validation Error
logger.exception("Error in chat completion stream generator.")
data = self.create_streaming_error_response(str(e))
yield f"data: {data}\n\n"
# Send the final done message after all response.n are finished
yield "data: [DONE]\n\n"
async def chat_completion_full_generator(
self,
request: ChatCompletionRequest,
@@ -553,6 +1299,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
for omni_outputs in final_outputs:
choices_data = []
if omni_outputs.request_output is not None and not getattr(omni_outputs.request_output, "finished", False):
continue
if omni_outputs.final_output_type == "text":
(
choices_data,
@@ -562,9 +1311,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
kv_transfer_params,
) = self._create_text_choice(request, omni_outputs, tokenizer, conversation, role)
elif omni_outputs.final_output_type == "audio":
choices_data = self._create_audio_choice(omni_outputs, role)
choices_data = self._create_audio_choice(omni_outputs, role, request, stream=False)
elif omni_outputs.final_output_type == "image":
choices_data = self._create_image_choice(omni_outputs, role)
choices_data = self._create_image_choice(omni_outputs, role, request, stream=False)
else:
logger.warning(f"Unsupported final output type: {omni_outputs.final_output_type}")
continue
@@ -857,7 +1606,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
return choices, usage, prompt_logprobs, prompt_token_ids, kv_transfer_params
def _create_audio_choice(self, omni_outputs: OmniRequestOutput, role: str):
def _create_audio_choice(
self, omni_outputs: OmniRequestOutput, role: str, request: ChatCompletionRequest, stream: bool = False
):
choices: list[ChatCompletionResponseChoice] = []
final_res = omni_outputs.request_output
audio_tensor = final_res.multimodal_output["audio"].float().detach().cpu().numpy()
@@ -893,17 +1644,29 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
)
for output in final_res.outputs:
choice_data = ChatCompletionResponseChoice(
index=output.index,
message=ChatMessage(role=role, audio=audio_obj),
logprobs=None,
finish_reason="stop",
stop_reason=None,
)
if stream:
choice_data = ChatCompletionResponseStreamChoice(
index=output.index,
delta=DeltaMessage(role=role, content=audio_base64),
logprobs=None,
finish_reason="stop",
stop_reason=output.stop_reason,
token_ids=(as_list(output.token_ids) if request.return_token_ids else None),
)
else:
choice_data = ChatCompletionResponseChoice(
index=output.index,
message=ChatMessage(role=role, audio=audio_obj),
logprobs=None,
finish_reason="stop",
stop_reason=None,
)
choices.append(choice_data)
return choices
def _create_image_choice(self, omni_outputs: OmniRequestOutput, role: str):
def _create_image_choice(
self, omni_outputs: OmniRequestOutput, role: str, request: ChatCompletionRequest, stream: bool = False
):
"""Create chat completion response choices for image output.
Converts image tensor or PIL Image output from diffusion models
@@ -1152,9 +1915,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
result = await self._diffusion_engine.generate(**gen_kwargs)
# Extract images from result
# Handle nested OmniRequestOutput structure where images might be in request_output
images: list[Image.Image] = []
if result.request_output["images"]:
images = result.request_output["images"]
images = getattr(result.request_output, "images", [])
# Convert images to base64 content
image_contents: list[dict[str, Any]] = []
@@ -449,10 +449,6 @@ class Qwen3OmniMoeForConditionalGeneration(
)
)
# # Remove EOS token if present
# if code[-1] == TALKER_CODEC_EOS_TOKEN_ID:
# code = code[:-1]
# Generate audio from codec codes
audio_tensors = []
for code in codes: