revert doc deletion

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>

.buildkite/pipeline.yml

@@ -1,6 +0,0 @@
-steps:
-  - label: ":toolbox: Qwen2.5 omni offline example"
-    commands:
-      - ".buildkite/scripts/run_qwen_offline_example.sh"
-    agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU

.buildkite/scripts/run_qwen_offline_example.sh

@@ -1,33 +0,0 @@
-#!/usr/bin/env bash
-
-set -euo pipefail
-
-# Move to repo root (script lives in .buildkite/scripts/)
-ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
-cd "${ROOT_DIR}"
-
-UV_INSTALL_DIR="${ROOT_DIR}/.local/bin"
-case ":${PATH}:" in
-  *":${UV_INSTALL_DIR}:"*) ;;
-  *) export PATH="${UV_INSTALL_DIR}:${PATH}" ;;
-esac
-
-if ! command -v uv >/dev/null 2>&1; then
-  mkdir -p "${UV_INSTALL_DIR}"
-  curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR="${UV_INSTALL_DIR}" sh
-fi
-
-UV_BIN=(uv)
-
-"${UV_BIN[@]}" python install 3.12
-"${UV_BIN[@]}" venv --python 3.12 .venv
-source .venv/bin/activate
-
-"${UV_BIN[@]}" pip install -e .
-"${UV_BIN[@]}" pip install vllm==0.11.0 --torch-backend=auto
-
-EXAMPLE_DIR="examples/offline_inference/qwen2_5_omni"
-cd "${EXAMPLE_DIR}"
-
-python end2end.py --query-type text \
-                  --no-save-results

.github/workflows/ci-all-files.yml

@@ -0,0 +1,60 @@
+name: CI All Files
+
+on:
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  test-changed:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0 # Need full history to detect changes
+
+      - name: Free disk space
+        run: |
+          # Remove large, preinstalled toolchains we don't need on the runner
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/hostedtoolcache/CodeQL
+          # Prune docker to reclaim space (safe on GH runners)
+          docker system prune -af || true
+          df -h
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Prepare environment with uv
+        run: |
+          python -m pip install --upgrade pip uv
+          uv venv --python "${{ matrix.python-version }}" --seed
+          echo "VIRTUAL_ENV=${PWD}/.venv" >> "$GITHUB_ENV"
+          echo "PATH=${PWD}/.venv/bin:$PATH" >> "$GITHUB_ENV"
+          source .venv/bin/activate
+          uv pip install -e ".[dev]"
+
+      - name: Run tests
+        run: |
+          if [ ! -d tests ]; then
+            echo "tests/ directory not found. Skipping pytest run."
+
+
+            exit 0
+          fi
+          TEST_FILES=$(find tests -type f \( -name "test_*.py" -o -name "*_test.py" \))
+          if [ -z "$TEST_FILES" ]; then
+            echo "No pytest-compatible test files detected under tests/. Skipping pytest run."
+            exit 0
+          fi
+          pytest tests/ -v -m "not slow"

.github/workflows/ci.yml

@@ -0,0 +1,124 @@
+name: CI
+
+on:
+  pull_request:
+    branches: [main]
+  workflow_dispatch:
+    inputs:
+      test_changed_only:
+        description: 'Only test changed files (true) or all files (false)'
+        required: false
+        default: 'true'
+        type: choice
+        options:
+          - 'true'
+          - 'false'
+      enable_changed_files_detection:
+        description: 'Enable changed files detection feature'
+        required: false
+        default: 'true'
+        type: choice
+        options:
+          - 'true'
+          - 'false'
+  workflow_call:
+    inputs:
+      test_changed_only:
+        description: 'Only test changed files (true) or all files (false)'
+        required: false
+        type: string
+        default: 'true'
+      enable_changed_files_detection:
+        description: 'Enable changed files detection feature'
+        required: false
+        type: string
+        default: 'true'
+
+permissions:
+  contents: read
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Free disk space
+        run: |
+          # Remove large, preinstalled toolchains we don't need on the runner
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/hostedtoolcache/CodeQL
+          # Prune docker to reclaim space (safe on GH runners)
+          docker system prune -af || true
+          df -h
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Prepare environment with uv
+        run: |
+          python -m pip install --upgrade pip uv
+          uv venv --python "${{ matrix.python-version }}" --seed
+          echo "VIRTUAL_ENV=${PWD}/.venv" >> "$GITHUB_ENV"
+          echo "PATH=${PWD}/.venv/bin:${PATH}" >> "$GITHUB_ENV"
+          # Activate env to ensure subsequent commands succeed in this step
+          source .venv/bin/activate
+          uv pip install -e ".[dev]"
+
+      - name: Determine test scope
+        id: test-scope
+        env:
+          ENABLE_CHANGED_FILES_DETECTION: ${{ inputs.enable_changed_files_detection || 'true' }}
+          TEST_CHANGED_ONLY: ${{ inputs.test_changed_only || 'true' }}
+        run: |
+          # Determine if we should test only changed files
+          # If detection is disabled, always test all files
+          if [ "$ENABLE_CHANGED_FILES_DETECTION" != "true" ]; then
+            TEST_CHANGED_ONLY="false"
+          fi
+
+          echo "test_changed_only=$TEST_CHANGED_ONLY" >> "$GITHUB_OUTPUT"
+          echo "enable_detection=$ENABLE_CHANGED_FILES_DETECTION" >> "$GITHUB_OUTPUT"
+          echo "Test scope: $TEST_CHANGED_ONLY"
+          echo "Changed files detection: $ENABLE_CHANGED_FILES_DETECTION"
+
+      - name: Get changed test files
+        id: changed-tests
+        if: steps.test-scope.outputs.test_changed_only == 'true' && steps.test-scope.outputs.enable_detection == 'true'
+        env:
+          GITHUB_BASE_REF: ${{ github.event.pull_request.base.sha || '' }}
+          GITHUB_SHA: ${{ github.sha }}
+        run: |
+          chmod +x .github/scripts/detect_changed_tests.sh
+          .github/scripts/detect_changed_tests.sh \
+            "${{ github.event.pull_request.base.sha || github.event.before || '' }}" \
+            "${{ github.sha }}" \
+            "${{ github.event_name }}"
+
+      - name: Run tests (changed files only)
+        if: steps.test-scope.outputs.test_changed_only == 'true' && steps.changed-tests.outputs.has_tests == 'true'
+        run: |
+          pytest "${{ steps.changed-tests.outputs.test_files }}" -v -m "not slow"
+
+      - name: Run tests (all files)
+        if: steps.test-scope.outputs.test_changed_only == 'false' || steps.changed-tests.outputs.has_tests == 'false'
+        run: |
+          if [ ! -d tests ]; then
+            echo "tests/ directory not found. Skipping pytest run."
+            exit 0
+          fi
+          TEST_FILES=$(find tests -type f \( -name "test_*.py" -o -name "*_test.py" \))
+          if [ -z "$TEST_FILES" ]; then
+            echo "No pytest-compatible test files detected under tests/. Skipping pytest run."
+            exit 0
+          fi
+          pytest tests/ -v -m "not slow"

.github/workflows/pre-commit.yml

@@ -3,8 +3,6 @@ name: pre-commit
 on:
   pull_request:
     branches: [main]
-  push:
-    branches: [main]
   workflow_dispatch: # for manual trigger
 
 concurrency:

.gitignore

@@ -160,7 +160,7 @@ cython_debug/
 docker/
 
 # scripts
-/scripts/
+scripts/
 
 # tests
 tests/

.pre-commit-config.yaml

@@ -1,6 +1,3 @@
-default_install_hook_types:
-  - pre-commit
-  - commit-msg
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.6.0
@@ -33,28 +30,3 @@ repos:
     hooks:
       - id: actionlint
         files: ^\.github/workflows/.*\.ya?ml$
-
-
-  - repo: local
-    hooks:
-    - id: signoff-commit
-      name: Sign-off Commit
-      entry: bash
-      args:
-        - -c
-        - |
-          if ! grep -q "^Signed-off-by: $(git config user.name) <$(git config user.email)>" "$(git rev-parse --git-path COMMIT_EDITMSG)"; then
-            printf "\nSigned-off-by: $(git config user.name) <$(git config user.email)>\n" >> "$(git rev-parse --git-path COMMIT_EDITMSG)"
-          fi
-      language: system
-      verbose: true
-      stages: [commit-msg]
-
-    # Keep `suggestion` last
-    - id: suggestion
-      name: Suggestion
-      entry: bash -c 'echo "To bypass all the pre-commit hooks, add --no-verify to git commit. To skip a specific hook, prefix the commit command with SKIP=<hook-id>."'
-      language: system
-      verbose: true
-      pass_filenames: false
-    # Insert new entries above the `suggestion` entry

README.md

@@ -1,65 +1,40 @@
-<p align="center">
-  <picture>
-    <source media="(prefers-color-scheme: dark)" src="./docs/source/logos/vllm-omni-logo-text-dark.png">
-    <img alt="vllm-omni" src="./docs/source/logos/vllm-omni-logo-text-dark.png" width=55%>
-  </picture>
-</p>
-<h3 align="center">
-Easy, fast, and cheap omni-modality model serving for everyone
-</h3>
+# vLLM-omni: Multi-modal Extension for vLLM
 
-<p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://tinyurl.com/vllm-omni-meeting"><b>Weekly Meeting</b></a> |
-</p>
+vLLM-omni is designed to extend vLLM capabilities to support multi-modality model inference and serving, particularly focusing on non-autoregressive structures and non-textual outputs.
 
----
+## 🎯 Overview
 
-*Latest News* 🔥
+Traditional vLLM systems are limited to text-based, autoregressive generation. vLLM-omni addresses this limitation by enabling support for:
 
-- [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
+- **Multi-modal Models**: Text, image, video, audio, and sensor data processing
+- **Non-autoregressive Architectures**: Diffusion Transformers (DiT) and other parallel generation models
+- **Heterogeneous Outputs**: Beyond traditional text generation to multimodal outputs
 
----
+## 🏗️ Architecture
 
-## About
+vLLM-omni is built on a modular architecture that extends vLLM's core functionality:
 
-[vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni extends its support for omni-modality model inference and serving:
 
-- **Omni-modality**: Text, image, video, and audio data processing
-- **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models
-- **Heterogeneous outputs**: from traditional text generation to multimodal outputs
+## 🚀 Key Features
 
-<p align="center">
-  <picture>
-    <img alt="vllm-omni" src="./docs/source/architecture/omni-modality model architecture.png" width=55%>
-  </picture>
-</p>
+### Multi-Engine Support
 
-vLLM-omni is fast with:
+- **Autoregressive Engine**: Traditional text generation with enhanced KV-caching
+- **Diffusion Engine**: Support for DiT models and iterative generation
+- **Hybrid Engine**: Combined AR+DiT processing pipelines
 
-- Seamless AR support for efficient KV Cache management with vLLM
-- Pipelined stage execution overlapping
-- Fully disaggregation based on omniConnector and dynamic resource allocation across stages
+### Modality Processing
 
-vLLM-omni is flexible and easy to use with:
+- **Text**: Advanced tokenization and embedding generation
+- **Image**: Vision encoder integration (CLIP, etc.)
+- **Audio**: Speech processing and audio embedding
 
-- Seamless integration with popular Hugging Face models
-- High-throughput serving with various decoding algorithms, including *parallel sampling*, *beam search*, and more
-- Tensor, pipeline, data and expert parallelism support for distributed inference
-- Streaming outputs
-- OpenAI-compatible API server
+## 📋 Supported Models
 
-vLLM seamlessly supports most popular open-source models on HuggingFace, including:
-
-- Multi-modality generation models (e.g. Qwen-image)
-- Omni-modality models (e.g. Qwen-omni)
-
-## Getting Started
-
-Install vLLM with `pip` or [from source](https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html#build-wheel-from-source):
-
-```bash
-pip install vllm
-```
+### AR + Diffusion Transformer (DiT) Models
+- Qwen-omni (Thinker-Talker-Codec structure)
+- HunyunaImage 3.0 (Ongoing)
+- Qwen-Image (Ongoing)
 
 ## 🛠️ Installation
 
@@ -109,6 +84,12 @@ uv pip install -r requirements/build.txt
 uv pip install --no-build-isolation --editable .
 ```
 
+### Verification for successful installation of vLLM
+Just run the command below. If no error, it demonstrates that the installation is successful.
+```bash
+python -c "import vllm._C"
+```
+
 ### Installation of vLLM-omni
 Install additional requirements for vllm-omni
 ```bash
@@ -117,21 +98,10 @@ cd vllm_omni
 uv pip install -e .
 ```
 
-
 ## Run examples (Qwen2.5-omni)
 
 Please check the folder of [examples](examples)
 
-## Contributing
+## Further details
 
-We welcome and value any contributions and collaborations.
-Please check out [Contributing to vLLM-omni](https://docs.vllm.ai/en/latest/contributing/index.html) for how to get involved.
-
-## Weekly Meeting
-
-- vLLM Omni Weekly Meeting: https://tinyurl.com/vllm-omni-meeting
-- Wednesday, 11:30 - 12:30 (UTC+8, [Convert to your timezone](https://dateful.com/convert/gmt8?t=15))
-
-## License
-
-Apache License 2.0, as found in the [LICENSE](./LICENSE) file.
+For detailed architecture design, see [vllm_omni_design.md](docs/architecture/vllm_omni_design.md) and [high_level_arch_design.md](docs/architecture/high_level_arch_design.md).

examples/offline_inference/qwen2_5_omni/README.md

@@ -32,11 +32,3 @@ Then run the command below.
 ```bash
 bash run_single_prompt.sh
 ```
-
-### FAQ
-
-If you encounter error about backend of librosa, try to install ffmpeg with command below.
-```
-sudo apt update
-sudo apt install ffmpeg
-```

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -1,146 +1,166 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use vLLM-omni for running offline inference
-with the correct prompt format on Qwen2.5-Omni
-"""
-
 import argparse
 import os
-from typing import NamedTuple
+import os as _os_env_toggle
+import random
 
+import numpy as np
 import soundfile as sf
-from vllm.assets.audio import AudioAsset
-from vllm.assets.image import ImageAsset
-from vllm.assets.video import VideoAsset
-from vllm.multimodal.image import convert_image_mode
+import torch
+from utils import make_omni_prompt
 from vllm.sampling_params import SamplingParams
-from vllm.utils import FlexibleArgumentParser
 
-from vllm_omni import OmniLLM
+from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+_os_env_toggle.environ["VLLM_USE_V1"] = "1"
 
 SEED = 42
+# Set all random seeds
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+torch.cuda.manual_seed(SEED)
+torch.cuda.manual_seed_all(SEED)
+
+# Make PyTorch deterministic
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+
+# Set environment variables for deterministic behavior
+os.environ["PYTHONHASHSEED"] = str(SEED)
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
 
 
-class QueryResult(NamedTuple):
-    inputs: dict
-    limit_mm_per_prompt: dict[str, int]
-
-
-# NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on
-# lower-end GPUs.
-# Unless specified, these settings have been tested to work on a single L4.
-
-default_system = (
-    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
-    "Group, capable of perceiving auditory and visual inputs, as well as "
-    "generating text and speech."
-)
-
-
-def get_text_query(question: str = None) -> QueryResult:
-    if question is None:
-        question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
-    prompt = (
-        f"<|im_start|>system\n{default_system}<|im_end|>\n"
-        "<|im_start|>user\n"
-        f"{question}<|im_end|>\n"
-        f"<|im_start|>assistant\n"
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model",
+        required=True,
+        help="Path to merged model directory (will be created if downloading).",
     )
-    return QueryResult(
-        inputs={
-            "prompt": prompt,
-        },
-        limit_mm_per_prompt={},
+    parser.add_argument("--thinker-model", type=str, default=None)
+    parser.add_argument("--talker-model", type=str, default=None)
+    parser.add_argument("--code2wav-model", type=str, default=None)
+    parser.add_argument(
+        "--hf-hub-id",
+        default="Qwen/Qwen2.5-Omni-7B",
+        help="Hugging Face repo id to download if needed.",
+    )
+    parser.add_argument("--hf-revision", default=None, help="Optional HF revision (branch/tag/commit).")
+    parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
+    parser.add_argument("--voice-type", default="default", help="Voice type, e.g., m02, f030, default.")
+    parser.add_argument(
+        "--code2wav-dir",
+        default=None,
+        help="Path to code2wav folder (contains spk_dict.pt).",
+    )
+    parser.add_argument("--dit-ckpt", default=None, help="Path to DiT checkpoint file (e.g., dit.pt).")
+    parser.add_argument("--bigvgan-ckpt", default=None, help="Path to BigVGAN checkpoint file.")
+    parser.add_argument("--dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"])
+    parser.add_argument("--max-model-len", type=int, default=32768)
+    parser.add_argument(
+        "--init-sleep-seconds",
+        type=int,
+        default=20,
+        help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
     )
 
-
-def get_mixed_modalities_query() -> QueryResult:
-    question = "What is recited in the audio? What is the content of this image? Why is this video funny?"
-    prompt = (
-        f"<|im_start|>system\n{default_system}<|im_end|>\n"
-        "<|im_start|>user\n<|audio_bos|><|AUDIO|><|audio_eos|>"
-        "<|vision_bos|><|IMAGE|><|vision_eos|>"
-        "<|vision_bos|><|VIDEO|><|vision_eos|>"
-        f"{question}<|im_end|>\n"
-        f"<|im_start|>assistant\n"
+    parser.add_argument("--thinker-only", action="store_true")
+    parser.add_argument("--text-only", action="store_true")
+    parser.add_argument("--do-wave", action="store_true")
+    parser.add_argument(
+        "--prompt_type",
+        choices=[
+            "text",
+            "audio",
+            "audio-long",
+            "audio-long-chunks",
+            "audio-long-expand-chunks",
+            "image",
+            "video",
+            "video-frames",
+            "audio-in-video",
+            "audio-in-video-v2",
+            "audio-multi-round",
+            "badcase-vl",
+            "badcase-text",
+            "badcase-image-early-stop",
+            "badcase-two-audios",
+            "badcase-two-videos",
+            "badcase-multi-round",
+            "badcase-voice-type",
+            "badcase-voice-type-v2",
+            "badcase-audio-tower-1",
+            "badcase-audio-only",
+        ],
+        default="text",
     )
-    return QueryResult(
-        inputs={
-            "prompt": prompt,
-            "multi_modal_data": {
-                "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate,
-                "image": convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
-                "video": VideoAsset(name="baby_reading", num_frames=16).np_ndarrays,
-            },
-        },
-        limit_mm_per_prompt={"audio": 1, "image": 1, "video": 1},
+    parser.add_argument("--use-torchvision", action="store_true")
+    parser.add_argument("--tokenize", action="store_true")
+    parser.add_argument(
+        "--output-wav",
+        default="output.wav",
+        help="[Deprecated] Output wav directory (use --output-dir).",
     )
-
-
-def get_use_audio_in_video_query() -> QueryResult:
-    question = "Describe the content of the video, then convert what the baby say into text."
-    prompt = (
-        f"<|im_start|>system\n{default_system}<|im_end|>\n"
-        "<|im_start|>user\n<|vision_bos|><|VIDEO|><|vision_eos|>"
-        f"{question}<|im_end|>\n"
-        f"<|im_start|>assistant\n"
+    parser.add_argument(
+        "--output-dir",
+        default="outputs",
+        help="Output directory to save text and wav files together.",
     )
-    asset = VideoAsset(name="baby_reading", num_frames=16)
-    audio = asset.get_audio(sampling_rate=16000)
-
-    return QueryResult(
-        inputs={
-            "prompt": prompt,
-            "multi_modal_data": {
-                "video": asset.np_ndarrays,
-                "audio": audio,
-            },
-            "mm_processor_kwargs": {
-                "use_audio_in_video": True,
-            },
-        },
-        limit_mm_per_prompt={"audio": 1, "video": 1},
+    parser.add_argument(
+        "--thinker-hidden-states-dir",
+        default="thinker_hidden_states",
+        help="Path to thinker hidden states directory.",
     )
-
-
-def get_multi_audios_query() -> QueryResult:
-    question = "Are these two audio clips the same?"
-    prompt = (
-        f"<|im_start|>system\n{default_system}<|im_end|>\n"
-        "<|im_start|>user\n<|audio_bos|><|AUDIO|><|audio_eos|>"
-        "<|audio_bos|><|AUDIO|><|audio_eos|>"
-        f"{question}<|im_end|>\n"
-        f"<|im_start|>assistant\n"
+    parser.add_argument(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Timeout for batching in seconds (default: 5)",
     )
-    return QueryResult(
-        inputs={
-            "prompt": prompt,
-            "multi_modal_data": {
-                "audio": [
-                    AudioAsset("winning_call").audio_and_sample_rate,
-                    AudioAsset("mary_had_lamb").audio_and_sample_rate,
-                ],
-            },
-        },
-        limit_mm_per_prompt={
-            "audio": 2,
-        },
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing stages in seconds (default: 300)",
     )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Threshold for using shared memory in bytes (default: 65536)",
+    )
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=False,
+        help="Enable writing detailed statistics (default: disabled)",
+    )
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+    )
+    args = parser.parse_args()
+    return args
 
 
-query_map = {
-    "mixed_modalities": get_mixed_modalities_query,
-    "use_audio_in_video": get_use_audio_in_video_query,
-    "multi_audios": get_multi_audios_query,
-    "text": get_text_query,
-}
-
-
-def main(args):
-    model_name = "Qwen/Qwen2.5-Omni-7B"
-    query_result = query_map[args.query_type]()
+def main():
+    args = parse_args()
+    model_name = args.model
+    try:
+        # Preferred: load from txt file (one prompt per line)
+        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
+            with open(args.txt_prompts, encoding="utf-8") as f:
+                lines = [ln.strip() for ln in f.readlines()]
+            args.prompts = [ln for ln in lines if ln != ""]
+            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
+    except Exception as e:
+        print(f"[Error] Failed to load prompts: {e}")
+        raise
 
+    if args.prompts is None:
+        raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts <file.txt> (with --prompt_type text)")
     omni_llm = OmniLLM(
         model=model_name,
         log_stats=args.enable_stats,
@@ -185,123 +205,39 @@ def main(args):
         code2wav_sampling_params,
     ]
 
-    if args.txt_prompts is None:
-        prompts = [query_result.inputs for _ in range(args.num_prompts)]
-    else:
-        assert args.query_type == "text", "txt-prompts is only supported for text query type"
-        with open(args.txt_prompts, encoding="utf-8") as f:
-            lines = [ln.strip() for ln in f.readlines()]
-            prompts = [get_text_query(ln).inputs for ln in lines if ln != ""]
-            print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
+    prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
+    omni_outputs = omni_llm.generate(prompt, sampling_params_list)
 
-    omni_outputs = omni_llm.generate(prompts, sampling_params_list)
-
-    output_dir = None
-    if args.save_results:
-        # 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)
+    # 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:
         if stage_outputs.final_output_type == "text":
             for output in stage_outputs.request_output:
                 request_id = int(output.request_id)
                 text_output = output.outputs[0].text
-                if args.save_results:
-                    # Save aligned text file per request
-                    prompt_text = prompts[request_id]["prompt"]
-                    out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
-                    lines = []
-                    lines.append("Prompt:\n")
-                    lines.append(str(prompt_text) + "\n")
-                    lines.append("vllm_text_output:\n")
-                    lines.append(str(text_output).strip() + "\n")
-                    try:
-                        with open(out_txt, "w", encoding="utf-8") as f:
-                            f.writelines(lines)
-                    except Exception as e:
-                        print(f"[Warn] Failed writing text file {out_txt}: {e}")
-                    print(f"Request ID: {request_id}, Text saved to {out_txt}")
-                else:
-                    print(f"Request ID: {request_id}, Text output:\n{text_output}")
+                # Save aligned text file per request
+                prompt_text = args.prompts[request_id]
+                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
+                lines = []
+                lines.append("Prompt:\n")
+                lines.append(str(prompt_text) + "\n")
+                lines.append("vllm_text_output:\n")
+                lines.append(str(text_output).strip() + "\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        f.writelines(lines)
+                except Exception as e:
+                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
+                print(f"Request ID: {request_id}, Text saved to {out_txt}")
         elif stage_outputs.final_output_type == "audio":
             for output in stage_outputs.request_output:
                 request_id = int(output.request_id)
                 audio_tensor = output.multimodal_output["audio"]
-                if args.save_results:
-                    output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
-                    sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
-                    print(f"Request ID: {request_id}, Saved audio to {output_wav}")
-                else:
-                    print(f"Request ID: {request_id}, Received audio output (not saved).")
-
-
-def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
-    parser.add_argument(
-        "--query-type",
-        "-q",
-        type=str,
-        default="mixed_modalities",
-        choices=query_map.keys(),
-        help="Query type.",
-    )
-    parser.add_argument(
-        "--enable-stats",
-        action="store_true",
-        default=False,
-        help="Enable writing detailed statistics (default: disabled)",
-    )
-    parser.add_argument(
-        "--init-sleep-seconds",
-        type=int,
-        default=20,
-        help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
-    )
-    parser.add_argument(
-        "--batch-timeout",
-        type=int,
-        default=5,
-        help="Timeout for batching in seconds (default: 5)",
-    )
-    parser.add_argument(
-        "--init-timeout",
-        type=int,
-        default=300,
-        help="Timeout for initializing stages in seconds (default: 300)",
-    )
-    parser.add_argument(
-        "--shm-threshold-bytes",
-        type=int,
-        default=65536,
-        help="Threshold for using shared memory in bytes (default: 65536)",
-    )
-    parser.add_argument(
-        "--save-results",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help="Save text/audio outputs to files (default: enabled). Use --no-save-results to disable.",
-    )
-    parser.add_argument(
-        "--output-wav",
-        default="output_audio",
-        help="[Deprecated] Output wav directory (use --output-dir).",
-    )
-    parser.add_argument(
-        "--num-prompts",
-        type=int,
-        default=1,
-        help="Number of prompts to generate.",
-    )
-    parser.add_argument(
-        "--txt-prompts",
-        type=str,
-        default=None,
-        help="Path to a .txt file with one prompt per line (preferred).",
-    )
-
-    return parser.parse_args()
+                output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
+                sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
+                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
 
 if __name__ == "__main__":
-    args = parse_args()
-    main(args)
+    main()

examples/offline_inference/qwen2_5_omni/processing_omni.py

@@ -0,0 +1,367 @@
+from __future__ import annotations
+
+import base64
+import logging
+import math
+import os
+import time
+import warnings
+from functools import lru_cache
+from io import BytesIO
+
+import requests
+import torch
+import torchvision
+from packaging import version
+from PIL import Image
+from torchvision import io, transforms
+from torchvision.transforms import InterpolationMode
+
+logger = logging.getLogger(__name__)
+
+IMAGE_FACTOR = 28
+MIN_PIXELS = 4 * 28 * 28
+MAX_PIXELS = 16384 * 28 * 28
+MAX_RATIO = 200
+
+VIDEO_MIN_PIXELS = 128 * 28 * 28
+VIDEO_MAX_PIXELS = 768 * 28 * 28
+VIDEO_TOTAL_PIXELS = 24576 * 28 * 28
+FRAME_FACTOR = 2
+FPS = 2.0
+FPS_MIN_FRAMES = 4
+FPS_MAX_FRAMES = 768
+
+temporal_patch_size = 2
+spatial_patch_size = 14
+spatial_merge_size = 2
+
+
+def round_by_factor(number: int, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is
+    divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is
+    divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int = IMAGE_FACTOR,
+    min_pixels: int = MIN_PIXELS,
+    max_pixels: int = MAX_PIXELS,
+) -> tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    if max(height, width) / min(height, width) > MAX_RATIO:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+
+def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image:
+    if "image" in ele:
+        image = ele["image"]
+    else:
+        image = ele["image_url"]
+    image_obj = None
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif image.startswith("http://") or image.startswith("https://"):
+        image_obj = Image.open(requests.get(image, stream=True).raw)
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            image_obj = Image.open(BytesIO(data))
+    else:
+        image_obj = Image.open(image)
+    if image_obj is None:
+        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
+    image = image_obj.convert("RGB")
+    # resize
+    if "resized_height" in ele and "resized_width" in ele:
+        resized_height, resized_width = smart_resize(
+            ele["resized_height"],
+            ele["resized_width"],
+            factor=size_factor,
+        )
+    else:
+        width, height = image.size
+        min_pixels = ele.get("min_pixels", MIN_PIXELS)
+        max_pixels = ele.get("max_pixels", MAX_PIXELS)
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor=size_factor,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+        )
+    image = image.resize((resized_width, resized_height))
+
+    return image
+
+
+def smart_nframes(
+    ele: dict,
+    total_frames: int,
+    video_fps: int | float,
+) -> int:
+    """calculate the number of frames for video used for model inputs.
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+            support either `fps` or `nframes`:
+                - nframes: the number of frames to extract for model inputs.
+                - fps: the fps to extract frames for model inputs.
+                    - min_frames: the minimum number of frames of the video,
+                      only used when fps is provided.
+                    - max_frames: the maximum number of frames of the video,
+                      only used when fps is provided.
+        total_frames (int): the original total number of frames of the video.
+        video_fps (int | float): the original fps of the video.
+
+    Raises:
+        ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
+
+    Returns:
+        int: the number of frames for video used for model inputs.
+    """
+    assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`"
+    if "nframes" in ele:
+        nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
+    else:
+        fps = ele.get("fps", FPS)
+        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
+        max_frames = floor_by_factor(
+            ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)),
+            FRAME_FACTOR,
+        )
+        nframes = total_frames / video_fps * fps
+        nframes = min(max(nframes, min_frames), max_frames)
+        nframes = round_by_factor(nframes, FRAME_FACTOR)
+    if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
+        raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.")
+    return nframes
+
+
+def _read_video_torchvision(
+    ele: dict,
+) -> torch.Tensor:
+    """read video using torchvision.io.read_video
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://",
+              "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        torch.Tensor: the video tensor with shape (T, C, H, W).
+    """
+    video_path = ele["video"]
+    if version.parse(torchvision.__version__) < version.parse("0.19.0"):
+        if "http://" in video_path or "https://" in video_path:
+            warnings.warn("torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.")
+        if "file://" in video_path:
+            video_path = video_path[7:]
+    st = time.time()
+    video, audio, info = io.read_video(
+        video_path,
+        start_pts=ele.get("video_start", 0.0),
+        end_pts=ele.get("video_end", None),
+        pts_unit="sec",
+        output_format="TCHW",
+    )
+    total_frames, video_fps = video.size(0), info["video_fps"]
+    total_duration = round(total_frames / video_fps, 3)
+    logger.info(
+        f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, duration={total_duration}s, time={time.time() - st:.3f}s"
+    )
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(0, total_frames - 1, nframes).round().long()
+    video = video[idx]
+    return video, total_duration, nframes
+
+
+def is_decord_available() -> bool:
+    import importlib.util
+
+    return importlib.util.find_spec("decord") is not None
+
+
+def _read_video_decord(
+    ele: dict,
+) -> torch.Tensor:
+    """read video using decord.VideoReader
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://",
+              "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        torch.Tensor: the video tensor with shape (T, C, H, W).
+    """
+    import decord
+
+    video_path = ele["video"]
+    st = time.time()
+    vr = decord.VideoReader(video_path)
+    # TODO: support start_pts and end_pts
+    if "video_start" in ele or "video_end" in ele:
+        raise NotImplementedError("not support start_pts and end_pts in decord for now.")
+    total_frames, video_fps = len(vr), vr.get_avg_fps()
+    total_duration = round(total_frames / video_fps, 3)
+    logger.info(f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
+    video = vr.get_batch(idx).asnumpy()
+    video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
+    return video, total_duration, nframes
+
+
+VIDEO_READER_BACKENDS = {
+    "decord": _read_video_decord,
+    "torchvision": _read_video_torchvision,
+}
+
+FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
+
+
+@lru_cache(maxsize=1)
+def get_video_reader_backend() -> str:
+    if FORCE_QWENVL_VIDEO_READER is not None:
+        video_reader_backend = FORCE_QWENVL_VIDEO_READER
+    elif is_decord_available():
+        video_reader_backend = "decord"
+    else:
+        video_reader_backend = "torchvision"
+    # print(f"qwen-vl-utils using {video_reader_backend} to read video.",
+    #       file=sys.stderr)
+    return video_reader_backend
+
+
+def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR) -> torch.Tensor | list[Image.Image]:
+    if isinstance(ele["video"], str):
+        video_reader_backend = get_video_reader_backend()
+        video, total_dur, nframes = VIDEO_READER_BACKENDS[video_reader_backend](ele)
+        frame_timestamps = total_dur * torch.arange(1, nframes + 1) / nframes
+        grid_timestamps = frame_timestamps[::FRAME_FACTOR]
+        second_per_grid = grid_timestamps[1] - grid_timestamps[0]
+        nframes, _, height, width = video.shape
+        min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
+        total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
+        max_pixels = max(
+            min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR),
+            int(min_pixels * 1.05),
+        )
+        max_pixels = ele.get("max_pixels", max_pixels)
+        # min_pixels = (factor ** 2) * 52
+        # max_pixels = (factor ** 2) * min(768, (16384 / nframes * temporal_patch_size))
+        if "resized_height" in ele and "resized_width" in ele:
+            resized_height, resized_width = smart_resize(
+                ele["resized_height"],
+                ele["resized_width"],
+                factor=image_factor,
+            )
+        else:
+            resized_height, resized_width = smart_resize(
+                height,
+                width,
+                factor=image_factor,
+                min_pixels=min_pixels,
+                max_pixels=max_pixels,
+            )
+        video = transforms.functional.resize(
+            video,
+            [resized_height, resized_width],
+            interpolation=InterpolationMode.BICUBIC,
+            antialias=True,
+        ).float()
+        return video, total_dur, nframes, second_per_grid
+    else:
+        assert isinstance(ele["video"], (list, tuple))
+        process_info = ele.copy()
+        process_info.pop("type", None)
+        process_info.pop("video", None)
+        images = [
+            fetch_image({"image": video_element, **process_info}, size_factor=image_factor)
+            for video_element in ele["video"]
+        ]
+        nframes = ceil_by_factor(len(images), FRAME_FACTOR)
+        if len(images) < nframes:
+            images.extend([images[-1]] * (nframes - len(images)))
+        return images, None, None, None
+
+
+def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]:
+    vision_infos = []
+    if isinstance(conversations[0], dict):
+        conversations = [conversations]
+    for conversation in conversations:
+        for message in conversation:
+            if isinstance(message["content"], list):
+                for ele in message["content"]:
+                    if (
+                        "image" in ele
+                        or "image_url" in ele
+                        or "video" in ele
+                        or ele["type"] in ("image", "image_url", "video")
+                    ):
+                        vision_infos.append(ele)
+    return vision_infos
+
+
+def process_vision_info(
+    conversations: list[dict] | list[list[dict]],
+) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None]:
+    vision_infos = extract_vision_info(conversations)
+    # Read images or videos
+    image_inputs = []
+    video_inputs = []
+    for vision_info in vision_infos:
+        if "image" in vision_info or "image_url" in vision_info:
+            image_inputs.append(fetch_image(vision_info))
+        elif "video" in vision_info:
+            video_inputs.append(fetch_video(vision_info))
+        else:
+            raise ValueError("image, image_url or video should in content.")
+    if len(image_inputs) == 0:
+        image_inputs = None
+    if len(video_inputs) == 0:
+        video_inputs = None
+    return image_inputs, video_inputs

examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh

@@ -1,3 +1,8 @@
-python end2end.py --output-wav output_audio \
-                  --query-type text \
-                  --txt-prompts top10.txt
+python end2end.py --model Qwen/Qwen2.5-Omni-7B \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --txt-prompts top100.txt

examples/offline_inference/qwen2_5_omni/run_single_prompt.sh

@@ -1,2 +1,8 @@
-python end2end.py --output-wav output_audio \
-                  --query-type use_audio_in_video
+python end2end.py --model Qwen/Qwen2.5-Omni-7B \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."

examples/offline_inference/qwen2_5_omni/utils.py

@@ -0,0 +1,312 @@
+import tempfile
+from typing import Optional, Union
+from urllib.request import urlopen
+
+import librosa
+import requests
+import resampy
+import soundfile as sf
+import torch
+import torchvision.io
+from processing_omni import fetch_image, fetch_video
+from transformers import AutoConfig, AutoProcessor
+from vllm.inputs import TextPrompt
+
+from vllm_omni.inputs.data import OmniTokensPrompt
+
+# Simple caches to avoid repeated heavy HF loads per prompt
+_PROCESSOR_CACHE: dict[str, "AutoProcessor"] = {}
+_CONFIG_CACHE: dict[str, "AutoConfig"] = {}
+
+
+def get_system_prompt():
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, "
+                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
+                    "as well as generating text and speech."
+                ),
+            }
+        ],
+    }
+
+
+def resample_wav_to_16khz(input_filepath):
+    data, original_sample_rate = sf.read(input_filepath)
+    # Only use the first channel
+    if len(data.shape) > 1:
+        data = data[:, 0]
+    # resample to 16kHz
+    data_resampled = resampy.resample(data, sr_orig=original_sample_rate, sr_new=16000)
+    return data_resampled
+
+
+def fetch_and_read_video(args, video_url: str, fps=2):
+    def read_video_with_torchvision(video_file_name: str):
+        video, audio, info = torchvision.io.read_video(
+            video_file_name,
+            start_pts=0.0,
+            end_pts=None,
+            pts_unit="sec",
+            output_format="TCHW",
+        )
+
+        total_frames, video_fps = video.size(0), info["video_fps"]
+        total_duration = round(total_frames / video_fps, 3)
+        nframes = int(total_frames / video_fps * fps)
+
+        frame_timestamps = total_duration * torch.arange(1, nframes + 1) / nframes
+        grid_timestamps = frame_timestamps[::2]
+        second_per_grid = grid_timestamps[1] - grid_timestamps[0]
+
+        idx = torch.linspace(0, video.size(0) - 1, nframes).round().long()
+        video = video[idx]
+
+        if args.legacy_omni_video:
+            return [video, total_duration, nframes, second_per_grid.item()]
+        else:
+            return video
+
+    def read_video_with_transformers(video_file_name: Union[str, list[str]]):
+        video, total_duration, nframes, second_per_grid = fetch_video({"video": video_file_name})
+        if total_duration is None and nframes is None:
+            nframes = len(video)
+            total_duration = 0.5 * nframes
+            second_per_grid = 1.0
+        if args.legacy_omni_video:
+            return [video, total_duration, nframes, second_per_grid]
+        else:
+            return video
+
+    def read_video(video_file_name: str):
+        if args.use_torchvision:
+            return read_video_with_torchvision(video_file_name)
+        else:
+            return read_video_with_transformers(video_file_name)
+
+    if isinstance(video_url, str) and video_url.startswith("http"):
+        with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
+            resp = requests.get(video_url)
+            assert resp.status_code == requests.codes.ok, (
+                f"Failed to fetch video from {video_url}, status_code:{resp.status_code}, resp:{resp}"
+            )
+
+            temp_video_file.write(urlopen(video_url).read())
+            temp_video_file_path = temp_video_file.name
+            video_file_name = temp_video_file_path
+            return read_video(video_file_name)
+    else:
+        video_file_name = video_url
+        return read_video(video_file_name)
+
+
+def make_inputs_qwen2_omni(
+    args,
+    messages: list[dict[str, Union[str, list[dict[str, str]]]]],
+    use_audio_in_video: Optional[bool] = False,
+    tokenize: bool = False,
+) -> Union[OmniTokensPrompt, TextPrompt]:
+    from transformers import AutoConfig, AutoProcessor
+
+    # Cached processor/config to prevent per-prompt reloading and repeated warnings
+    if args.model not in _PROCESSOR_CACHE:
+        _PROCESSOR_CACHE[args.model] = AutoProcessor.from_pretrained(args.model)
+    processor = _PROCESSOR_CACHE[args.model]
+
+    config = _CONFIG_CACHE.get(args.model)
+    if config is None:
+        try:
+            config = AutoConfig.from_pretrained(args.model)
+        except Exception:
+            config = None
+        _CONFIG_CACHE[args.model] = config  # cache even if None to avoid retry storms
+
+    # Decide legacy flag only once based on config (default True if unknown)
+    if getattr(args, "legacy_omni_video", None) is None:
+        if config is not None and hasattr(config, "architectures"):
+            args.legacy_omni_video = "Qwen2_5OmniModel" not in config.architectures
+        else:
+            args.legacy_omni_video = True
+
+    audios, images, videos = [], [], []
+    for message in messages:
+        if not isinstance(message["content"], list):
+            message["content"] = [
+                {
+                    "type": "text",
+                    "text": message["content"],
+                }
+            ]
+        index, num_contents = 0, len(message["content"])
+        while index < num_contents:
+            ele = message["content"][index]
+            if "type" not in ele:
+                if "text" in ele:
+                    ele["type"] = "text"
+                elif "audio" in ele:
+                    ele["type"] = "audio"
+                elif "audio_url" in ele:
+                    ele["type"] = "audio_url"
+                elif "image" in ele:
+                    ele["type"] = "image"
+                elif "image_url" in ele:
+                    ele["type"] = "image_url"
+                elif "video" in ele:
+                    ele["type"] = "video"
+                elif "video_url" in ele:
+                    ele["type"] = "video_url"
+                else:
+                    raise ValueError(f"Unknown ele: {ele}")
+
+            if ele["type"] == "audio" or ele["type"] == "audio_url":
+                if "audio_url" in ele:
+                    audio_key = "audio_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_audio_file:
+                        temp_audio_file.write(urlopen(ele[audio_key]).read())
+                        temp_audio_file_path = temp_audio_file.name
+                        audios.append(resample_wav_to_16khz(temp_audio_file_path))
+                        ele["audio"] = temp_audio_file_path
+                elif "audio" in ele:
+                    audio_key = "audio"
+                    audios.append(resample_wav_to_16khz(ele[audio_key]))
+                else:
+                    raise ValueError(f"Unknown ele {ele}")
+            elif use_audio_in_video and (ele["type"] == "video" or ele["type"] == "video_url"):
+                # use video as audio as well
+                if "video_url" in ele:
+                    audio_key = "video_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
+                        temp_video_file.write(urlopen(ele[audio_key]).read())
+                        temp_video_file_path = temp_video_file.name
+                        ele[audio_key] = temp_video_file_path
+                        audios.append(librosa.load(temp_video_file_path, sr=16000)[0])
+                        videos.append(fetch_and_read_video(args, temp_video_file_path))
+                        ele["video"] = temp_video_file_path
+                elif "video" in ele:
+                    audio_key = "video"
+                    audios.append(librosa.load(ele[audio_key], sr=16000)[0])
+                    videos.append(fetch_and_read_video(args, audio_key))
+                else:
+                    raise ValueError(f"Unknown ele {ele}")
+                # insert a audio after the video
+                message["content"].insert(
+                    index + 1,
+                    {
+                        "type": "audio",
+                        "audio": ele[audio_key],
+                    },
+                )
+                # no need to load the added audio again
+                index += 1
+            elif ele["type"] == "video" or ele["type"] == "video_url":
+                if "video_url" in ele:
+                    video_key = "video_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
+                        temp_video_file.write(urlopen(ele["video_url"]).read())
+                        temp_video_file_path = temp_video_file.name
+                        videos.append(fetch_and_read_video(args, temp_video_file))
+                        ele["video"] = temp_video_file_path
+                else:
+                    video_key = "video"
+                    videos.append(fetch_and_read_video(args, ele[video_key]))
+            elif ele["type"] == "image" or ele["type"] == "image_url":
+                images.append(fetch_image(ele))
+
+            # move to the next content
+            index += 1
+
+    prompt = processor.apply_chat_template(
+        messages,
+        tokenize=tokenize,
+        add_generation_prompt=True,
+        add_vision_id=True,
+    )
+
+    audios = audios if len(audios) > 0 else None
+    images = images if len(images) > 0 else None
+    videos = videos if len(videos) > 0 else None
+
+    multi_modal_data = {}
+    if audios:
+        multi_modal_data["audio"] = audios
+    if images:
+        multi_modal_data["image"] = images
+    if videos:
+        multi_modal_data["video"] = videos
+
+    if isinstance(prompt, list) and isinstance(prompt[0], (list, str)):
+        prompt = prompt[0]
+
+    if tokenize:
+        return OmniTokensPrompt(
+            prompt_token_ids=prompt,
+            multi_modal_data=multi_modal_data,
+        )
+    else:
+        return TextPrompt(
+            prompt=prompt,
+            multi_modal_data=multi_modal_data,
+        )
+
+
+def make_text_prompt(args, prompt):
+    messages = [
+        get_system_prompt(),
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": prompt},
+            ],
+        },
+    ]
+
+    prompt = make_inputs_qwen2_omni(args, messages, tokenize=args.tokenize)
+    return prompt
+
+
+def make_audio_in_video_v2_prompt(args):
+    messages = [
+        {
+            "role": "system",
+            "content": [
+                {
+                    "type": "text",
+                    "text": (
+                        "You are Qwen, a virtual human developed by the Qwen Team, "
+                        "Alibaba Group, capable of perceiving auditory and visual "
+                        "inputs, as well as generating text and speech."
+                    ),
+                }
+            ],
+        },
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "video_url",
+                    "video_url": ("https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-Omni/draw_small.mp4"),
+                },
+            ],
+        },
+    ]
+    prompt = make_inputs_qwen2_omni(
+        args,
+        messages,
+        use_audio_in_video=True,
+        tokenize=args.tokenize,
+    )
+    return prompt
+
+
+def make_omni_prompt(args, prompt=None) -> Union[OmniTokensPrompt, list[OmniTokensPrompt]]:
+    if args.prompt_type == "text":
+        prompt = make_text_prompt(args, prompt)
+    elif args.prompt_type == "audio-in-video-v2":
+        prompt = make_audio_in_video_v2_prompt(args)
+    else:
+        raise ValueError(f"Unsupported prompt type: {args.prompt_type}")
+    return prompt

examples/online_serving/README.md

@@ -23,55 +23,10 @@ cd examples/online_serving
 
 Send request via python
 ```bash
-python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
+python openai_chat_completion_client_for_multimodal_generation.py
 ```
 
 Send request via curl
 ```bash
-bash run_curl_multimodal_generation.sh mixed_modalities
+bash run_curl_multimodal_generation.sh
 ```
-
-### FAQ
-
-If you encounter error about backend of librosa, try to install ffmpeg with command below.
-```
-sudo apt update
-sudo apt install ffmpeg
-```
-
-## Run Local Web UI Demo
-
-This Web UI demo allows users to interact with the model through a web browser.
-
-### Running Gradio Demo
-
-Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
-
-```bash
-python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861
-```
-
-Then open `http://localhost:7861/` on your local browser to interact with the web UI.
-
-
-### Options
-
-The gradio demo also supports running with an existing API server and can be customized with the following arguments.
-
-
-```bash
-python gradio_demo.py \
-    --model Qwen/Qwen2.5-Omni-7B \
-    --use-api-server \
-    --api-base http://localhost:8091/v1 \
-    --ip 127.0.0.1 \
-    --port 7861
-```
-
-- `--model`: Model name
-- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmniLLM locally.
-- `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1)
-- `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
-- `--port`: Port for Gradio server (default: 7861)
-- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
-- `--share`: Share the Gradio demo publicly (creates a public link)

examples/online_serving/gradio_demo.py

@@ -1,535 +0,0 @@
-import argparse
-import os
-import random
-import signal
-import sys
-from pathlib import Path
-from types import SimpleNamespace
-from typing import Any, Optional
-
-import gradio as gr
-import numpy as np
-import soundfile as sf
-import torch
-from PIL import Image
-from vllm.assets.video import video_get_metadata, video_to_ndarrays
-from vllm.sampling_params import SamplingParams
-
-from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
-
-# Import utils from offline inference example
-sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../offline_inference/qwen2_5_omni"))
-try:
-    from utils import make_omni_prompt
-except ImportError:
-    # Fallback if utils doesn't exist - we'll build prompts directly
-    make_omni_prompt = None
-
-SEED = 42
-ASYNC_INIT_TIMEOUT = 600
-
-SUPPORTED_MODELS: dict[str, dict[str, Any]] = {
-    "Qwen/Qwen2.5-Omni-7B": {
-        "sampling_params": {
-            "thinker": {
-                "temperature": 0.0,
-                "top_p": 1.0,
-                "top_k": -1,
-                "max_tokens": 2048,
-                "detokenize": True,
-                "repetition_penalty": 1.1,
-            },
-            "talker": {
-                "temperature": 0.0,
-                "top_p": 1.0,
-                "top_k": -1,
-                "max_tokens": 2048,
-                "detokenize": True,
-                "repetition_penalty": 1.1,
-                "stop_token_ids": [8294],
-            },
-            "code2wav": {
-                "temperature": 0.0,
-                "top_p": 1.0,
-                "top_k": -1,
-                "max_tokens": 2048,
-                "detokenize": True,
-                "repetition_penalty": 1.1,
-            },
-        },
-    },
-}
-# Ensure deterministic behavior across runs.
-random.seed(SEED)
-np.random.seed(SEED)
-torch.manual_seed(SEED)
-torch.cuda.manual_seed(SEED)
-torch.cuda.manual_seed_all(SEED)
-torch.backends.cudnn.deterministic = True
-torch.backends.cudnn.benchmark = False
-os.environ["PYTHONHASHSEED"] = str(SEED)
-os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Gradio demo for Qwen2.5-Omni online inference.")
-    parser.add_argument(
-        "--model",
-        default="Qwen/Qwen2.5-Omni-7B",
-        help="Path to local model checkpoint.",
-    )
-    parser.add_argument(
-        "--ip",
-        default="127.0.0.1",
-        help="Host/IP for gradio `launch`.",
-    )
-    parser.add_argument("--port", type=int, default=7861, help="Port for gradio `launch`.")
-    parser.add_argument("--share", action="store_true", help="Share the Gradio demo publicly.")
-    parser.add_argument(
-        "--stage-configs-path",
-        type=str,
-        default=None,
-        help="Path to custom stage configs YAML file (optional).",
-    )
-    return parser.parse_args()
-
-
-def build_sampling_params(seed: int, model_key: str) -> list[SamplingParams]:
-    """Build SamplingParams objects by reusing the dict definitions."""
-    return [SamplingParams(**params_dict) for params_dict in build_sampling_params_dict(seed, model_key)]
-
-
-def build_sampling_params_dict(seed: int, model_key: str) -> list[dict]:
-    """Build sampling params as dict for HTTP API mode."""
-    model_conf = SUPPORTED_MODELS.get(model_key)
-    if model_conf is None:
-        raise ValueError(f"Unsupported model '{model_key}'")
-
-    sampling_templates: dict[str, dict[str, Any]] = model_conf["sampling_params"]
-    sampling_params: list[dict] = []
-    for stage_name, template in sampling_templates.items():
-        params = dict(template)
-        params["seed"] = seed
-        sampling_params.append(params)
-    return sampling_params
-
-
-def create_prompt_args(base_args: argparse.Namespace) -> SimpleNamespace:
-    # The prompt builder expects a minimal namespace with these attributes.
-    return SimpleNamespace(
-        model=base_args.model,
-        prompt_type="text",
-        tokenize=True,
-        use_torchvision=True,
-        legacy_omni_video=False,
-    )
-
-
-def process_audio_file(
-    audio_file: Optional[Any],
-) -> Optional[tuple[np.ndarray, int]]:
-    """Normalize Gradio audio input to (np.ndarray, sample_rate)."""
-    if audio_file is None:
-        return None
-
-    sample_rate: Optional[int] = None
-    audio_np: Optional[np.ndarray] = None
-
-    def _load_from_path(path_str: str) -> Optional[tuple[np.ndarray, int]]:
-        if not path_str:
-            return None
-        path = Path(path_str)
-        if not path.exists():
-            return None
-        data, sr = sf.read(path)
-        if data.ndim > 1:
-            data = data[:, 0]
-        return data.astype(np.float32), int(sr)
-
-    if isinstance(audio_file, tuple):
-        if len(audio_file) == 2:
-            first, second = audio_file
-            # Case 1: (sample_rate, np.ndarray)
-            if isinstance(first, (int, float)) and isinstance(second, np.ndarray):
-                sample_rate = int(first)
-                audio_np = second
-            # Case 2: (filepath, (sample_rate, np.ndarray or list))
-            elif isinstance(first, str):
-                if isinstance(second, tuple) and len(second) == 2:
-                    sr_candidate, data_candidate = second
-                    if isinstance(sr_candidate, (int, float)) and isinstance(data_candidate, np.ndarray):
-                        sample_rate = int(sr_candidate)
-                        audio_np = data_candidate
-                if audio_np is None:
-                    loaded = _load_from_path(first)
-                    if loaded is not None:
-                        audio_np, sample_rate = loaded
-            # Case 3: (None, (sample_rate, np.ndarray))
-            elif first is None and isinstance(second, tuple) and len(second) == 2:
-                sr_candidate, data_candidate = second
-                if isinstance(sr_candidate, (int, float)) and isinstance(data_candidate, np.ndarray):
-                    sample_rate = int(sr_candidate)
-                    audio_np = data_candidate
-        elif len(audio_file) == 1 and isinstance(audio_file[0], str):
-            loaded = _load_from_path(audio_file[0])
-            if loaded is not None:
-                audio_np, sample_rate = loaded
-    elif isinstance(audio_file, str):
-        loaded = _load_from_path(audio_file)
-        if loaded is not None:
-            audio_np, sample_rate = loaded
-
-    if audio_np is None or sample_rate is None:
-        return None
-
-    if audio_np.ndim > 1:
-        audio_np = audio_np[:, 0]
-
-    return audio_np.astype(np.float32), sample_rate
-
-
-def process_image_file(image_file: Optional[Image.Image]) -> Optional[Image.Image]:
-    """Process image file from Gradio input.
-
-    Returns:
-        PIL Image in RGB mode or None if no image provided.
-    """
-    if image_file is None:
-        return None
-    # Convert to RGB if needed
-    if image_file.mode != "RGB":
-        image_file = image_file.convert("RGB")
-    return image_file
-
-
-def process_video_file(
-    video_file: Optional[str],
-    enable_audio_in_video: bool = False,
-    max_frames: int = 32,
-) -> Optional[tuple[np.ndarray, dict[str, Any], Optional[tuple[np.ndarray, int]]]]:
-    """Process video file and optionally extract audio track."""
-    if video_file is None:
-        return None
-
-    video_path = Path(video_file)
-    if not video_path.exists():
-        print(f"Video file not found: {video_path}")
-        return None
-
-    try:
-        frames = video_to_ndarrays(str(video_path), num_frames=max_frames)
-        metadata = video_get_metadata(str(video_path), num_frames=frames.shape[0])
-    except Exception as exc:  # pylint: disable=broad-except
-        print(f"Failed to decode video {video_path}: {exc}")
-        return None
-
-    audio_tuple: Optional[tuple[np.ndarray, int]] = None
-    if enable_audio_in_video:
-        try:
-            import librosa  # type: ignore import
-
-            audio_signal, sampling_rate = librosa.load(str(video_path), sr=16000)
-            audio_tuple = (audio_signal.astype(np.float32), sampling_rate)
-        except Exception as exc:  # pylint: disable=broad-except
-            print(f"Failed to extract audio from video {video_path}: {exc}")
-
-    return frames, metadata, audio_tuple
-
-
-async def run_inference_async_omni_llm(
-    omni_llm: AsyncOmniLLM,
-    sampling_params: list[SamplingParams],
-    prompt_args_template: SimpleNamespace,
-    user_prompt: str,
-    audio_file: Optional[tuple[str, tuple[int, np.ndarray]]] = None,
-    image_file: Optional[Image.Image] = None,
-    video_file: Optional[str] = None,
-    use_audio_in_video: bool = False,
-):
-    """Run inference using AsyncOmniLLM directly with multimodal support."""
-    if not user_prompt.strip() and not audio_file and not image_file and not video_file:
-        return "Please provide at least a text prompt or multimodal input.", None
-
-    try:
-        # Build prompt with multimodal data
-        prompt_args = SimpleNamespace(**prompt_args_template.__dict__)
-
-        # Process multimodal inputs
-        multi_modal_data = {}
-        mm_processor_kwargs = {}
-
-        # Process audio
-        audio_data = process_audio_file(audio_file)
-        if audio_data is not None:
-            multi_modal_data["audio"] = audio_data
-
-        # Process image
-        image_data = process_image_file(image_file)
-        if image_data is not None:
-            multi_modal_data["image"] = image_data
-
-        # Process video
-        video_payload = process_video_file(video_file, enable_audio_in_video=use_audio_in_video)
-        if video_payload is not None:
-            video_frames, video_metadata, extracted_audio = video_payload
-            video_entry: Any
-            if video_metadata:
-                video_entry = (video_frames, video_metadata)
-            else:
-                video_entry = video_frames
-            multi_modal_data["video"] = video_entry
-            if use_audio_in_video and extracted_audio is not None and "audio" not in multi_modal_data:
-                multi_modal_data["audio"] = extracted_audio
-                mm_processor_kwargs["use_audio_in_video"] = True
-
-        # Build the prompt input
-        if make_omni_prompt is not None:
-            omni_prompt = make_omni_prompt(prompt_args, user_prompt)
-            # Add multimodal data if present
-            if multi_modal_data:
-                if isinstance(omni_prompt, dict):
-                    omni_prompt["multi_modal_data"] = multi_modal_data
-                    if mm_processor_kwargs:
-                        omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs
-                else:
-                    # If make_omni_prompt returns a string, we need to wrap it
-                    omni_prompt = {
-                        "prompt": omni_prompt,
-                        "multi_modal_data": multi_modal_data,
-                    }
-                    if mm_processor_kwargs:
-                        omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs
-        else:
-            # Fallback: build prompt directly
-            default_system = (
-                "You are Qwen, a virtual human developed by the Qwen Team, "
-                "Alibaba Group, capable of perceiving auditory and visual inputs, "
-                "as well as generating text and speech."
-            )
-            prompt = f"<|im_start|>system\n{default_system}<|im_end|>\n<|im_start|>user\n"
-            if audio_data:
-                prompt += "<|audio_bos|><|AUDIO|><|audio_eos|>"
-            if image_data:
-                prompt += "<|vision_bos|><|IMAGE|><|vision_eos|>"
-            if video_payload is not None:
-                prompt += "<|vision_bos|><|VIDEO|><|vision_eos|>"
-            if user_prompt.strip():
-                prompt += f"{user_prompt}"
-            prompt += "<|im_end|>\n<|im_start|>assistant\n"
-
-            omni_prompt = {
-                "prompt": prompt,
-                "multi_modal_data": multi_modal_data,
-            }
-            if mm_processor_kwargs:
-                omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs
-
-        request_id = "0"
-        text_outputs: list[str] = []
-        audio_output = None
-
-        async for stage_outputs in omni_llm.generate(
-            prompt=omni_prompt,
-            request_id=request_id,
-            sampling_params_list=sampling_params,
-        ):
-            # stage_outputs.request_output is a RequestOutput object, not a list
-            request_output = stage_outputs.request_output
-            if stage_outputs.final_output_type == "text":
-                if request_output.outputs:
-                    for output in request_output.outputs:
-                        if output.text:
-                            text_outputs.append(output.text)
-            elif stage_outputs.final_output_type == "audio":
-                # multimodal_output is on the RequestOutput object
-                # See vllm_omni/entrypoints/openai/serving_chat.py:680 for reference
-                if hasattr(request_output, "multimodal_output") and request_output.multimodal_output:
-                    if "audio" in request_output.multimodal_output:
-                        audio_tensor = request_output.multimodal_output["audio"]
-                        # Ensure audio is 1D (flatten if needed)
-                        if hasattr(audio_tensor, "ndim") and audio_tensor.ndim > 1:
-                            audio_tensor = audio_tensor.flatten()
-                        audio_np = audio_tensor.detach().cpu().numpy()
-                        audio_output = (
-                            24000,  # sampling rate in Hz
-                            audio_np,
-                        )
-
-        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
-
-
-def build_interface(
-    omni_llm: AsyncOmniLLM,
-    sampling_params: list[SamplingParams],
-    prompt_args_template: SimpleNamespace,
-    model: str,
-):
-    """Build Gradio interface for AsyncOmniLLM mode."""
-
-    async def run_inference(
-        user_prompt: str,
-        audio_file: Optional[tuple[str, tuple[int, np.ndarray]]],
-        image_file: Optional[Image.Image],
-        video_file: Optional[str],
-        use_audio_in_video: bool,
-    ):
-        return await run_inference_async_omni_llm(
-            omni_llm,
-            sampling_params,
-            prompt_args_template,
-            user_prompt,
-            audio_file,
-            image_file,
-            video_file,
-            use_audio_in_video,
-        )
-
-    css = """
-    .media-input-container {
-        display: flex;
-        gap: 10px;
-    }
-    .media-input-container > div {
-        flex: 1;
-    }
-    .media-input-container .image-input,
-    .media-input-container .audio-input {
-        height: 300px;
-    }
-    .media-input-container .video-column {
-        height: 300px;
-        display: flex;
-        flex-direction: column;
-    }
-    .media-input-container .video-input {
-        flex: 1;
-        min-height: 0;
-    }
-    #generate-btn button {
-        width: 100%;
-    }
-    """
-
-    with gr.Blocks(css=css) as demo:
-        gr.Markdown("# vLLM-Omni Online Serving Demo")
-        gr.Markdown(f"**Model:** {model} \n\n")
-
-        with gr.Column():
-            with gr.Row():
-                input_box = gr.Textbox(
-                    label="Text Prompt",
-                    placeholder="For example: Describe what happens in the media inputs.",
-                    lines=4,
-                    scale=1,
-                )
-            with gr.Row(elem_classes="media-input-container"):
-                image_input = gr.Image(
-                    label="Image Input (optional)",
-                    type="pil",
-                    sources=["upload"],
-                    scale=1,
-                    elem_classes="image-input",
-                )
-                with gr.Column(scale=1, elem_classes="video-column"):
-                    video_input = gr.Video(
-                        label="Video Input (optional)",
-                        sources=["upload"],
-                        elem_classes="video-input",
-                    )
-                    use_audio_in_video_checkbox = gr.Checkbox(
-                        label="Use audio from video",
-                        value=False,
-                        info="Extract the video's audio track when provided.",
-                    )
-                audio_input = gr.Audio(
-                    label="Audio Input (optional)",
-                    type="numpy",
-                    sources=["upload", "microphone"],
-                    scale=1,
-                    elem_classes="audio-input",
-                )
-
-        with gr.Row():
-            generate_btn = gr.Button(
-                "Generate",
-                variant="primary",
-                size="lg",
-                elem_id="generate-btn",
-            )
-
-        with gr.Row():
-            text_output = gr.Textbox(label="Text Output", lines=10, scale=2)
-            audio_output = gr.Audio(label="Audio Output", interactive=False, scale=1)
-
-        generate_btn.click(
-            fn=run_inference,
-            inputs=[input_box, audio_input, image_input, video_input, use_audio_in_video_checkbox],
-            outputs=[text_output, audio_output],
-        )
-        demo.queue()
-    return demo
-
-
-def main():
-    args = parse_args()
-    omni_llm = None
-
-    model_name = "/".join(args.model.split("/")[-2:])
-    assert model_name in SUPPORTED_MODELS, (
-        f"Unsupported model '{model_name}'. Supported models: {SUPPORTED_MODELS.keys()}"
-    )
-
-    # Register signal handlers for graceful shutdown
-    def signal_handler(sig, frame):
-        print("\nReceived interrupt signal, shutting down...")
-        if omni_llm is not None:
-            try:
-                omni_llm.shutdown()
-            except Exception as e:
-                print(f"Error during shutdown: {e}")
-        sys.exit(0)
-
-    signal.signal(signal.SIGINT, signal_handler)
-    signal.signal(signal.SIGTERM, signal_handler)
-
-    print(f"Initializing AsyncOmniLLM with model: {args.model}")
-    if args.stage_configs_path:
-        print(f"Using custom stage configs: {args.stage_configs_path}")
-
-    sampling_params = build_sampling_params(SEED, model_name)
-    omni_llm = AsyncOmniLLM(
-        model=args.model,
-        stage_configs_path=args.stage_configs_path,
-        init_timeout=ASYNC_INIT_TIMEOUT,
-    )
-    print("✓ AsyncOmniLLM initialized successfully")
-    prompt_args_template = create_prompt_args(args)
-
-    demo = build_interface(
-        omni_llm,
-        sampling_params,
-        prompt_args_template,
-        args.model,
-    )
-    try:
-        demo.launch(
-            server_name=args.ip,
-            server_port=args.port,
-            share=args.share,
-        )
-    except KeyboardInterrupt:
-        print("\nShutting down...")
-    finally:
-        # Cleanup
-        if omni_llm is not None:
-            try:
-                omni_llm.shutdown()
-            except Exception as e:
-                print(f"Error during cleanup: {e}")
-
-
-if __name__ == "__main__":
-    main()

examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py

@@ -1,9 +1,6 @@
 import base64
 
-import requests
 from openai import OpenAI
-from vllm.assets.audio import AudioAsset
-from vllm.utils import FlexibleArgumentParser
 
 # Modify OpenAI's API key and API base to use vLLM's API server.
 openai_api_key = "EMPTY"
@@ -18,16 +15,6 @@ client = OpenAI(
 SEED = 42
 
 
-def encode_base64_content_from_url(content_url: str) -> str:
-    """Encode a content retrieved from a remote url to base64 format."""
-
-    with requests.get(content_url) as response:
-        response.raise_for_status()
-        result = base64.b64encode(response.content).decode("utf-8")
-
-    return result
-
-
 def get_system_prompt():
     return {
         "role": "system",
@@ -44,106 +31,7 @@ def get_system_prompt():
     }
 
 
-def get_text_query():
-    question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
-    prompt = {
-        "role": "user",
-        "content": [
-            {
-                "type": "text",
-                "text": f"{question}",
-            }
-        ],
-    }
-    return prompt
-
-
-def get_mixed_modalities_query():
-    question = "What is recited in the audio? What is the content of this image? Why is this video funny?"
-    prompt = {
-        "role": "user",
-        "content": [
-            {
-                "type": "audio_url",
-                "audio_url": {"url": AudioAsset("mary_had_lamb").url},
-            },
-            {
-                "type": "image_url",
-                "image_url": {
-                    "url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
-                },
-            },
-            {
-                "type": "video_url",
-                "video_url": {
-                    "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
-                },
-            },
-            {
-                "type": "text",
-                "text": f"{question}",
-            },
-        ],
-    }
-
-    return prompt
-
-
-def get_use_audio_in_video_query():
-    question = "Describe the content of the video, then convert what the baby say into text."
-
-    prompt = {
-        "role": "user",
-        "content": [
-            {
-                "type": "video_url",
-                "video_url": {
-                    "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4",
-                    "num_frames": 16,
-                },
-            },
-            {
-                "type": "text",
-                "text": f"{question}",
-            },
-        ],
-    }
-
-    return prompt
-
-
-def get_multi_audios_query():
-    question = "Are these two audio clips the same?"
-    prompt = {
-        "role": "user",
-        "content": [
-            {
-                "type": "audio_url",
-                "audio_url": {"url": AudioAsset("mary_had_lamb").url},
-            },
-            {
-                "type": "audio_url",
-                "audio_url": {"url": AudioAsset("winning_call").url},
-            },
-            {
-                "type": "text",
-                "text": f"{question}",
-            },
-        ],
-    }
-    return prompt
-
-
-query_map = {
-    "mixed_modalities": get_mixed_modalities_query,
-    "use_audio_in_video": get_use_audio_in_video_query,
-    "multi_audios": get_multi_audios_query,
-    "text": get_text_query,
-}
-
-
-def run_multimodal_generation(args) -> None:
-    model_name = "Qwen/Qwen2.5-Omni-7B"
+def run_text_to_audio(model: str) -> None:
     thinker_sampling_params = {
         "temperature": 0.0,  # Deterministic - no randomness
         "top_p": 1.0,  # Disable nucleus sampling
@@ -179,21 +67,23 @@ def run_multimodal_generation(args) -> None:
         code2wav_sampling_params,
     ]
 
-    prompt = query_map[args.query_type]()
-    extra_body = {
-        "sampling_params_list": sampling_params_list  # Optional, it has a default setting in stage_configs of the corresponding model.
-    }
-
-    if args.query_type == "use_audio_in_video":
-        extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}
-
     chat_completion = client.chat.completions.create(
         messages=[
             get_system_prompt(),
-            prompt,
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words.",
+                    },
+                ],
+            },
         ],
-        model=model_name,
-        extra_body=extra_body,
+        model=model,
+        extra_body={
+            "sampling_params_list": sampling_params_list
+        },  # Optional, it has a default setting in stage_configs of the corresponding model.
     )
 
     count = 0
@@ -209,20 +99,5 @@ def run_multimodal_generation(args) -> None:
             print("Chat completion output from text:", choice.message.content)
 
 
-def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
-    parser.add_argument(
-        "--query-type",
-        "-q",
-        type=str,
-        default="mixed_modalities",
-        choices=query_map.keys(),
-        help="Query type.",
-    )
-
-    return parser.parse_args()
-
-
 if __name__ == "__main__":
-    args = parse_args()
-    run_multimodal_generation(args)
+    run_text_to_audio("Qwen/Qwen2.5-Omni-7B")

examples/online_serving/run_curl_multimodal_generation.sh

@@ -1,20 +1,6 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
-# Default query type
-QUERY_TYPE="${1:-mixed_modalities}"
-
-# Validate query type
-if [[ ! "$QUERY_TYPE" =~ ^(mixed_modalities|use_audio_in_video|multi_audios|text)$ ]]; then
-    echo "Error: Invalid query type '$QUERY_TYPE'"
-    echo "Usage: $0 [mixed_modalities|use_audio_in_video|multi_audios|text]"
-    echo "  mixed_modalities: Audio + Image + Video + Text query"
-    echo "  use_audio_in_video: Video + Text query (with audio extraction from video)"
-    echo "  multi_audios: Two audio clips + Text query"
-    echo "  text: Text query"
-    exit 1
-fi
-
 SEED=42
 
 thinker_sampling_params='{
@@ -49,121 +35,18 @@ code2wav_sampling_params='{
 }'
 # Above is optional, it has a default setting in stage_configs of the corresponding model.
 
-# Define URLs for assets
-MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
-WINNING_CALL_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/winning_call.ogg"
-CHERRY_BLOSSOM_IMAGE_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
-SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
-
-# Build user content and extra fields based on query type
-case "$QUERY_TYPE" in
-  text)
-    user_content='[
-      {
-        "type": "text",
-        "text": "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
-      }
-    ]'
-    sampling_params_list='[
-      '"$thinker_sampling_params"',
-      '"$talker_sampling_params"',
-      '"$code2wav_sampling_params"'
-    ]'
-    mm_processor_kwargs="{}"
-    ;;
-  mixed_modalities)
-    user_content='[
-        {
-          "type": "audio_url",
-          "audio_url": {
-            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
-          }
-        },
-        {
-          "type": "image_url",
-          "image_url": {
-            "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
-          }
-        },
-        {
-          "type": "video_url",
-          "video_url": {
-            "url": "'"$SAMPLE_VIDEO_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "What is recited in the audio? What is the content of this image? Why is this video funny?"
-        }
-      ]'
-    sampling_params_list='[
-      '"$thinker_sampling_params"',
-      '"$talker_sampling_params"',
-      '"$code2wav_sampling_params"'
-    ]'
-    mm_processor_kwargs="{}"
-    ;;
-  use_audio_in_video)
-    user_content='[
-        {
-          "type": "video_url",
-          "video_url": {
-            "url": "'"$SAMPLE_VIDEO_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "Describe the content of the video, then convert what the baby say into text."
-        }
-      ]'
-    sampling_params_list='[
-      '"$thinker_sampling_params"',
-      '"$talker_sampling_params"',
-      '"$code2wav_sampling_params"'
-    ]'
-    mm_processor_kwargs='{
-      "use_audio_in_video": true
-    }'
-    ;;
-  multi_audios)
-    user_content='[
-        {
-          "type": "audio_url",
-          "audio_url": {
-            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
-          }
-        },
-        {
-          "type": "audio_url",
-          "audio_url": {
-            "url": "'"$WINNING_CALL_AUDIO_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "Are these two audio clips the same?"
-        }
-      ]'
-    sampling_params_list='[
-      '"$thinker_sampling_params"',
-      '"$talker_sampling_params"',
-      '"$code2wav_sampling_params"'
-    ]'
-    mm_processor_kwargs="{}"
-    ;;
-esac
-
-echo "Running query type: $QUERY_TYPE"
-echo ""
-
-
 output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d @- <<EOF
+  -H "Content-Type: application/json" \
+  -d @- <<EOF
 {
   "model": "Qwen/Qwen2.5-Omni-7B",
-  "sampling_params_list": $sampling_params_list,
-  "mm_processor_kwargs": $mm_processor_kwargs,
+  "extra_body": {
+    "sampling_params_list": [
+      $thinker_sampling_params,
+      $talker_sampling_params,
+      $code2wav_sampling_params
+    ]
+  },
   "messages": [
     {
       "role": "system",
@@ -176,12 +59,17 @@ output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
     },
     {
       "role": "user",
-      "content": $user_content
+      "content": [
+        {
+          "type": "text",
+          "text": "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+        }
+      ]
     }
   ]
 }
 EOF
-  )
+)
 
 # Here it only shows the text content of the first choice. Audio content has many binaries, so it's not displayed here.
 echo "Output of request: $(echo "$output" | jq '.choices[0].message.content')"

pyproject.toml

@@ -30,10 +30,10 @@ classifiers = [
 dependencies = [
     # Core runtime dependencies (required for actual usage)
     "omegaconf>=2.3.0",
-    # "decord>=0.6.0", # install manually
+    "decord>=0.6.0",
     "librosa>=0.11.0",
     "resampy>=0.4.3",
-    # "vllm==0.11.0",  # TODO: fix the entrypoints overwrite problem
+    # "vllm==0.11.0", # install manually
 ]
 
 [project.optional-dependencies]
@@ -46,8 +46,6 @@ dev = [
     "isort>=5.12.0",
     "mypy>=1.0.0",
     "pre-commit>=3.0.0",
-    "gradio>=5.49.1",
-    "soundfile>=0.13.1"
 ]
 
 docs = [
@@ -75,14 +73,10 @@ Documentation = "https://vllm-omni.readthedocs.io"
 vllm = "vllm_omni.entrypoints.cli.main:main"
 vllm-omni = "vllm_omni.entrypoints.cli.main:main"
 
-
 [tool.setuptools.packages.find]
 where = ["."]
 include = ["vllm_omni*"]
 
-[tool.setuptools.package-data]
-"vllm_omni.model_executor.stage_configs" = ["*.yaml"]
-
 [tool.ruff]
 line-length = 120
 exclude = [

scripts/README.md

@@ -0,0 +1,108 @@
+# vLLM-omni Testing Scripts
+
+This directory contains automated testing scripts for vLLM-omni serving functionality.
+
+## Scripts Overview
+
+### 1. `test_serving.sh` - Comprehensive Test Suite
+Full-featured testing script that performs complete validation of the vLLM-omni serving functionality.
+
+**Features:**
+- Model existence validation
+- Environment setup verification
+- Import functionality testing
+- Server startup and health checks
+- Text generation testing
+- Performance benchmarking
+- API client integration testing
+- Comprehensive logging and error handling
+
+**Usage:**
+```bash
+# Use default model and port
+./scripts/test_serving.sh
+
+# Use specific model and port
+./scripts/test_serving.sh ./models/Qwen3-0.6B 8000
+
+# Use HuggingFace model
+./scripts/test_serving.sh Qwen/Qwen3-0.6B 8001
+
+# Show help
+./scripts/test_serving.sh --help
+```
+
+**Test Coverage:**
+- ✅ Model loading and server startup
+- ✅ Health and info endpoints
+- ✅ Text generation functionality
+- ✅ Performance metrics
+- ✅ API client integration
+- ✅ All imports working correctly
+
+### 2. `quick_test.sh` - Fast Validation
+Lightweight script for quick validation after making changes.
+
+**Features:**
+- Fast import testing
+- Basic server startup
+- Health endpoint validation
+- Simple text generation test
+- Retry mechanism for reliability
+
+**Usage:**
+```bash
+# Use default port (8000)
+./scripts/quick_test.sh
+
+# Use specific port
+./scripts/quick_test.sh 8001
+```
+
+**Test Coverage:**
+- ✅ Import functionality
+- ✅ Server startup
+- ✅ Health endpoint
+- ✅ Basic text generation
+
+## Prerequisites
+
+1. **Conda Environment**: Ensure `vllm_omni` environment is activated
+2. **Model Available**: Qwen3-0.6B model should be available in `./models/Qwen3-0.6B/`
+3. **Dependencies**: All vLLM-omni dependencies should be installed
+
+## Environment Setup
+
+```bash
+# Activate conda environment
+conda activate vllm_omni
+
+# Verify installation
+python -c "import vllm_omni; print('vLLM-omni ready')"
+```
+
+## Model Setup
+
+The scripts expect the Qwen3-0.6B model to be available. You can:
+
+1. **Use local model**: Place model in `./models/Qwen3-0.6B/`
+2. **Use HuggingFace model**: Pass `Qwen/Qwen3-0.6B` as model path
+3. **Download model**: Use the `download_models.py` script
+
+```bash
+# Download model using the provided script
+python scripts/download_models.py
+```
+
+## Output and Logging
+
+### Test Results
+- **Success**: Green `[SUCCESS]` messages
+- **Info**: Blue `[INFO]` messages  
+- **Warnings**: Yellow `[WARNING]` messages
+- **Errors**: Red `[ERROR]` messages
+
+### Log Files
+- `server.log`: Server startup and runtime logs
+- `api_client_test.log`: API client test results
+- `simple_usage_test.log`: Simple usage test results

scripts/build_wheel.sh

@@ -1,139 +0,0 @@
-#!/usr/bin/env bash
-
-set -euo pipefail
-
-SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
-REPO_ROOT="$(cd -- "${SCRIPT_DIR}/.." && pwd)"
-PROJECT_NAME="vllm-omni"
-RUN_QUALITY=false
-SKIP_CLEAN=false
-CREATE_VENV=false
-VENV_DIR=".venv-build"
-PYTHON_BIN="python"
-UV_BIN="uv"
-
-log() {
-  local level="$1"
-  shift
-  printf '[%s] %s\n' "${level}" "$*"
-}
-
-abort() {
-  log "ERROR" "$*"
-  exit 1
-}
-
-usage() {
-  cat <<EOF
-Usage: $(basename "$0") [options]
-
-Options:
-  --run-quality    Run pre-commit, install dev deps, and pytest before building
-  --skip-clean     Skip removing previous build artifacts
-  --create-venv    Build inside a fresh virtual environment (default path: .venv-build)
-  --venv-dir PATH  Custom directory for the virtual environment (implies --create-venv)
-  --python PATH    Python executable to use (default: python)
-  -h, --help       Show this help message
-EOF
-}
-
-while [[ $# -gt 0 ]]; do
-  case "$1" in
-    --run-quality)
-      RUN_QUALITY=true
-      ;;
-    --skip-clean)
-      SKIP_CLEAN=true
-      ;;
-    --create-venv)
-      CREATE_VENV=true
-      ;;
-    --venv-dir)
-      CREATE_VENV=true
-      shift
-      [[ $# -gt 0 ]] || abort "--venv-dir requires a path"
-      VENV_DIR="$1"
-      ;;
-    --python)
-      shift
-      [[ $# -gt 0 ]] || abort "--python requires a path"
-      PYTHON_BIN="$1"
-      ;;
-    -h|--help)
-      usage
-      exit 0
-      ;;
-    *)
-      usage
-      abort "Unknown option: $1"
-      ;;
-  esac
-  shift
-done
-
-HOST_PYTHON="${PYTHON_BIN}"
-
-log "INFO" "Switching to repository root: ${REPO_ROOT}"
-cd "${REPO_ROOT}" || abort "Cannot enter repository root"
-
-[[ -f pyproject.toml ]] || abort "pyproject.toml not found, please ensure correct script location"
-
-ensure_uv() {
-  if ! command -v "${UV_BIN}" >/dev/null 2>&1; then
-    log "INFO" "uv not found, installing via ${HOST_PYTHON}"
-    "${HOST_PYTHON}" -m pip install --upgrade pip
-    "${HOST_PYTHON}" -m pip install uv
-  fi
-}
-
-ensure_uv
-
-if [[ "${CREATE_VENV}" == "true" ]]; then
-  log "INFO" "Creating fresh virtual environment at ${VENV_DIR} via uv"
-  "${UV_BIN}" venv --python "${HOST_PYTHON}" --seed "${VENV_DIR}"
-  PYTHON_BIN="${VENV_DIR}/bin/python"
-  [[ -x "${PYTHON_BIN}" ]] || abort "Failed to locate python inside ${VENV_DIR}"
-  log "INFO" "Installing build module inside virtual environment"
-  "${UV_BIN}" pip install --python "${PYTHON_BIN}" build
-else
-  log "INFO" "Ensuring build module is available via uv pip"
-  "${UV_BIN}" pip install --python "${PYTHON_BIN}" build
-fi
-
-log "INFO" "Checking build module"
-if ! "${PYTHON_BIN}" -m build --version >/dev/null 2>&1; then
-  abort "${PYTHON_BIN} -m build is not available, install build first"
-fi
-
-run_quality_steps() {
-  log "INFO" "Running quality checks"
-  "${UV_BIN}" pip install --python "${PYTHON_BIN}" -e ".[dev]"
-  "${PYTHON_BIN}" -m pre_commit run --all-files
-  "${PYTHON_BIN}" -m pytest tests/ -v -m "not slow"
-}
-
-if [[ "${RUN_QUALITY}" == "true" ]]; then
-  run_quality_steps
-else
-  log "INFO" "Quality steps available via --run-quality"
-  log "INFO" "  - pre-commit run --all-files"
-  log "INFO" "  - pip install -e '.[dev]'"
-  log "INFO" "  - pytest tests/ -v -m \"not slow\""
-fi
-
-cleanup_artifacts() {
-  log "INFO" "Cleaning previous build artifacts"
-  rm -rf build dist "${PROJECT_NAME}.egg-info" "${PROJECT_NAME//-/_}.egg-info"
-}
-
-if [[ "${SKIP_CLEAN}" == "true" ]]; then
-  log "INFO" "Skipping cleanup as requested"
-else
-  cleanup_artifacts
-fi
-
-log "INFO" "Building source and wheel distributions"
-"${PYTHON_BIN}" -m build
-
-log "INFO" "Build finished, artifacts:"
-ls -lh dist

scripts/test_serving.sh

@@ -0,0 +1,363 @@
+#!/bin/bash
+
+# vLLM-omni Serving Functionality Test Script
+# This script tests the complete serving functionality of vLLM-omni
+# Usage: ./scripts/test_serving.sh [model_path] [port]
+
+set -e  # Exit on any error
+
+# Configuration
+MODEL_PATH=${1:-"./models/Qwen3-0.6B"}
+PORT=${2:-8000}
+HOST="localhost"
+TIMEOUT=30
+SERVER_STARTUP_WAIT=15
+
+# Colors for output
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+BLUE='\033[0;34m'
+NC='\033[0m' # No Color
+
+# Logging functions
+log_info() {
+    echo -e "${BLUE}[INFO]${NC} $1"
+}
+
+log_success() {
+    echo -e "${GREEN}[SUCCESS]${NC} $1"
+}
+
+log_warning() {
+    echo -e "${YELLOW}[WARNING]${NC} $1"
+}
+
+log_error() {
+    echo -e "${RED}[ERROR]${NC} $1"
+}
+
+# Cleanup function
+cleanup() {
+    log_info "Cleaning up..."
+    if [ ! -z "$SERVER_PID" ]; then
+        log_info "Stopping server (PID: $SERVER_PID)..."
+        kill $SERVER_PID 2>/dev/null || true
+        wait $SERVER_PID 2>/dev/null || true
+    fi
+    # Kill any remaining vllm processes
+    pkill -f "vllm serve" 2>/dev/null || true
+    log_info "Cleanup completed"
+}
+
+# Set up trap for cleanup on exit
+trap cleanup EXIT
+
+# Check if model exists
+check_model() {
+    log_info "Checking if model exists at: $MODEL_PATH"
+    if [ ! -d "$MODEL_PATH" ]; then
+        log_error "Model directory not found: $MODEL_PATH"
+        log_info "Available models:"
+        ls -la models/ 2>/dev/null || log_warning "No models directory found"
+        exit 1
+    fi
+    log_success "Model found: $MODEL_PATH"
+}
+
+# Check if conda environment is activated
+check_environment() {
+    log_info "Checking conda environment..."
+    if [[ "$CONDA_DEFAULT_ENV" != "vllm_omni" ]]; then
+        log_warning "vllm_omni conda environment not activated"
+        log_info "Attempting to activate vllm_omni environment..."
+        source $(conda info --base)/etc/profile.d/conda.sh
+        conda activate vllm_omni
+        if [[ "$CONDA_DEFAULT_ENV" != "vllm_omni" ]]; then
+            log_error "Failed to activate vllm_omni environment"
+            exit 1
+        fi
+    fi
+    log_success "Environment check passed: $CONDA_DEFAULT_ENV"
+}
+
+# Test import functionality
+test_imports() {
+    log_info "Testing imports..."
+    python -c "
+import vllm_omni
+from vllm_omni.entrypoints.omni_llm import OmniLLM, AsyncOmniLLM
+print('✅ All imports successful')
+" || {
+        log_error "Import test failed"
+        exit 1
+    }
+    log_success "Import test passed"
+}
+
+# Start the server
+start_server() {
+    log_info "Starting vLLM-omni server on port $PORT..."
+    log_info "Command: vllm serve $MODEL_PATH --omni --port $PORT"
+
+    # Start server in background
+    vllm serve "$MODEL_PATH" --omni --port "$PORT" > server.log 2>&1 &
+    SERVER_PID=$!
+
+    log_info "Server started with PID: $SERVER_PID"
+    log_info "Waiting $SERVER_STARTUP_WAIT seconds for server to initialize..."
+    sleep $SERVER_STARTUP_WAIT
+
+    # Check if server is still running
+    if ! kill -0 $SERVER_PID 2>/dev/null; then
+        log_error "Server failed to start"
+        log_info "Server logs:"
+        cat server.log
+        exit 1
+    fi
+
+    log_success "Server appears to be running"
+}
+
+# Test health endpoint
+test_health() {
+    log_info "Testing health endpoint..."
+    for attempt in {1..10}; do
+        local response=$(curl -s -w "%{http_code}" http://$HOST:$PORT/health 2>/dev/null || echo "000")
+        local http_code="${response: -3}"
+        local body="${response%???}"
+
+        if [ "$http_code" = "200" ]; then
+            log_success "Health check passed: $body"
+            return 0
+        fi
+
+        log_warning "Health check attempt $attempt failed (HTTP $http_code). Retrying in 3s..."
+        sleep 3
+    done
+
+    log_error "Health check failed after multiple attempts"
+    return 1
+}
+
+# Test info endpoint
+test_info() {
+    log_info "Testing info endpoint..."
+    local response=$(curl -s -w "%{http_code}" http://$HOST:$PORT/info 2>/dev/null || echo "000")
+    local http_code="${response: -3}"
+    local body="${response%???}"
+
+    if [ "$http_code" = "200" ]; then
+        log_success "Info endpoint working"
+        echo "$body" | python -m json.tool 2>/dev/null || log_warning "Info response not valid JSON"
+    else
+        log_error "Info endpoint failed (HTTP $http_code): $body"
+        return 1
+    fi
+}
+
+# Test text generation
+test_generation() {
+    log_info "Testing text generation..."
+
+    # Test 1: Simple generation
+    local response=$(curl -s -X POST http://$HOST:$PORT/generate \
+        -H "Content-Type: application/json" \
+        -d '{"prompts": ["Test the server functionality"], "max_tokens": 20, "temperature": 0.7}' \
+        2>/dev/null || echo "{}")
+
+    if echo "$response" | python -c "import sys, json; data=json.load(sys.stdin); exit(0 if 'outputs' in data and len(data['outputs']) > 0 else 1)" 2>/dev/null; then
+        log_success "Text generation test passed"
+        echo "$response" | python -c "import sys, json; data=json.load(sys.stdin); print('Generated text:', data['outputs'][0]['text'][:100] + '...' if len(data['outputs'][0]['text']) > 100 else data['outputs'][0]['text'])" 2>/dev/null
+    else
+        log_error "Text generation test failed"
+        echo "Response: $response"
+        return 1
+    fi
+}
+
+# Test API client example
+test_api_client() {
+    log_info "Testing API client example..."
+    if [ -f "examples/basic/api_client.py" ]; then
+        python examples/basic/api_client.py > api_client_test.log 2>&1 || {
+            log_error "API client test failed"
+            log_info "API client logs:"
+            cat api_client_test.log
+            return 1
+        }
+        log_success "API client test passed"
+    else
+        log_warning "API client example not found, skipping"
+    fi
+}
+
+# Test AR → DiT pipeline example
+test_ar_dit_pipeline() {
+    log_info "Testing AR → DiT diffusers pipeline example..."
+
+    if ! python - 2>/dev/null <<'PY'
+import importlib
+for pkg in ("torch", "diffusers"):
+    if importlib.util.find_spec(pkg) is None:
+        raise ImportError(pkg)
+PY
+    then
+        log_warning "Required packages for diffusers pipeline not found, skipping AR → DiT test"
+        return 0
+    fi
+
+    if [ ! -f "examples/omni/ar_dit_diffusers.py" ]; then
+        log_warning "AR → DiT example not found, skipping"
+        return 0
+    fi
+
+    local output_path="logs/ar_dit_pipeline.png"
+    local log_path="logs/ar_dit_pipeline.log"
+    rm -f "$output_path"
+
+    local cmd=(python examples/omni/ar_dit_diffusers.py
+        --prompt "Smoke test prompt of a lighthouse"
+        --seed 0
+        --output "$output_path")
+
+    local env_prefix=()
+    if [ -n "$AR_MODEL" ]; then
+        env_prefix+=("AR_MODEL=$AR_MODEL")
+    fi
+    if [ -n "$DIT_MODEL" ]; then
+        env_prefix+=("DIT_MODEL=$DIT_MODEL")
+    fi
+
+    if ! env "${env_prefix[@]}" "${cmd[@]}" >"$log_path" 2>&1; then
+        log_error "AR → DiT pipeline test failed"
+        log_info "AR → DiT logs:"
+        cat "$log_path"
+        return 1
+    fi
+
+    if [ ! -f "$output_path" ]; then
+        log_error "AR → DiT pipeline did not produce an image"
+        log_info "AR → DiT logs:"
+        cat "$log_path"
+        return 1
+    fi
+
+    log_success "AR → DiT pipeline test passed (output: $output_path)"
+}
+
+# Test simple usage example
+test_simple_usage() {
+    log_info "Testing simple usage example..."
+    if [ -f "examples/basic/simple_usage.py" ]; then
+        # Run with timeout to avoid hanging
+        timeout 60 python examples/basic/simple_usage.py > simple_usage_test.log 2>&1 || {
+            local exit_code=$?
+            if [ $exit_code -eq 124 ]; then
+                log_warning "Simple usage test timed out (60s), but this might be normal for model loading"
+            else
+                log_error "Simple usage test failed (exit code: $exit_code)"
+                log_info "Simple usage logs:"
+                cat simple_usage_test.log
+                return 1
+            fi
+        }
+        log_success "Simple usage test passed"
+    else
+        log_warning "Simple usage example not found, skipping"
+    fi
+}
+
+# Performance test
+test_performance() {
+    log_info "Running performance test..."
+
+    local start_time=$(date +%s)
+    local response=$(curl -s -X POST http://$HOST:$PORT/generate \
+        -H "Content-Type: application/json" \
+        -d '{"prompts": ["Performance test prompt"], "max_tokens": 50, "temperature": 0.7}' \
+        2>/dev/null)
+    local end_time=$(date +%s)
+    local duration=$((end_time - start_time))
+
+    if echo "$response" | python -c "import sys, json; data=json.load(sys.stdin); exit(0 if 'outputs' in data else 1)" 2>/dev/null; then
+        log_success "Performance test passed (${duration}s response time)"
+    else
+        log_error "Performance test failed"
+        return 1
+    fi
+}
+
+# Main test function
+run_tests() {
+    log_info "Starting vLLM-omni serving functionality tests..."
+    log_info "Model: $MODEL_PATH"
+    log_info "Port: $PORT"
+    log_info "Host: $HOST"
+    echo "=========================================="
+
+    # Run all tests
+    check_model
+    check_environment
+    test_imports
+    start_server
+
+    # Wait a bit more for server to be fully ready
+    sleep 5
+
+    test_health
+    test_info
+    test_generation
+    test_performance
+    test_api_client
+    test_ar_dit_pipeline
+
+    # Note: Simple usage test is commented out as it starts its own server
+    # test_simple_usage
+
+    log_success "All tests completed successfully!"
+    echo "=========================================="
+    log_info "Test Summary:"
+    log_info "✅ Model loading and server startup"
+    log_info "✅ Health and info endpoints"
+    log_info "✅ Text generation functionality"
+    log_info "✅ Performance metrics"
+    log_info "✅ API client integration"
+    log_info "✅ AR → DiT diffusers pipeline"
+    log_info "✅ All imports working correctly"
+}
+
+# Show usage
+show_usage() {
+    echo "Usage: $0 [model_path] [port]"
+    echo ""
+    echo "Arguments:"
+    echo "  model_path    Path to the model directory (default: ./models/Qwen3-0.6B)"
+    echo "  port          Port to run the server on (default: 8000)"
+    echo ""
+    echo "Examples:"
+    echo "  $0                                    # Use default model and port"
+    echo "  $0 ./models/Qwen3-0.6B 8001          # Use specific model and port"
+    echo "  $0 Qwen/Qwen3-0.6B 8000              # Use HuggingFace model"
+}
+
+# Main execution
+main() {
+    # Check for help flag
+    if [[ "$1" == "-h" || "$1" == "--help" ]]; then
+        show_usage
+        exit 0
+    fi
+
+    # Change to script directory
+    cd "$(dirname "$0")/.."
+
+    # Create logs directory if it doesn't exist
+    mkdir -p logs
+
+    # Run tests
+    run_tests
+}
+
+# Run main function with all arguments
+main "$@"

vllm_omni/config/model.py

@@ -15,6 +15,7 @@ class OmniModelConfig(ModelConfig):
     stage_id: int = 0
     model_stage: str = "thinker"
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
+    ode_solver_class: Optional[str] = None
     engine_output_type: Optional[str] = None
 
     @property

vllm_omni/engine/arg_utils.py

@@ -12,6 +12,7 @@ class OmniEngineArgs(EngineArgs):
     stage_id: int = 0
     model_stage: str = "thinker"
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
+    ode_solver_class: Optional[str] = None
     engine_output_type: Optional[str] = None
 
     def create_model_config(self) -> OmniModelConfig:
@@ -27,6 +28,7 @@ class OmniEngineArgs(EngineArgs):
         config_dict["model_stage"] = self.model_stage
         config_dict["model_arch"] = self.model_arch
         config_dict["engine_output_type"] = self.engine_output_type
+        config_dict["ode_solver_class"] = self.ode_solver_class
 
         # Create and return the OmniModelConfig instance
         omni_config = OmniModelConfig(**config_dict)

vllm_omni/entrypoints/chat_utils.py

@@ -1,242 +0,0 @@
-from collections.abc import Awaitable, Iterable
-from typing import Any, Optional, Union, cast
-
-import numpy as np
-from openai.types.chat import ChatCompletionContentPartTextParam
-from vllm.config import ModelConfig
-from vllm.entrypoints.chat_utils import (
-    AsyncMultiModalContentParser,
-    AsyncMultiModalItemTracker,
-    BaseMultiModalContentParser,
-    BaseMultiModalItemTracker,
-    ChatCompletionContentPartParam,
-    ChatCompletionMessageParam,
-    ConversationMessage,
-    MultiModalDataDict,
-    MultiModalUUIDDict,
-    _AssistantParser,
-    _ChatTemplateContentFormat,
-    _ContentPart,
-    _get_full_multimodal_text_prompt,
-    _parse_chat_message_content_part,
-    _postprocess_messages,
-    _ToolParser,
-)
-from vllm.transformers_utils.tokenizer import AnyTokenizer
-
-
-class OmniAsyncMultiModalItemTracker(AsyncMultiModalItemTracker):
-    def create_parser(self) -> "BaseMultiModalContentParser":
-        return OmniAsyncMultiModalContentParser(self)
-
-
-class OmniAsyncMultiModalContentParser(AsyncMultiModalContentParser):
-    def __init__(self, tracker: AsyncMultiModalItemTracker) -> None:
-        super().__init__(tracker=tracker)
-        self._mm_processor_kwargs: Optional[dict[str, Any]] = None
-
-    def set_mm_processor_kwargs(self, mm_processor_kwargs: Optional[dict[str, Any]]) -> None:
-        """Set mm_processor_kwargs for use in parsing."""
-        self._mm_processor_kwargs = mm_processor_kwargs
-
-    def parse_video(self, video_url: Optional[str], uuid: Optional[str] = None) -> None:
-        video = self._connector.fetch_video_async(video_url=video_url) if video_url else None
-
-        placeholder = self._tracker.add("video", video, uuid)
-        self._add_placeholder("video", placeholder)
-
-        # Extract audio from video if use_audio_in_video is True
-        if video_url and self._mm_processor_kwargs and self._mm_processor_kwargs.get("use_audio_in_video", False):
-            audio_coro = self._extract_audio_from_video_async(video_url)
-            audio_placeholder = self._tracker.add("audio", audio_coro, uuid)
-            self._add_placeholder("audio", audio_placeholder)
-
-    async def _extract_audio_from_video_async(self, video_url: str) -> tuple[np.ndarray, Union[int, float]]:
-        """
-        Extract audio from video URL using librosa.
-        Returns tuple of (audio_array, sample_rate) compatible with audio format.
-
-        All blocking I/O operations are run in a thread pool to avoid blocking the event loop.
-        """
-        import asyncio
-        import os
-        import tempfile
-        from urllib.parse import urlparse
-
-        # Parse URL to determine type
-        parsed_url = urlparse(video_url)
-        temp_video_file_path = None
-
-        def _download_video_sync(url: str) -> bytes:
-            """Synchronous video download - runs in thread pool."""
-            from urllib.request import urlopen
-
-            return urlopen(url).read()
-
-        def _write_temp_file_sync(data: bytes, suffix: str) -> str:
-            """Synchronous temp file write - runs in thread pool."""
-            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
-                temp_file.write(data)
-                return temp_file.name
-
-        def _load_audio_sync(file_path: str) -> tuple[np.ndarray, Union[int, float]]:
-            """Synchronous audio loading with librosa - runs in thread pool."""
-            import librosa
-
-            return librosa.load(file_path, sr=16000)
-
-        def _cleanup_file_sync(file_path: str) -> None:
-            """Synchronous file deletion - runs in thread pool."""
-            try:
-                if os.path.exists(file_path):
-                    os.unlink(file_path)
-            except OSError:
-                pass
-
-        try:
-            if parsed_url.scheme in ("http", "https"):
-                # Download video from HTTP/HTTPS URL asynchronously
-                video_data = await asyncio.to_thread(_download_video_sync, video_url)
-                # Write temp file asynchronously
-                temp_video_file_path = await asyncio.to_thread(_write_temp_file_sync, video_data, ".mp4")
-            elif parsed_url.scheme == "file":
-                # Use file path directly (handle Windows paths)
-                from urllib.request import url2pathname
-
-                temp_video_file_path = url2pathname(parsed_url.path)
-            elif parsed_url.scheme == "data":
-                # Handle data URL (base64 encoded video)
-                import base64
-
-                header, data = video_url.split(",", 1)
-                video_data = base64.b64decode(data)
-                # Write temp file asynchronously
-                temp_video_file_path = await asyncio.to_thread(_write_temp_file_sync, video_data, ".mp4")
-            else:
-                # Assume it's a local file path
-                temp_video_file_path = video_url
-
-            # Extract audio using librosa asynchronously (CPU-intensive, runs in thread pool)
-            audio_array, sample_rate = await asyncio.to_thread(_load_audio_sync, temp_video_file_path)
-
-            return audio_array, sample_rate
-        finally:
-            # Clean up temporary file if we created one (asynchronously)
-            if temp_video_file_path and parsed_url.scheme in ("http", "https", "data"):
-                await asyncio.to_thread(_cleanup_file_sync, temp_video_file_path)
-
-
-def parse_chat_messages_futures(
-    messages: list[ChatCompletionMessageParam],
-    model_config: ModelConfig,
-    tokenizer: AnyTokenizer,
-    content_format: _ChatTemplateContentFormat,
-    mm_processor_kwargs: Optional[dict[str, Any]] = None,
-) -> tuple[
-    list[ConversationMessage],
-    Awaitable[Optional[MultiModalDataDict]],
-    Optional[MultiModalUUIDDict],
-]:
-    conversation: list[ConversationMessage] = []
-    mm_tracker = OmniAsyncMultiModalItemTracker(model_config, tokenizer)
-
-    for msg in messages:
-        sub_messages = _parse_chat_message_content(
-            msg,
-            mm_tracker,
-            content_format,
-            interleave_strings=(
-                content_format == "string"
-                and model_config.multimodal_config is not None
-                and model_config.multimodal_config.interleave_mm_strings
-            ),
-            mm_processor_kwargs=mm_processor_kwargs,
-        )
-
-        conversation.extend(sub_messages)
-
-    _postprocess_messages(conversation)
-
-    return conversation, mm_tracker.all_mm_data(), mm_tracker.all_mm_uuids()
-
-
-def _parse_chat_message_content(
-    message: ChatCompletionMessageParam,
-    mm_tracker: BaseMultiModalItemTracker,
-    content_format: _ChatTemplateContentFormat,
-    interleave_strings: bool,
-    mm_processor_kwargs: Optional[dict[str, Any]] = None,
-) -> list[ConversationMessage]:
-    role = message["role"]
-    content = message.get("content")
-
-    if content is None:
-        content = []
-    elif isinstance(content, str):
-        content = [ChatCompletionContentPartTextParam(type="text", text=content)]
-    result = _parse_chat_message_content_parts(
-        role,
-        content,  # type: ignore
-        mm_tracker,
-        wrap_dicts=(content_format == "openai"),
-        interleave_strings=interleave_strings,
-        mm_processor_kwargs=mm_processor_kwargs,
-    )
-
-    for result_msg in result:
-        if role == "assistant":
-            parsed_msg = _AssistantParser(message)
-
-            # The 'tool_calls' is not None check ensures compatibility.
-            # It's needed only if downstream code doesn't strictly
-            # follow the OpenAI spec.
-            if "tool_calls" in parsed_msg and parsed_msg["tool_calls"] is not None:
-                result_msg["tool_calls"] = list(parsed_msg["tool_calls"])
-        elif role == "tool":
-            parsed_msg = _ToolParser(message)
-            if "tool_call_id" in parsed_msg:
-                result_msg["tool_call_id"] = parsed_msg["tool_call_id"]
-
-        if "name" in message and isinstance(message["name"], str):
-            result_msg["name"] = message["name"]
-
-    return result
-
-
-def _parse_chat_message_content_parts(
-    role: str,
-    parts: Iterable[ChatCompletionContentPartParam],
-    mm_tracker: BaseMultiModalItemTracker,
-    *,
-    wrap_dicts: bool,
-    interleave_strings: bool,
-    mm_processor_kwargs: Optional[dict[str, Any]] = None,
-) -> list[ConversationMessage]:
-    content = list[_ContentPart]()
-
-    mm_parser = mm_tracker.create_parser()
-    # Set mm_processor_kwargs if parser supports it
-    if hasattr(mm_parser, "set_mm_processor_kwargs"):
-        mm_parser.set_mm_processor_kwargs(mm_processor_kwargs)
-
-    for part in parts:
-        parse_res = _parse_chat_message_content_part(
-            part,
-            mm_parser,
-            wrap_dicts=wrap_dicts,
-            interleave_strings=interleave_strings,
-        )
-        if parse_res:
-            content.append(parse_res)
-
-    if wrap_dicts:
-        # Parsing wraps images and texts as interleaved dictionaries
-        return [ConversationMessage(role=role, content=content)]  # type: ignore
-    texts = cast(list[str], content)
-    mm_placeholder_storage = mm_parser.mm_placeholder_storage()
-    if mm_placeholder_storage:
-        text_prompt = _get_full_multimodal_text_prompt(mm_placeholder_storage, texts, interleave_strings)
-    else:
-        text_prompt = "\n".join(texts)
-
-    return [ConversationMessage(role=role, content=text_prompt)]

vllm_omni/entrypoints/omni_stage.py

@@ -50,7 +50,6 @@ class OmniStage:
         self.stage_id = stage_config.stage_id
         self.engine_args = stage_config.engine_args
         self.model_stage = stage_config.engine_args.model_stage
-        self.requires_multimodal_data = getattr(stage_config.runtime, "requires_multimodal_data", False)
         self.engine_input_source = getattr(stage_config, "engine_input_source", [])
         self.engine_output_type = stage_config.engine_args.engine_output_type
         self.engine_outputs = None
@@ -206,20 +205,14 @@ class OmniStage:
             for source_output in source_outputs:
                 engine_input = OmniTokensPrompt(
                     prompt_token_ids=source_output.outputs[0].token_ids,
-                    multi_modal_data=(
-                        multi_modal_data[source_output.request_id]
-                        if self.requires_multimodal_data and multi_modal_data
-                        else None
-                    ),
+                    multi_modal_data=(multi_modal_data[source_output.request_id] if multi_modal_data else None),
                 )
                 engine_inputs.append(engine_input)
             return engine_inputs
 
         else:
             engine_input_source = self.engine_input_source
-            return self.custom_process_input_func(
-                stage_list, engine_input_source, prompt, self.requires_multimodal_data
-            )
+            return self.custom_process_input_func(stage_list, engine_input_source, prompt)
 
 
 def _stage_worker(
@@ -234,13 +227,13 @@ def _stage_worker(
     import logging as _logging
     import time as _time
 
-    from vllm_omni.entrypoints.log_utils import (
+    from vllm_omni.entrypoints.log_utils import (  # noqa: WPS433
         compute_and_log_stage_request_stats,
         count_tokens_from_outputs,
         log_stage_batch_stats,
         log_stage_running_avg,
     )
-    from vllm_omni.entrypoints.omni_llm import OmniStageLLM
+    from vllm_omni.entrypoints.omni_llm import OmniStageLLM  # noqa: WPS433
 
     # no inline JSONL/serialization imports; logging handled by utilities
 
@@ -508,8 +501,8 @@ async def _stage_worker_async(
     import logging as _logging
     import time as _time
 
-    from vllm_omni.entrypoints.async_omni_llm import AsyncOmniStageLLM
-    from vllm_omni.entrypoints.log_utils import (
+    from vllm_omni.entrypoints.async_omni_llm import AsyncOmniStageLLM  # noqa: WPS433
+    from vllm_omni.entrypoints.log_utils import (  # noqa: WPS433
         compute_and_log_stage_request_stats,
         count_tokens_from_outputs,
         log_stage_batch_stats,

vllm_omni/entrypoints/openai/serving_chat.py

@@ -3,10 +3,10 @@ import base64
 import json
 import time
 import uuid
-from collections.abc import AsyncGenerator, AsyncIterator, Sequence
+from collections.abc import AsyncGenerator, AsyncIterator
 from datetime import datetime, timedelta, timezone
 from io import BytesIO
-from typing import Any, Callable, Optional, Union
+from typing import Optional, Union
 
 import jinja2
 from fastapi import Request
@@ -18,16 +18,7 @@ except ImportError:
     soundfile = None
 
 from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio
-from vllm.entrypoints.chat_utils import (
-    ChatCompletionMessageParam,
-    ChatTemplateContentFormatOption,
-    ConversationMessage,
-    apply_hf_chat_template,
-    apply_mistral_chat_template,
-    get_history_tool_calls_cnt,
-    make_tool_call_id,
-    resolve_chat_template_content_format,
-)
+from vllm.entrypoints.chat_utils import ConversationMessage, get_history_tool_calls_cnt, make_tool_call_id
 from vllm.entrypoints.harmony_utils import parse_chat_output
 from vllm.entrypoints.openai.protocol import (
     ChatCompletionNamedToolChoiceParam,
@@ -44,16 +35,7 @@ from vllm.entrypoints.openai.protocol import (
     UsageInfo,
 )
 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
-from vllm.entrypoints.openai.serving_engine import (
-    ChatLikeRequest,
-    EngineTokensPrompt,
-    RequestPrompt,
-    ResponsesRequest,
-    TextTokensPrompt,
-    clamp_prompt_logprobs,
-    is_list_of,
-)
-from vllm.entrypoints.openai.tool_parsers import ToolParser
+from vllm.entrypoints.openai.serving_engine import RequestPrompt, clamp_prompt_logprobs
 from vllm.entrypoints.openai.tool_parsers.mistral_tool_parser import MistralToolCall
 from vllm.inputs.data import PromptType
 from vllm.logger import init_logger
@@ -68,7 +50,6 @@ from vllm.transformers_utils.tokenizers import (
 )
 from vllm.utils import as_list
 
-from vllm_omni.entrypoints.chat_utils import parse_chat_messages_futures
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -234,123 +215,6 @@ class OmniOpenAIServingChat(OpenAIServingChat):
             # TODO: Use a vllm-specific Validation Error
             return self.create_error_response(str(e))
 
-    async def _preprocess_chat(
-        self,
-        request: Union[ChatLikeRequest, ResponsesRequest],
-        tokenizer: AnyTokenizer,
-        messages: list[ChatCompletionMessageParam],
-        chat_template: Optional[str],
-        chat_template_content_format: ChatTemplateContentFormatOption,
-        add_generation_prompt: bool = True,
-        continue_final_message: bool = False,
-        tool_dicts: Optional[list[dict[str, Any]]] = None,
-        documents: Optional[list[dict[str, str]]] = None,
-        chat_template_kwargs: Optional[dict[str, Any]] = None,
-        tool_parser: Optional[Callable[[AnyTokenizer], ToolParser]] = None,
-        add_special_tokens: bool = False,
-    ) -> tuple[
-        list[ConversationMessage],
-        Sequence[RequestPrompt],
-        list[EngineTokensPrompt],
-    ]:
-        model_config = self.model_config
-
-        resolved_content_format = resolve_chat_template_content_format(
-            chat_template,
-            tool_dicts,
-            chat_template_content_format,
-            tokenizer,
-            model_config=model_config,
-        )
-        conversation, mm_data_future, mm_uuids = parse_chat_messages_futures(
-            messages,
-            model_config,
-            tokenizer,
-            content_format=resolved_content_format,
-            mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None),
-        )
-
-        _chat_template_kwargs: dict[str, Any] = dict(
-            chat_template=chat_template,
-            add_generation_prompt=add_generation_prompt,
-            continue_final_message=continue_final_message,
-            tools=tool_dicts,
-            documents=documents,
-        )
-        _chat_template_kwargs.update(chat_template_kwargs or {})
-
-        request_prompt: Union[str, list[int]]
-
-        if tokenizer is None:
-            request_prompt = "placeholder"
-        elif isinstance(tokenizer, MistralTokenizer):
-            request_prompt = apply_mistral_chat_template(
-                tokenizer,
-                messages=messages,
-                **_chat_template_kwargs,
-            )
-        else:
-            request_prompt = apply_hf_chat_template(
-                tokenizer=tokenizer,
-                conversation=conversation,
-                model_config=model_config,
-                **_chat_template_kwargs,
-            )
-
-        mm_data = await mm_data_future
-
-        # tool parsing is done only if a tool_parser has been set and if
-        # tool_choice is not "none" (if tool_choice is "none" but a tool_parser
-        # is set, we want to prevent parsing a tool_call hallucinated by the LLM
-        should_parse_tools = tool_parser is not None and (
-            hasattr(request, "tool_choice") and request.tool_choice != "none"
-        )
-
-        if should_parse_tools:
-            if not isinstance(request, ChatCompletionRequest):
-                msg = "Tool usage is only supported for Chat Completions API"
-                raise NotImplementedError(msg)
-
-            request = tool_parser(tokenizer).adjust_request(  # type: ignore
-                request=request
-            )
-
-        if tokenizer is None:
-            assert isinstance(request_prompt, str), (
-                "Prompt has to be a string",
-                "when the tokenizer is not initialised",
-            )
-            prompt_inputs = TextTokensPrompt(prompt=request_prompt, prompt_token_ids=[1])
-        elif isinstance(request_prompt, str):
-            prompt_inputs = await self._tokenize_prompt_input_async(
-                request,
-                tokenizer,
-                request_prompt,
-                add_special_tokens=add_special_tokens,
-            )
-        else:
-            # For MistralTokenizer
-            assert is_list_of(request_prompt, int), "Prompt has to be either a string or a list of token ids"
-            prompt_inputs = TextTokensPrompt(
-                prompt=tokenizer.decode(request_prompt),
-                prompt_token_ids=request_prompt,
-            )
-
-        engine_prompt = EngineTokensPrompt(prompt_token_ids=prompt_inputs["prompt_token_ids"])
-        if mm_data is not None:
-            engine_prompt["multi_modal_data"] = mm_data
-
-        if mm_uuids is not None:
-            engine_prompt["multi_modal_uuids"] = mm_uuids
-
-        if request.mm_processor_kwargs is not None:
-            engine_prompt["mm_processor_kwargs"] = request.mm_processor_kwargs
-
-        if hasattr(request, "cache_salt") and request.cache_salt is not None:
-            engine_prompt["cache_salt"] = request.cache_salt
-
-        return conversation, [request_prompt], [engine_prompt]
-
     def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[SamplingParams]:
         final_sampling_params_list = []
         for sampling_params in sampling_params_list:

vllm_omni/model_executor/models/qwen2_5_omni.py

@@ -7,7 +7,6 @@ from typing import NamedTuple, Optional, Union
 import numpy as np
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
     Qwen2_5OmniConfig,
     Qwen2_5OmniTalkerConfig,
@@ -15,7 +14,13 @@ from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
 )
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.qwen2_5_omni_thinker import (
+    Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniThinkerDummyInputsBuilder,
+    Qwen2_5OmniThinkerMultiModalProcessor,
+    Qwen2_5OmniThinkerProcessingInfo,
+)
 from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
 
 # from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav
@@ -26,14 +31,7 @@ from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
-from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
-    Qwen2_5OmniConditionalGenerationMixin,
-    Qwen2_5OmniThinkerDummyInputsBuilder,
-    Qwen2_5OmniThinkerMultiModalProcessor,
-    Qwen2_5OmniThinkerProcessingInfo,
-)
-from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
-from vllm_omni.model_executor.models.vision import get_llm_pos_ids_for_vision
+from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
 
 TALKER_CODEC_EOS_TOKEN_ID = 8294
 TALKER_CODEC_BOS_TOKEN_ID = 8293
@@ -56,7 +54,7 @@ logger = init_logger(__name__)
     dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
 class Qwen2_5OmniForConditionalGeneration(
-    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin, SupportsMRoPE
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
 ):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
@@ -461,170 +459,6 @@ class Qwen2_5OmniForConditionalGeneration(
             multimodal_outputs=None,
         )
 
-    def get_mrope_input_positions(
-        self,
-        input_tokens: list[int],
-        hf_config: PretrainedConfig,
-        image_grid_thw: list[list[int]] | torch.Tensor,
-        video_grid_thw: list[list[int]] | torch.Tensor,
-        second_per_grid_ts: list[float] | None = None,
-        context_len: int = 0,
-        seq_len: int | None = None,
-        audio_feature_lengths: torch.Tensor | None = None,
-        use_audio_in_video: bool = False,
-    ) -> tuple[torch.Tensor, int]:
-        """Get mrope input positions and delta value (Qwen2.5-Omni version).
-
-        Differences from MRotaryEmbedding:
-            1. Add audio support (and related `audio_feature_lengths`).
-            2. Add `use_audio_in_video` option to read audio from video inputs.
-                In this case, audio and vision position ids will be split into
-                chunks and interleaved.
-
-        Example:
-
-            (V_i are vision position ids, A_i are audio position ids)
-
-            |V_1 ...    V_n|A_1 ...   A_n|V_n+1 ... V_2n|A_n+1 ... A_2n|...
-            |vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |...
-        """
-
-        # TODO(fyabc): refactor and share more code with
-        #  _vl_get_input_positions_tensor.
-
-        thinker_config = hf_config.thinker_config
-        audio_token_id = thinker_config.audio_token_index
-        image_token_id = thinker_config.image_token_index
-        video_token_id = thinker_config.video_token_index
-        audio_start_token_id = thinker_config.audio_start_token_id
-        audio_end_token_id = thinker_config.audio_end_token_id
-        vision_start_token_id = thinker_config.vision_start_token_id
-        vision_end_token_id = thinker_config.vision_end_token_id
-        seconds_per_chunk = thinker_config.seconds_per_chunk
-        spatial_merge_size = thinker_config.vision_config.spatial_merge_size
-        tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25)
-
-        if isinstance(image_grid_thw, list):
-            image_grid_thw = torch.tensor(image_grid_thw)
-        if isinstance(video_grid_thw, list):
-            video_grid_thw = torch.tensor(video_grid_thw)
-
-        src_item = input_tokens
-        audio_seqlens = audio_feature_lengths
-        if not second_per_grid_ts:
-            second_per_grid_ts = [1] * video_grid_thw.shape[0]
-        audio_idx = 0
-        video_idx = 0
-        image_idx = 0
-        new_src_item: list[int] = []
-        llm_pos_ids_list: list[torch.Tensor] = []
-
-        idx = 0
-        while idx < len(src_item):
-            new_src_item_len = len(new_src_item)
-            start_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
-            if src_item[idx] not in [audio_token_id, video_token_id, image_token_id]:
-                if use_audio_in_video and idx > 0:
-                    if src_item[idx] == vision_end_token_id and src_item[idx - 1] == audio_end_token_id:
-                        # processing the <|audio_eos|> before <|vision_eos|>
-                        start_idx -= 1
-                    elif src_item[idx] == audio_start_token_id and src_item[idx - 1] == vision_start_token_id:
-                        # processing the <|audio_bos|> after <|vision_eos|>
-                        start_idx -= 1
-                new_src_item.append(src_item[idx])
-                llm_pos_ids = torch.tensor([start_idx], dtype=torch.long).expand(3, -1)
-                llm_pos_ids_list.append(llm_pos_ids)
-            elif src_item[idx] == audio_token_id:
-                assert audio_seqlens is not None
-                audio_seqlen = audio_seqlens[audio_idx]
-                place_num = ((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1
-                new_src_item.extend([audio_token_id] * place_num)
-                llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx
-                llm_pos_ids_list.append(llm_pos_ids)
-                audio_idx += 1
-            elif src_item[idx] == image_token_id:
-                grid_t = image_grid_thw[image_idx][0]
-                grid_hs = image_grid_thw[:, 1]
-                grid_ws = image_grid_thw[:, 2]
-                t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long()
-                llm_pos_ids = get_llm_pos_ids_for_vision(
-                    start_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
-                )
-                llm_pos_ids_list.append(llm_pos_ids)
-                vision_seqlen = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
-                new_src_item.extend([image_token_id] * vision_seqlen)
-                image_idx += 1
-            elif src_item[idx] == video_token_id and not use_audio_in_video:
-                grid_t = video_grid_thw[video_idx][0]
-                grid_hs = video_grid_thw[:, 1]
-                grid_ws = video_grid_thw[:, 2]
-                t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long()
-                llm_pos_ids = get_llm_pos_ids_for_vision(
-                    start_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
-                )
-                llm_pos_ids_list.append(llm_pos_ids)
-                vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
-                new_src_item.extend([video_token_id] * vision_seqlen)
-                video_idx += 1
-            else:
-                # read audio from video
-                assert audio_seqlens is not None
-                audio_seqlen = audio_seqlens[audio_idx]
-                vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
-                grid_t = video_grid_thw[video_idx][0]
-                grid_h = video_grid_thw[video_idx][1]
-                grid_w = video_grid_thw[video_idx][2]
-                grid_hs = video_grid_thw[:, 1]
-                grid_ws = video_grid_thw[:, 2]
-                t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
-                t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long()
-                t_index_split_chunk = split_list_into_ranges(t_index, t_ntoken_per_chunk)
-                place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2
-                pure_audio_len = place_num - 2
-                added_audio_len = 0
-                audio_llm_pos_ids_list: list[torch.Tensor] = []
-                for t_chunk in t_index_split_chunk:
-                    vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
-                    new_src_item.extend([video_token_id] * vision_ntoken_per_chunk)
-                    vision_llm_pos_ids_list = get_llm_pos_ids_for_vision(
-                        start_idx,
-                        video_idx,
-                        spatial_merge_size,
-                        t_chunk,
-                        grid_hs,
-                        grid_ws,
-                    ).split(1, dim=1)
-                    llm_pos_ids_list.extend(vision_llm_pos_ids_list)
-                    new_src_item.extend(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) * [audio_token_id])
-                    audio_start_idx = (
-                        start_idx if len(audio_llm_pos_ids_list) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1
-                    )
-                    if min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) > 0:
-                        audio_llm_pos_ids_list = (
-                            torch.arange(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)).expand(3, -1)
-                            + audio_start_idx
-                        ).split(1, dim=1)
-                    else:
-                        audio_llm_pos_ids_list = []
-                    added_audio_len += min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)
-                    llm_pos_ids_list.extend(audio_llm_pos_ids_list)
-                if added_audio_len < pure_audio_len:
-                    new_src_item.extend((pure_audio_len - added_audio_len) * [audio_token_id])
-                    audio_llm_pos_ids_list = (
-                        torch.arange(pure_audio_len - added_audio_len).expand(3, -1) + llm_pos_ids_list[-1].max() + 1
-                    ).split(1, dim=1)
-                    llm_pos_ids_list.extend(audio_llm_pos_ids_list)
-                audio_idx += 1
-                video_idx += 1
-            # move to the next token
-            idx += len(new_src_item) - new_src_item_len
-
-        llm_positions = torch.cat(llm_pos_ids_list, dim=1)
-        mrope_position_delta = torch.cat(llm_pos_ids_list, dim=1).max() + 1 - len(src_item)
-        llm_positions = llm_positions[:, context_len:seq_len]
-
-        return llm_positions, mrope_position_delta
-
     def generate_audio(self, code, voice_type):
         token2wav_dev = self._module_device(self.token2wav)
         if isinstance(code, torch.Tensor):

vllm_omni/model_executor/models/qwen2_5_omni_thinker.py

@@ -1,7 +1,7 @@
-from collections import defaultdict
 from collections.abc import Iterable, Mapping, Sequence
+from copy import copy
 from functools import partial
-from typing import Annotated, Any, Callable, Literal, Optional, Union
+from typing import Any, Optional, Union
 
 import torch
 import torch.nn as nn
@@ -12,7 +12,7 @@ from transformers.models.qwen2_5_omni.processing_qwen2_5_omni import Qwen2_5Omni
 from transformers.models.whisper import WhisperFeatureExtractor
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMRoPE, SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
 from vllm.model_executor.models.qwen2_5_vl import (
     Qwen2_5_VisionTransformer,
     Qwen2_5_VLImageEmbeddingInputs,
@@ -24,6 +24,7 @@ from vllm.model_executor.models.qwen2_5_vl import (
     Qwen2_5_VLVideoPixelInputs,
 )
 from vllm.model_executor.models.qwen2_audio import (
+    Qwen2AudioInputs,
     Qwen2AudioProcessingInfo,
     _get_feat_extract_output_lengths,
 )
@@ -41,7 +42,7 @@ from vllm.multimodal.inputs import (
     ModalityData,
     MultiModalDataDict,
     MultiModalFieldConfig,
-    MultiModalKwargsItems,
+    MultiModalKwargs,
     NestedTensors,
 )
 from vllm.multimodal.parse import (
@@ -51,19 +52,10 @@ from vllm.multimodal.parse import (
     MultiModalDataItems,
     MultiModalDataParser,
 )
-from vllm.multimodal.processing import (
-    BaseMultiModalProcessor,
-    MultiModalPromptUpdates,
-    PlaceholderFeaturesInfo,
-    PromptReplacement,
-    PromptUpdate,
-    PromptUpdateDetails,
-    ResolvedPromptUpdate,
-)
+from vllm.multimodal.processing import BaseMultiModalProcessor, PlaceholderFeaturesInfo, PromptReplacement, PromptUpdate
 from vllm.multimodal.profiling import BaseDummyInputsBuilder
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.tokenizer import decode_tokens, encode_tokens
-from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
 from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
 
@@ -75,27 +67,6 @@ except (ImportError, ModuleNotFoundError):
 logger = init_logger(__name__)
 
 
-class Qwen2_5OmniAudioFeatureInputs(TensorSchema):
-    """
-    Dimensions:
-        - na: Number of audios
-        - nmb: Number of mel bins
-        - msl: Maximum sequence length
-        - tsl: Total sequence length
-    """
-
-    type: Literal["audio_features"]
-    input_features: Annotated[
-        Union[torch.Tensor, list[torch.Tensor]],
-        TensorShape("nmb", "tsl"),
-    ]
-
-    feature_attention_mask: Annotated[
-        torch.Tensor,
-        TensorShape("na", "msl"),
-    ]
-
-
 def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
     audio_feature_lengths = hf_inputs.get("audio_feature_lengths", torch.empty((0,)))
 
@@ -119,39 +90,6 @@ def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
     )
 
 
-def create_qwen2_5_omni_thinker_field_factory(
-    spatial_merge_size: int,
-) -> Callable[[Mapping[str, torch.Tensor]], Mapping[str, MultiModalFieldConfig]]:
-    def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
-        audio_feature_lengths = hf_inputs.get("audio_feature_lengths", torch.empty((0,)))
-
-        image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))
-        image_pixel_grid_sizes = image_grid_thw.prod(-1)
-        image_embed_grid_sizes = image_pixel_grid_sizes // spatial_merge_size // spatial_merge_size
-
-        video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3)))
-        video_grid_sizes = video_grid_thw.prod(-1)
-        video_embed_grid_sizes = video_grid_sizes // spatial_merge_size // spatial_merge_size
-
-        num_videos = len(video_grid_sizes)
-
-        return dict(
-            input_audio_features=MultiModalFieldConfig.flat_from_sizes("audio", audio_feature_lengths, dim=1),
-            feature_attention_mask=MultiModalFieldConfig.batched("audio"),
-            audio_feature_lengths=MultiModalFieldConfig.batched("audio"),
-            pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_pixel_grid_sizes),
-            image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_embed_grid_sizes),
-            image_grid_thw=MultiModalFieldConfig.batched("image"),
-            pixel_values_videos=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes),
-            video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_embed_grid_sizes),
-            video_grid_thw=MultiModalFieldConfig.batched("video"),
-            second_per_grid_ts=MultiModalFieldConfig.batched("video"),
-            use_audio_in_video=MultiModalFieldConfig.shared("video", num_videos),
-        )
-
-    return _qwen2_5_omni_thinker_field_config
-
-
 class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
     def _parse_audio_data(
         self,
@@ -228,7 +166,10 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(BaseDummyInputsBuilder[Qwen2_5OmniThi
             "audio": self._get_dummy_audios(length=target_audio_length, num_audios=num_audios),
             "image": self._get_dummy_images(width=target_width, height=target_height, num_images=num_images),
             "video": self._get_dummy_videos(
-                width=target_width, height=target_height, num_frames=target_num_frames, num_videos=num_videos
+                width=target_width,
+                height=target_height,
+                num_frames=target_num_frames,
+                num_videos=num_videos,
             ),
         }
 
@@ -238,10 +179,7 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(BaseDummyInputsBuilder[Qwen2_5OmniThi
 class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]):
     def _get_data_parser(self) -> MultiModalDataParser:
         feature_extractor = self.info.get_feature_extractor()
-        return Qwen2_5OmniThinkerMultiModalDataParser(
-            spatial_merge_size=self.info.get_hf_config().vision_config.spatial_merge_size,
-            target_sr=feature_extractor.sampling_rate,
-        )
+        return Qwen2_5OmniThinkerMultiModalDataParser(target_sr=feature_extractor.sampling_rate)
 
     def _call_hf_processor(
         self,
@@ -276,14 +214,6 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
             hf_inputs["input_audio_features"] = input_features
         if "audio_feature_lengths" not in hf_inputs and feature_attention_mask is not None:
             hf_inputs["audio_feature_lengths"] = feature_attention_mask.sum(-1)
-
-        video_second_per_grid = hf_inputs.get("video_second_per_grid", None)
-        if video_second_per_grid is not None:
-            hf_inputs["second_per_grid_ts"] = video_second_per_grid
-
-        use_audio_in_video = mm_kwargs.get("use_audio_in_video", False)
-        hf_inputs["use_audio_in_video"] = torch.tensor(use_audio_in_video)
-
         return hf_inputs
 
     def _get_mm_fields_config(
@@ -291,177 +221,66 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
         hf_inputs: BatchFeature,
         hf_processor_mm_kwargs: Mapping[str, object],
     ) -> Mapping[str, MultiModalFieldConfig]:
-        return create_qwen2_5_omni_thinker_field_factory(self.info.get_hf_config().vision_config.spatial_merge_size)(
-            hf_inputs
-        )
-
-    def _apply_prompt_updates(
-        self,
-        token_ids: list[int],
-        mm_prompt_updates: MultiModalPromptUpdates,
-    ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
-        tokenizer = self.info.get_tokenizer()
-
-        new_token_ids, match_result = self._apply_token_matches(
-            token_ids,
-            mm_prompt_updates,
-        )
-
-        # If the search text does not represent a special token,
-        # it may have different token IDs in the prompt, because
-        # the tokens may go across the boundaries of the search text.
-        # ----
-        # e.g. when searching for "foo" in "food", if "food" itself makes
-        # up a token, then the token ID of "foo" will not appear at all
-        # ----
-        # Since it is inefficient to search for all possible tokenizations
-        # of the search text in the prompt, we instead perform string-based
-        # updates on the decoded token IDs, then encode them back.
-        if not all(all(update_idx is not None for update_idx in update_idxs) for update_idxs in match_result.values()):
-            new_text, match_result = self._apply_text_matches(
-                decode_tokens(tokenizer, token_ids),
-                mm_prompt_updates,
-            )
-
-            new_token_ids = encode_tokens(
-                tokenizer,
-                new_text,
-                add_special_tokens=False,
-            )
-
-        matched_updates = defaultdict[str, list[Sequence[ResolvedPromptUpdate]]](list)
-        unmatched_audio_items = []
-
-        for modality, update_idxs in match_result.items():
-            for item_idx, update_idx in enumerate(update_idxs):
-                if update_idx is None:
-                    # Check if this is audio that might be embedded in video
-                    if modality == "audio":
-                        unmatched_audio_items.append(item_idx)
-                        continue
-                    else:
-                        assert False, f"Failed to apply prompt replacement for mm_items[{modality!r}][{item_idx}]"
-
-                matched_updates[modality].append([mm_prompt_updates[modality][item_idx][update_idx]])
-
-        # If there are unmatched audio items, check if we can derive them from video
-        if unmatched_audio_items:
-            # Check if video exists in the original updates (not just matched ones)
-            if "video" in mm_prompt_updates:
-                num_videos = len(mm_prompt_updates.get("video", []))
-                num_audios = len(mm_prompt_updates.get("audio", []))
-                # If counts match, audio might be embedded in video - skip assertion
-                # The placeholders will be derived from video in _maybe_apply_prompt_updates
-                if num_audios == num_videos and "video" in matched_updates:
-                    pass  # Will be handled by deriving from video placeholders
-                else:
-                    # Audio items exist but can't be matched
-                    # Check if video was also matched
-                    if "video" not in matched_updates:
-                        # Neither audio nor video matched - this is a real error
-                        assert False, (
-                            f"Failed to apply prompt replacement for {len(unmatched_audio_items)} "
-                            f"audio item(s). Audio items cannot be matched in prompt."
-                        )
-                    elif num_audios != num_videos:
-                        # Audio and video counts don't match
-                        assert False, (
-                            f"Failed to apply prompt replacement for {len(unmatched_audio_items)} "
-                            f"audio item(s). Audio items cannot be matched in prompt and "
-                            f"audio count ({num_audios}) does not match video count ({num_videos})."
-                        )
-                    else:
-                        # Video matched but audio didn't, and counts match - allow it
-                        pass
-            else:
-                # Audio items can't be matched and there are no video items
-                assert False, (
-                    f"Failed to apply prompt replacement for {len(unmatched_audio_items)} "
-                    f"audio item(s). Audio items cannot be matched in prompt."
-                )
-
-        placeholders = self._find_mm_placeholders(
-            new_token_ids,
-            dict(matched_updates),
-        )
-
-        return new_token_ids, placeholders
+        return _qwen2_5_omni_thinker_field_config(hf_inputs)
 
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
         prompt_ids: list[int],
-        mm_kwargs: MultiModalKwargsItems,
-        mm_prompt_updates: MultiModalPromptUpdates,
+        mm_kwargs: MultiModalKwargs,
         is_update_applied: bool,
     ) -> tuple[list[int], str, Mapping[str, list[PlaceholderFeaturesInfo]]]:
         """
         Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.
         """
+        unbound_prompt_updates = self._get_prompt_updates(
+            mm_items,
+            hf_processor_mm_kwargs,
+            mm_kwargs,
+        )
+        mm_prompt_updates = self._bind_and_group_updates(unbound_prompt_updates)
+
         mm_item_counts = mm_items.get_all_counts()
         self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
 
-        use_audio_in_video = False
-        if "video" in mm_kwargs:
-            video_items = [item for item in mm_kwargs["video"] if item is not None]
-            # only check video items (if there are any)
-            if video_items:
-                use_audio_in_video = all(item["use_audio_in_video"].data for item in video_items)
+        use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)
 
         if is_update_applied:
             mm_placeholders = self._find_mm_placeholders(
+                mm_prompt_updates,
+                prompt_ids,
+                mm_item_counts,
+            )
+            self._validate_mm_placeholders(mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video)
+
+            tokenizer = self.info.get_tokenizer()
+            prompt = decode_tokens(tokenizer, prompt_ids)
+        else:
+            (
+                prompt_ids,
+                prompt,
+                mm_placeholders,
+            ) = self._apply_prompt_updates(
                 prompt_ids,
                 mm_prompt_updates,
-            )
-            self._validate_mm_placeholders(
-                mm_placeholders,
-                mm_item_counts,
-            )
-
-        else:
-            # When use_audio_in_video=True, audio tokens are not in the prompt,
-            # so we need to filter out audio updates before applying replacements
-            if use_audio_in_video and "audio" in mm_prompt_updates:
-                # Remove audio from prompt updates (it won't match anything)
-                filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"}
-                prompt_ids, mm_placeholders = self._apply_prompt_updates(
-                    prompt_ids,
-                    filtered_updates,
-                )
-                # Derive audio placeholders from video placeholders
-                mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates)
-            else:
-                # Apply prompt updates normally
-                # _apply_prompt_updates will handle unmatched audio items
-                # when video exists with matching counts
-                prompt_ids, mm_placeholders = self._apply_prompt_updates(
-                    prompt_ids,
-                    mm_prompt_updates,
-                )
-
-                # Check if audio placeholders are missing but should exist
-                # This can happen when audio items are embedded in video tokens
-                if "audio" in mm_prompt_updates and "audio" not in mm_placeholders and "video" in mm_placeholders:
-                    num_audios = len(mm_prompt_updates.get("audio", []))
-                    num_videos = len(mm_placeholders.get("video", []))
-                    # If counts match, derive audio placeholders from video
-                    if num_audios == num_videos:
-                        mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates)
-            self._validate_mm_placeholders(
-                mm_placeholders,
                 mm_item_counts,
             )
+            self._validate_mm_placeholders(mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video)
 
         tokenizer = self.info.get_tokenizer()
         prompt = decode_tokens(tokenizer, prompt_ids)
 
+        if use_audio_in_video:
+            mm_kwargs["use_audio_in_video"] = True
+
         return prompt_ids, prompt, mm_placeholders
 
     def _get_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
         hf_processor_mm_kwargs: Mapping[str, Any],
-        out_mm_kwargs: MultiModalKwargsItems,
+        out_mm_kwargs: MultiModalKwargs,
     ) -> Sequence[PromptUpdate]:
         processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
         tokenizer = self.info.get_tokenizer()
@@ -475,9 +294,8 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
         image_token_id = vocab[image_token]
         video_token_id = vocab[video_token]
 
-        out_mm_data = out_mm_kwargs.get_data()
-        audio_feature_lengths = out_mm_data.get("audio_feature_lengths")
-        feature_attention_mask = out_mm_data.get("feature_attention_mask")
+        audio_feature_lengths = out_mm_kwargs.get("audio_feature_lengths")
+        feature_attention_mask = out_mm_kwargs.get("feature_attention_mask")
         if audio_feature_lengths is None and feature_attention_mask is None:
             audio_output_lengths = []
         elif audio_feature_lengths is not None:
@@ -505,7 +323,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
             return [audio_token_id] * num_features
 
         def get_replacement_qwen2_vision(item_idx: int, modality: str):
-            grid_thw = out_mm_data[f"{modality}_grid_thw"][item_idx]
+            grid_thw = out_mm_kwargs[f"{modality}_grid_thw"][item_idx]
             assert isinstance(grid_thw, torch.Tensor)
             merge_length = image_processor.merge_size**2
 
@@ -519,7 +337,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
             nonlocal audio_in_video_item_idx
 
             audio_num_features = audio_output_lengths[audio_in_video_item_idx + item_idx]
-            video_grid_thw = out_mm_data["video_grid_thw"][item_idx]
+            video_grid_thw = out_mm_kwargs["video_grid_thw"][item_idx]
 
             audio_in_video_item_idx += 1
 
@@ -529,15 +347,13 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
             else:
                 video_second_per_grid_t = 1.0
 
-            placeholder = MRotaryEmbedding.omni_get_updates_use_audio_in_video(
+            return MRotaryEmbedding.omni_get_updates_use_audio_in_video(
                 thinker_config=thinker_config,
                 audio_len=audio_num_features,
                 video_grid_thw=video_grid_thw,
                 video_second_per_grid_t=video_second_per_grid_t,
             )
 
-            return PromptUpdateDetails.select_token_id(placeholder, embed_token_id=video_token_id)
-
         video_replacement_fn = (
             get_replacement_qwen2_use_audio_in_video
             if use_audio_in_video
@@ -562,50 +378,6 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
             ),
         ]
 
-    def _derive_audio_from_video_placeholders(
-        self,
-        placeholders: Mapping[str, list[PlaceholderFeaturesInfo]],
-        mm_prompt_updates: MultiModalPromptUpdates,
-    ) -> Mapping[str, list[PlaceholderFeaturesInfo]]:
-        """
-        Helper to derive audio placeholders from video placeholders when
-        use_audio_in_video=True.
-        """
-        if "video" not in placeholders:
-            return placeholders
-
-        # Validate audio and video counts match
-        num_videos = len(placeholders["video"])
-        num_audios = len(mm_prompt_updates.get("audio", []))
-        if num_audios != num_videos:
-            raise ValueError(
-                f"use_audio_in_video requires equal number of audio and video items, got {num_audios=}, {num_videos=}"
-            )
-
-        tokenizer = self.info.get_tokenizer()
-        processor = self.info.get_hf_processor()
-        audio_token_id = tokenizer.get_vocab()[processor.audio_token]
-
-        result_placeholders = dict(placeholders)
-        audio_placeholders = []
-
-        # Each video is paired with one audio
-        for video_idx, video_placeholder in enumerate(placeholders["video"]):
-            # Create is_embed mask selecting only audio tokens
-            audio_is_embed = torch.tensor(video_placeholder.tokens) == audio_token_id
-
-            audio_placeholder = PlaceholderFeaturesInfo(
-                modality="audio",
-                item_idx=video_idx,
-                start_idx=video_placeholder.start_idx,
-                tokens=video_placeholder.tokens,
-                is_embed=audio_is_embed,
-            )
-            audio_placeholders.append(audio_placeholder)
-
-        result_placeholders["audio"] = audio_placeholders
-        return result_placeholders
-
     def _apply_hf_processor_main(
         self,
         prompt: Union[str, list[int]],
@@ -664,19 +436,30 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
 
         return mm_processed_data
 
+    def _validate_mm_placeholders(
+        self,
+        mm_placeholders: Mapping[str, list[PlaceholderFeaturesInfo]],
+        mm_item_counts: Mapping[str, int],
+        use_audio_in_video: bool = False,
+    ) -> None:
+        if use_audio_in_video:
+            mm_item_counts = copy(mm_item_counts)
+            if "video" in mm_item_counts:
+                assert "audio" in mm_item_counts
+                mm_item_counts["audio"] -= mm_item_counts["video"]
+        super()._validate_mm_placeholders(mm_placeholders, mm_item_counts)
+
 
 class Qwen2_5OmniConditionalGenerationMixin:
     def _validate_and_reshape_mm_tensor(self, mm_input: object, name: str, dim: int = 0) -> torch.Tensor:
         if not isinstance(mm_input, (torch.Tensor, list)):
             raise ValueError(f"Incorrect type of {name}. Got type: {type(mm_input)}")
         if isinstance(mm_input, torch.Tensor):
-            if dim == 0:
-                return mm_input.reshape(-1, *mm_input.shape[2:])
             return torch.concat(list(mm_input), dim=dim)
         else:
             return torch.concat(mm_input, dim=dim)
 
-    def _parse_and_validate_audio_input(self, **kwargs: object) -> Optional[Qwen2_5OmniAudioFeatureInputs]:
+    def _parse_and_validate_audio_input(self, **kwargs: object) -> Optional[Qwen2AudioInputs]:
         input_audio_features = kwargs.pop("input_audio_features", None)
         audio_feature_lengths = kwargs.pop("audio_feature_lengths", None)
         feature_attention_mask = kwargs.pop("feature_attention_mask", None)
@@ -689,8 +472,7 @@ class Qwen2_5OmniConditionalGenerationMixin:
             )
         if not isinstance(input_audio_features, (torch.Tensor, list)):
             raise ValueError(f"Incorrect type of audio input features. Got type: {type(input_audio_features)}")
-        return Qwen2_5OmniAudioFeatureInputs(
-            type="audio_features",
+        return Qwen2AudioInputs(
             input_features=input_audio_features,
             audio_feature_lengths=audio_feature_lengths,
             feature_attention_mask=feature_attention_mask,
@@ -715,7 +497,9 @@ class Qwen2_5OmniConditionalGenerationMixin:
                 raise ValueError(f"Incorrect type of image pixel values. Got type: {type(pixel_values)}")
 
             return Qwen2_5_VLImagePixelInputs(
-                type="pixel_values", pixel_values=pixel_values, image_grid_thw=image_grid_thw
+                type="pixel_values",
+                pixel_values=pixel_values,
+                image_grid_thw=image_grid_thw,
             )
 
         if image_embeds is not None:
@@ -725,7 +509,9 @@ class Qwen2_5OmniConditionalGenerationMixin:
             if not isinstance(image_embeds, torch.Tensor):
                 raise ValueError(f"Incorrect type of image embeddings. Got type: {type(image_embeds)}")
             return Qwen2_5_VLImageEmbeddingInputs(
-                type="image_embeds", image_embeds=image_embeds, image_grid_thw=image_grid_thw
+                type="image_embeds",
+                image_embeds=image_embeds,
+                image_grid_thw=image_grid_thw,
             )
 
     def _parse_and_validate_video_input(
@@ -756,12 +542,14 @@ class Qwen2_5OmniConditionalGenerationMixin:
             if not isinstance(video_embeds, torch.Tensor):
                 raise ValueError(f"Incorrect type of video embeddings. Got type: {type(video_embeds)}")
             return Qwen2_5_VLVideoEmbeddingInputs(
-                type="video_embeds", video_embeds=video_embeds, video_grid_thw=video_grid_thw
+                type="video_embeds",
+                video_embeds=video_embeds,
+                video_grid_thw=video_grid_thw,
             )
 
     def _process_audio_input(
         self,
-        audio_input: Qwen2_5OmniAudioFeatureInputs,
+        audio_input: Qwen2AudioInputs,
         audio_hashes: list[str] = None,
         cached_audio_features: torch.Tensor = None,
     ) -> torch.Tensor:
@@ -786,7 +574,8 @@ class Qwen2_5OmniConditionalGenerationMixin:
             feature_lens=audio_feature_lengths,
             aftercnn_lens=audio_feat_lengths,
         )
-        return audio_outputs.last_hidden_state.split(audio_output_lengths.tolist())
+        audio_features = audio_outputs.last_hidden_state
+        return audio_features.split(audio_output_lengths.tolist())
 
     def _process_image_input(self, image_input: Qwen2_5_VLImageInputs) -> tuple[torch.Tensor, ...]:
         if image_input["type"] == "image_embeds":
@@ -830,7 +619,7 @@ class Qwen2_5OmniConditionalGenerationMixin:
     dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
 class Qwen2_5OmniThinkerForConditionalGeneration(
-    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin, SupportsMRoPE
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
 ):
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
@@ -871,21 +660,13 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
                 "exactly same result as the transformers implementation in the audio tower part."
             )
 
-        if multimodal_config.get_limit_per_prompt("audio"):
-            self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
-        else:
-            self.audio_tower = None
-
-        if multimodal_config.get_limit_per_prompt("image") or multimodal_config.get_limit_per_prompt("video"):
-            self.visual = Qwen2_5_VisionTransformer(
-                vision_config=thinker_config.vision_config,
-                norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
-                quant_config=quant_config,
-                prefix=maybe_prefix(prefix, "visual"),
-                # attn_backend_override=None,
-            )
-        else:
-            self.visual = None
+        self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
+        self.visual = Qwen2_5_VisionTransformer(
+            vision_config=thinker_config.vision_config,
+            norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "visual"),
+        )
         self.quant_config = quant_config
         self.language_model = init_vllm_registered_model(
             vllm_config=vllm_config,

vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py

@@ -23,6 +23,7 @@ from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniPr
 from transformers.utils.logging import get_logger as _hf_get_logger
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
+from vllm.model_executor.model_loader.utils import set_default_torch_dtype
 from vllm.model_executor.models.interfaces import SupportsPP
 from vllm.model_executor.models.utils import AutoWeightsLoader as _Vllm_AutoWeightsLoader
 from vllm.model_executor.models.utils import WeightsMapper as _Vllm_WeightsMapper
@@ -33,6 +34,8 @@ from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
+from vllm_omni.model_executor.ode_solver import get_ode_solver_class
+
 
 # Provide a no-op auto_docstring decorator to satisfy annotations if missing
 def auto_docstring(func=None, **_kwargs):
@@ -569,7 +572,8 @@ class DiTAttention(nn.Module):
         cos, sin = position_embeddings
         query[:, :1], key[:, :1] = apply_rotary_pos_emb(query[:, :1], key[:, :1], cos, sin)
 
-        attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+        attn_impl = self.config._attn_implementation or "sdpa"
+        attention_interface = ALL_ATTENTION_FUNCTIONS[attn_impl]
         attention_weights, _ = attention_interface(
             self,
             query,
@@ -1014,100 +1018,19 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
         return torch.clamp(output_waveform, min=-1.0, max=1.0).squeeze().cpu()
 
 
-class RungeKutta4ODESolver:
-    def __init__(self, function, initial_value):
-        self.function = function
-        self.initial_value = initial_value
-
-        self._one_third = 1 / 3
-        self._two_thirds = 2 / 3
-
-    def _rk4_step(
-        self,
-        function,
-        time_start,
-        time_step,
-        time_end,
-        value_start,
-        function_value_start=None,
-    ):
-        k1 = function_value_start if function_value_start is not None else function(time_start, value_start)
-        k2 = function(
-            time_start + time_step * self._one_third,
-            value_start + time_step * k1 * self._one_third,
-        )
-        k3 = function(
-            time_start + time_step * self._two_thirds,
-            value_start + time_step * (k2 - k1 * self._one_third),
-        )
-        k4 = function(time_end, value_start + time_step * (k1 - k2 + k3))
-        return (k1 + 3 * (k2 + k3) + k4) * time_step / 8
-
-    def _compute_step(self, function, time_start, time_step, time_end, value_start):
-        function_value_start = function(time_start, value_start)
-        return (
-            self._rk4_step(
-                function,
-                time_start,
-                time_step,
-                time_end,
-                value_start,
-                function_value_start=function_value_start,
-            ),
-            function_value_start,
-        )
-
-    def _linear_interpolation(self, time_start, time_end, value_start, value_end, time_point):
-        if time_point == time_start:
-            return value_start
-        if time_point == time_end:
-            return value_end
-        weight = (time_point - time_start) / (time_end - time_start)
-        return value_start + weight * (value_end - value_start)
-
-    def integrate(self, time_points):
-        solution = torch.empty(
-            len(time_points),
-            *self.initial_value.shape,
-            dtype=self.initial_value.dtype,
-            device=self.initial_value.device,
-        )
-        solution[0] = self.initial_value
-
-        current_index = 1
-        current_value = self.initial_value
-        for time_start, time_end in zip(time_points[:-1], time_points[1:]):
-            time_step = time_end - time_start
-            delta_value, _ = self._compute_step(self.function, time_start, time_step, time_end, current_value)
-            next_value = current_value + delta_value
-
-            while current_index < len(time_points) and time_end >= time_points[current_index]:
-                solution[current_index] = self._linear_interpolation(
-                    time_start,
-                    time_end,
-                    current_value,
-                    next_value,
-                    time_points[current_index],
-                )
-                current_index += 1
-
-            current_value = next_value
-
-        return solution
-
-
 @auto_docstring(
     custom_intro="""
     The full Qwen2.5Omni Token2WavDiT model. Which take speech tokens as
     input and predict mel spectrogram.
     """
 )
-class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
-    config: Qwen2_5OmniDiTConfig
-    _no_split_modules = ["DiTDecoderLayer"]
+class Qwen2_5OmniToken2WavDiTModel(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        model_config = vllm_config.model_config
+        config: Qwen2_5OmniDiTConfig = model_config.hf_config
 
-    def __init__(self, config: Qwen2_5OmniDiTConfig):
-        super().__init__(config)
         self.mel_dim = config.mel_dim
         self.repeats = config.repeats
         self.time_embed = DiTTimestepEmbedding(config.hidden_size)
@@ -1135,6 +1058,10 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         self.norm_out = Qwen2_5_OmniAdaLayerNormZero_Final(config.hidden_size)  # final modulation
         self.proj_out = nn.Linear(config.hidden_size, config.mel_dim)
 
+        self.ode_solver = get_ode_solver_class(vllm_config.model_config.ode_solver_class)(
+            function=None, initial_value=None
+        )  # placeholder
+
     def _create_block_diff(self, hidden_states):
         batch, seq_len = hidden_states.shape[0], hidden_states.shape[1]
         block_indices = torch.arange(seq_len, device=hidden_states.device) // self.block_size  # [seq_length]
@@ -1248,8 +1175,8 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         if sway_coefficient is not None:
             time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
 
-        ode_solver = RungeKutta4ODESolver(function=ode_function, initial_value=initial_state)
-        solution_trajectory = ode_solver.integrate(time_embedding)
+        self.ode_solver.reset_ode_solver(function=ode_function, initial_value=initial_state)
+        solution_trajectory = self.ode_solver.integrate(time_embedding)
 
         generated_waveform = solution_trajectory[-1]
         generated_mel_spectrogram = generated_waveform.permute(0, 2, 1)
@@ -1320,8 +1247,8 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         if sway_coefficient is not None:
             time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
 
-        ode_solver = RungeKutta4ODESolver(function=ode_function, initial_value=initial_state)
-        solution_trajectory = ode_solver.integrate(time_embedding)
+        self.ode_solver.reset_ode_solver(function=ode_function, initial_value=initial_state)
+        solution_trajectory = self.ode_solver.integrate(time_embedding)
 
         generated_waveform = solution_trajectory[-1]
         generated_mel_spectrogram = generated_waveform.permute(0, 2, 1)
@@ -1335,16 +1262,13 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
     mel spectrogram as input and predict waveform.
     """
 )
-class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
-    config: Qwen2_5OmniToken2WavConfig
-    base_model_prefix = "model"
-    _no_split_modules = [
-        "Qwen2_5OmniToken2WavDiTModel",
-        "Qwen2_5OmniToken2WavBigVGANModel",
-    ]
+class Qwen2_5OmniToken2WavModel(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix=""):
+        super().__init__()
+
+        self.model_config = vllm_config.model_config
+        config: Qwen2_5OmniToken2WavConfig = self.model_config.hf_config
 
-    def __init__(self, config: Qwen2_5OmniToken2WavConfig):
-        super().__init__(config)
         attn_impl = config._attn_implementation
         if config._attn_implementation == "flash_attention_2":
             logger.warning_once(
@@ -1359,9 +1283,10 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
                 "Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa"
             )
             attn_impl = "sdpa"
-        self.code2wav_dit_model = Qwen2_5OmniToken2WavDiTModel._from_config(
-            config.dit_config, attn_implementation=attn_impl
-        )
+        with set_default_torch_dtype(torch.float32):
+            self.code2wav_dit_model = Qwen2_5OmniToken2WavDiTModel(
+                vllm_config=vllm_config.with_hf_config(config.dit_config),
+            )
         self.code2wav_bigvgan_model = Qwen2_5OmniToken2WavBigVGANModel._from_config(
             config.bigvgan_config, attn_implementation=attn_impl
         )
@@ -1394,6 +1319,10 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
             # fallback to commonly used value
             self.vocoder_hop = 240
 
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
     def forward(
         self,
         code,

vllm_omni/model_executor/models/registry.py

@@ -2,31 +2,26 @@ from vllm.model_executor.models.registry import _VLLM_MODELS, _LazyRegisteredMod
 
 _OMNI_MODELS = {
     "Qwen2_5OmniForConditionalGeneration": (
-        "qwen2_5_omni",
         "qwen2_5_omni",
         "Qwen2_5OmniForConditionalGeneration",
     ),
     "Qwen2_5OmniThinkerModel": (
-        "qwen2_5_omni",
         "qwen2_5_omni_thinker",
         "Qwen2_5OmniThinkerForConditionalGeneration",
     ),  # noqa: E501
     "Qwen2_5OmniTalkerModel": (
-        "qwen2_5_omni",
         "qwen2_5_omni_talker",
         "Qwen2_5OmniTalkerForConditionalGeneration",
     ),  # noqa: E501
     "Qwen2_5OmniToken2WavModel": (
-        "qwen2_5_omni",
         "qwen2_5_omni_token2wav",
         "Qwen2_5OmniToken2WavForConditionalGenerationVLLM",
     ),
     "Qwen2_5OmniToken2WavDiTModel": (
-        "qwen2_5_omni",
         "qwen2_5_omni_token2wav",
         "Qwen2_5OmniToken2WavModel",
     ),
-    "Qwen2ForCausalLM_old": ("qwen2_5_omni", "qwen2_old", "Qwen2ForCausalLM"),  # need to discuss
+    "Qwen2ForCausalLM_old": ("qwen2_old", "Qwen2ForCausalLM"),  # need to discuss
 }
 
 _VLLM_OMNI_MODELS = {
@@ -46,10 +41,10 @@ OmniModelRegistry = _ModelRegistry(
         },
         **{
             model_arch: _LazyRegisteredModel(
-                module_name=f"vllm_omni.model_executor.models.{mod_folder}.{mod_relname}",
+                module_name=f"vllm_omni.model_executor.models.{mod_relname}",
                 class_name=cls_name,
             )
-            for model_arch, (mod_folder, mod_relname, cls_name) in _OMNI_MODELS.items()
+            for model_arch, (mod_relname, cls_name) in _OMNI_MODELS.items()
         },
     }
 )

vllm_omni/model_executor/models/utils.py

@@ -1,4 +1,3 @@
-import torch
 from vllm.model_executor.models.utils import maybe_prefix
 
 
@@ -7,24 +6,3 @@ def add_prefix_to_loaded_weights(weights: set[str], prefix: str) -> set[str]:
     Add a prefix to the names of the loaded weights.
     """
     return {maybe_prefix(prefix, name) for name in weights}
-
-
-def split_list_into_ranges(lst: torch.Tensor, interval: int) -> list[list[int]]:
-    if lst.numel() == 0:
-        return []
-
-    # Move to CPU and convert to list once (High Speedup)
-    # using .item() inside a loop is very slow.
-    data_list = lst.detach().cpu().tolist()
-
-    # Calculate max on the list or tensor (Tensor max is fast enough)
-    max_val = int(torch.max(lst).item())
-
-    # Pre-allocate buckets
-    ranges: list[list[int]] = [[] for _ in range((max_val // interval) + 1)]
-
-    for num in data_list:
-        index = int(num // interval)
-        ranges[index].append(num)
-
-    return ranges

vllm_omni/model_executor/models/vision.py

@@ -1,23 +0,0 @@
-import torch
-
-
-def get_llm_pos_ids_for_vision(
-    start_idx: int,
-    vision_idx: int,
-    spatial_merge_size: int,
-    t_index: list[int],
-    grid_hs: torch.Tensor,
-    grid_ws: torch.Tensor,
-) -> torch.Tensor:
-    llm_pos_ids_list = []
-    llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
-    llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
-    h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten()
-    w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten()
-    t_index_tensor = (
-        torch.Tensor(t_index).to(llm_grid_h.device).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).long().flatten()
-    )
-    _llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index])
-    llm_pos_ids_list.append(_llm_pos_ids + start_idx)
-    llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
-    return llm_pos_ids

vllm_omni/model_executor/ode_solver/__init__.py

@@ -0,0 +1,15 @@
+from vllm_omni.model_executor.ode_solver.runger_kutta_4 import RungeKutta4ODESolver
+
+_ODE_SOLVERS = {
+    "RungeKutta4ODESolver": RungeKutta4ODESolver,
+}
+
+
+def get_ode_solver_class(class_name: str):
+    solver = _ODE_SOLVERS.get(class_name)
+    if solver is None:
+        raise ValueError(f"ODE solver class '{class_name}' not found.")
+    return solver
+
+
+__all__ = ["RungeKutta4ODESolver"]

vllm_omni/model_executor/ode_solver/runger_kutta_4.py

@@ -0,0 +1,89 @@
+from typing import Callable
+
+import torch
+
+
+class RungeKutta4ODESolver:
+    def __init__(self, function: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], initial_value: torch.Tensor):
+        self.function = function
+        self.initial_value = initial_value
+
+        self._one_third = 1 / 3
+        self._two_thirds = 2 / 3
+
+    def reset_ode_solver(self, function, initial_value):
+        self.function = function
+        self.initial_value = initial_value
+
+    def _rk4_step(
+        self,
+        function,
+        time_start,
+        time_step,
+        time_end,
+        value_start,
+        function_value_start=None,
+    ):
+        k1 = function_value_start if function_value_start is not None else function(time_start, value_start)
+        k2 = function(
+            time_start + time_step * self._one_third,
+            value_start + time_step * k1 * self._one_third,
+        )
+        k3 = function(
+            time_start + time_step * self._two_thirds,
+            value_start + time_step * (k2 - k1 * self._one_third),
+        )
+        k4 = function(time_end, value_start + time_step * (k1 - k2 + k3))
+        return (k1 + 3 * (k2 + k3) + k4) * time_step / 8
+
+    def _compute_step(self, function, time_start, time_step, time_end, value_start):
+        function_value_start = function(time_start, value_start)
+        return (
+            self._rk4_step(
+                function,
+                time_start,
+                time_step,
+                time_end,
+                value_start,
+                function_value_start=function_value_start,
+            ),
+            function_value_start,
+        )
+
+    def _linear_interpolation(self, time_start, time_end, value_start, value_end, time_point):
+        if time_point == time_start:
+            return value_start
+        if time_point == time_end:
+            return value_end
+        weight = (time_point - time_start) / (time_end - time_start)
+        return value_start + weight * (value_end - value_start)
+
+    def integrate(self, time_points: torch.Tensor) -> torch.Tensor:
+        solution = torch.empty(
+            len(time_points),
+            *self.initial_value.shape,
+            dtype=self.initial_value.dtype,
+            device=self.initial_value.device,
+        )
+        solution[0] = self.initial_value
+
+        current_index = 1
+        current_value = self.initial_value
+        for time_start, time_end in zip(time_points[:-1], time_points[1:]):
+            time_step = time_end - time_start
+            delta_value, _ = self._compute_step(self.function, time_start, time_step, time_end, current_value)
+            next_value = current_value + delta_value
+
+            while current_index < len(time_points) and time_end >= time_points[current_index]:
+                solution[current_index] = self._linear_interpolation(
+                    time_start,
+                    time_end,
+                    current_value,
+                    next_value,
+                    time_points[current_index],
+                )
+                current_index += 1
+
+            current_value = next_value
+
+        return solution

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -60,6 +60,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
+      ode_solver_class: RungeKutta4ODESolver
       worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
       scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
       gpu_memory_utilization: 0.15

vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py

@@ -10,12 +10,7 @@ TALKER_CODEC_START_TOKEN_ID = 8293
 TALKER_CODEC_END_TOKEN_ID = 8294
 
 
-def thinker2talker(
-    stage_list,
-    engine_input_source,
-    prompt: Union[OmniTokensPrompt, TextPrompt] = None,
-    requires_multimodal_data: bool = False,
-):
+def thinker2talker(stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None):
     if not engine_input_source:
         raise ValueError("engine_input_source cannot be empty")
     source_stage_id = engine_input_source[0]
@@ -53,9 +48,7 @@ def thinker2talker(
                 + [TALKER_CODEC_END_TOKEN_ID],
                 additional_information=additional_information,
                 multi_modal_data=(
-                    multi_modal_data[thinker_output.request_id]
-                    if requires_multimodal_data and multi_modal_data is not None
-                    else None
+                    multi_modal_data[thinker_output.request_id] if multi_modal_data is not None else None
                 ),
                 mm_processor_kwargs=None,
             )

vllm_omni/worker/gpu_model_runner.py

@@ -6,7 +6,6 @@ from vllm.config import CUDAGraphMode
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces import supports_mrope
 from vllm.model_executor.models.interfaces_base import VllmModelForPooling
 from vllm.sampling_params import SamplingType
 from vllm.utils import LazyLoader, cdiv
@@ -30,42 +29,6 @@ logger = init_logger(__name__)
 
 
 class OmniGPUModelRunner(GPUModelRunner):
-    def _init_mrope_positions(self, req_state: CachedRequestState):
-        image_grid_thw = []
-        video_grid_thw = []
-        second_per_grid_ts = []
-        audio_feature_lengths = []
-        use_audio_in_video = False
-        for mm_feature in req_state.mm_features:
-            mm_item = mm_feature.data
-            if mm_item is None:
-                continue
-            mm_input = mm_item.get_data()
-            if (t := mm_input.get("image_grid_thw")) is not None:
-                image_grid_thw.append(t.tolist())
-            if (t := mm_input.get("video_grid_thw")) is not None:
-                video_grid_thw.append(t.tolist())
-            if (t := mm_input.get("second_per_grid_ts")) is not None:
-                second_per_grid_ts.append(t)
-            if (t := mm_input.get("audio_feature_lengths")) is not None:
-                audio_feature_lengths.append(t)
-            # Check for use_audio_in_video
-            use_audio_in_video_value = mm_input.get("use_audio_in_video")
-            if use_audio_in_video_value is not None:
-                use_audio_in_video = bool(use_audio_in_video_value.item())
-
-        assert supports_mrope(self.get_model()), "M-RoPE support is not implemented."
-
-        req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions(
-            req_state.prompt_token_ids,
-            hf_config=self.model_config.hf_config,
-            image_grid_thw=image_grid_thw,
-            video_grid_thw=video_grid_thw,
-            second_per_grid_ts=second_per_grid_ts,
-            audio_feature_lengths=audio_feature_lengths,
-            use_audio_in_video=use_audio_in_video,
-        )
-
     def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
         """Update the cached states and the persistent batch with the scheduler
         output.