update

Signed-off-by: Roger Wang <hey@rogerw.io>
test
2025-11-27 01:24:32 +00:00 · 2025-11-27 00:59:38 +00:00 · 2025-11-27 00:27:17 +00:00 · 2025-11-27 00:24:54 +00:00 · 2025-11-26 16:20:00 -08:00 · 2025-11-27 00:14:41 +00:00
44 changed files with 2448 additions and 1833 deletions
@@ -0,0 +1,6 @@
 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
@@ -0,0 +1,33 @@
 #!/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
@@ -1,60 +0,0 @@
 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"
@@ -1,124 +0,0 @@
 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"
@@ -3,6 +3,8 @@ name: pre-commit
 on:
  pull_request:
    branches: [main]
  push:
    branches: [main]
  workflow_dispatch: # for manual trigger
 concurrency:
@@ -160,7 +160,7 @@ cython_debug/
 docker/
 # scripts
-scripts/
+/scripts/
 # tests
 tests/
@@ -1,3 +1,6 @@
 default_install_hook_types:
  - pre-commit
  - commit-msg
 repos:
  - repo: https://github.com/pre-commit/pre-commit-hooks
    rev: v4.6.0
@@ -30,3 +33,28 @@ 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
@@ -1,40 +1,65 @@
-# vLLM-omni: Multi-modal Extension for vLLM
+<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 is designed to extend vLLM capabilities to support multi-modality model inference and serving, particularly focusing on non-autoregressive structures and non-textual outputs.
+<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>
-## 🎯 Overview
+---
-Traditional vLLM systems are limited to text-based, autoregressive generation. vLLM-omni addresses this limitation by enabling support for:
+*Latest News* 🔥
- **Multi-modal Models**: Text, image, video, audio, and sensor data processing
+- [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.
 - **Non-autoregressive Architectures**: Diffusion Transformers (DiT) and other parallel generation models
 - **Heterogeneous Outputs**: Beyond traditional text generation to multimodal outputs
-## 🏗️ Architecture
+---
-vLLM-omni is built on a modular architecture that extends vLLM's core functionality:
+## About
 [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:
-## 🚀 Key Features
+- **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
-### Multi-Engine Support
+<p align="center">
  <picture>
    <img alt="vllm-omni" src="./docs/source/architecture/omni-modality model architecture.png" width=55%>
  </picture>
 </p>
- **Autoregressive Engine**: Traditional text generation with enhanced KV-caching
+vLLM-omni is fast with:
 - **Diffusion Engine**: Support for DiT models and iterative generation
 - **Hybrid Engine**: Combined AR+DiT processing pipelines
-### Modality Processing
+- 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
- **Text**: Advanced tokenization and embedding generation
+vLLM-omni is flexible and easy to use with:
 - **Image**: Vision encoder integration (CLIP, etc.)
 - **Audio**: Speech processing and audio embedding
-## 📋 Supported Models
+- 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
-### AR + Diffusion Transformer (DiT) Models
+vLLM seamlessly supports most popular open-source models on HuggingFace, including:
- Qwen-omni (Thinker-Talker-Codec structure)
+
- HunyunaImage 3.0 (Ongoing)
+- Multi-modality generation models (e.g. Qwen-image)
- Qwen-Image (Ongoing)
+- 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
 ```
 ## 🛠️ Installation
@@ -84,12 +109,6 @@ 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
@@ -98,10 +117,21 @@ cd vllm_omni
 uv pip install -e .
 ```
 ## Run examples (Qwen2.5-omni)
 Please check the folder of [examples](examples)
-## Further details
+## Contributing
-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).
+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.
@@ -32,3 +32,11 @@ 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
 ```
@@ -1,166 +1,146 @@
 # 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
-import os as _os_env_toggle
+from typing import NamedTuple
 import random
 import numpy as np
 import soundfile as sf
-import torch
+from vllm.assets.audio import AudioAsset
-from utils import make_omni_prompt
+from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset
 from vllm.multimodal.image import convert_image_mode
 from vllm.sampling_params import SamplingParams
 from vllm.utils import FlexibleArgumentParser
-from vllm_omni.entrypoints.omni_llm import OmniLLM
+from vllm_omni 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"
-def parse_args():
+class QueryResult(NamedTuple):
-    parser = argparse.ArgumentParser()
+    inputs: dict
-    parser.add_argument(
+    limit_mm_per_prompt: dict[str, int]
        "--model",
        required=True,
        help="Path to merged model directory (will be created if downloading).",
    )
    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)",
    )
    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",
    )
    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).",
    )
    parser.add_argument(
        "--output-dir",
        default="outputs",
        help="Output directory to save text and wav files together.",
    )
    parser.add_argument(
        "--thinker-hidden-states-dir",
        default="thinker_hidden_states",
        help="Path to thinker hidden states directory.",
    )
    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(
        "--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
-def main():
+# NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on
-    args = parse_args()
+# lower-end GPUs.
-    model_name = args.model
+# Unless specified, these settings have been tested to work on a single L4.
-    try:
+
-        # Preferred: load from txt file (one prompt per line)
+default_system = (
-        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
+    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
-            with open(args.txt_prompts, encoding="utf-8") as f:
+    "Group, capable of perceiving auditory and visual inputs, as well as "
-                lines = [ln.strip() for ln in f.readlines()]
+    "generating text and speech."
-            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}")
+def get_text_query(question: str = None) -> QueryResult:
-        raise
+    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"
    )
    return QueryResult(
        inputs={
            "prompt": prompt,
        },
        limit_mm_per_prompt={},
    )
 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"
    )
    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},
    )
 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"
    )
    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},
    )
 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"
    )
    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,
        },
    )
 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]()
    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,
@@ -205,39 +185,123 @@ def main():
        code2wav_sampling_params,
    ]
-    prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
+    if args.txt_prompts is None:
-    omni_outputs = omni_llm.generate(prompt, sampling_params_list)
+        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}")
-    # Determine output directory: prefer --output-dir; fallback to --output-wav
+    omni_outputs = omni_llm.generate(prompts, sampling_params_list)
-    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+
-    os.makedirs(output_dir, exist_ok=True)
+    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)
    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
-                # Save aligned text file per request
+                if args.save_results:
-                prompt_text = args.prompts[request_id]
+                    # Save aligned text file per request
-                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
+                    prompt_text = prompts[request_id]["prompt"]
-                lines = []
+                    out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
-                lines.append("Prompt:\n")
+                    lines = []
-                lines.append(str(prompt_text) + "\n")
+                    lines.append("Prompt:\n")
-                lines.append("vllm_text_output:\n")
+                    lines.append(str(prompt_text) + "\n")
-                lines.append(str(text_output).strip() + "\n")
+                    lines.append("vllm_text_output:\n")
-                try:
+                    lines.append(str(text_output).strip() + "\n")
-                    with open(out_txt, "w", encoding="utf-8") as f:
+                    try:
-                        f.writelines(lines)
+                        with open(out_txt, "w", encoding="utf-8") as f:
-                except Exception as e:
+                            f.writelines(lines)
-                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
+                    except Exception as e:
-                print(f"Request ID: {request_id}, Text saved to {out_txt}")
+                        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}")
        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"]
-                output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
+                if args.save_results:
-                sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
+                    output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
-                print(f"Request ID: {request_id}, Saved audio to {output_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()
 if __name__ == "__main__":
-    main()
+    args = parse_args()
    main(args)
@@ -1,367 +0,0 @@
 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
@@ -1,8 +1,3 @@
-python end2end.py --model Qwen/Qwen2.5-Omni-7B \
+python end2end.py --output-wav output_audio \
-                                 --voice-type "m02" \
+                  --query-type text \
-                                 --dit-ckpt none \
+                  --txt-prompts top10.txt
                                 --bigvgan-ckpt none \
                                 --output-wav output_audio \
                                 --prompt_type text \
                                 --init-sleep-seconds 0 \
                                 --txt-prompts top100.txt
@@ -1,8 +1,2 @@
-python end2end.py --model Qwen/Qwen2.5-Omni-7B \
+python end2end.py --output-wav output_audio \
-                                 --voice-type "m02" \
+                  --query-type use_audio_in_video
                                 --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."
@@ -1,312 +0,0 @@
 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
@@ -23,10 +23,55 @@ cd examples/online_serving
 Send request via python
 ```bash
-python openai_chat_completion_client_for_multimodal_generation.py
+python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
 ```
 Send request via curl
 ```bash
-bash run_curl_multimodal_generation.sh
+bash run_curl_multimodal_generation.sh mixed_modalities
 ```
 ### 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)
@@ -0,0 +1,535 @@
 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()
@@ -1,6 +1,9 @@
 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"
@@ -15,6 +18,16 @@ 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",
@@ -31,7 +44,106 @@ def get_system_prompt():
    }
-def run_text_to_audio(model: str) -> None:
+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"
    thinker_sampling_params = {
        "temperature": 0.0,  # Deterministic - no randomness
        "top_p": 1.0,  # Disable nucleus sampling
@@ -67,23 +179,21 @@ def run_text_to_audio(model: str) -> 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,
+        model=model_name,
-        extra_body={
+        extra_body=extra_body,
            "sampling_params_list": sampling_params_list
        },  # Optional, it has a default setting in stage_configs of the corresponding model.
    )
    count = 0
@@ -99,5 +209,20 @@ def run_text_to_audio(model: str) -> 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__":
-    run_text_to_audio("Qwen/Qwen2.5-Omni-7B")
+    args = parse_args()
    run_multimodal_generation(args)
@@ -1,6 +1,20 @@
 #!/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='{
@@ -35,18 +49,121 @@ 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" \
+    -H "Content-Type: application/json" \
-  -d @- <<EOF
+    -d @- <<EOF
 {
  "model": "Qwen/Qwen2.5-Omni-7B",
-  "extra_body": {
+  "sampling_params_list": $sampling_params_list,
-    "sampling_params_list": [
+  "mm_processor_kwargs": $mm_processor_kwargs,
      $thinker_sampling_params,
      $talker_sampling_params,
      $code2wav_sampling_params
    ]
  },
  "messages": [
    {
      "role": "system",
@@ -59,17 +176,12 @@ output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
    },
    {
      "role": "user",
-      "content": [
+      "content": $user_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')"
@@ -30,10 +30,10 @@ classifiers = [
 dependencies = [
    # Core runtime dependencies (required for actual usage)
    "omegaconf>=2.3.0",
-    "decord>=0.6.0",
+    # "decord>=0.6.0", # install manually
    "librosa>=0.11.0",
    "resampy>=0.4.3",
-    # "vllm==0.11.0", # install manually
+    # "vllm==0.11.0",  # TODO: fix the entrypoints overwrite problem
 ]
 [project.optional-dependencies]
@@ -46,6 +46,8 @@ dev = [
    "isort>=5.12.0",
    "mypy>=1.0.0",
    "pre-commit>=3.0.0",
    "gradio>=5.49.1",
    "soundfile>=0.13.1"
 ]
 docs = [
@@ -73,10 +75,14 @@ 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 = [
@@ -1,108 +0,0 @@
 # 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
@@ -0,0 +1,139 @@
 #!/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
@@ -1,363 +0,0 @@
 #!/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 "$@"
@@ -15,7 +15,6 @@ 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
@@ -12,7 +12,6 @@ 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:
@@ -28,7 +27,6 @@ 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)
@@ -0,0 +1,242 @@
 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)]
@@ -50,6 +50,7 @@ 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
@@ -205,14 +206,20 @@ 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 multi_modal_data else None),
+                    multi_modal_data=(
                        multi_modal_data[source_output.request_id]
                        if self.requires_multimodal_data and 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)
+            return self.custom_process_input_func(
                stage_list, engine_input_source, prompt, self.requires_multimodal_data
            )
 def _stage_worker(
@@ -227,13 +234,13 @@ def _stage_worker(
    import logging as _logging
    import time as _time
-    from vllm_omni.entrypoints.log_utils import (  # noqa: WPS433
+    from vllm_omni.entrypoints.log_utils import (
        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  # noqa: WPS433
+    from vllm_omni.entrypoints.omni_llm import OmniStageLLM
    # no inline JSONL/serialization imports; logging handled by utilities
@@ -501,8 +508,8 @@ async def _stage_worker_async(
    import logging as _logging
    import time as _time
-    from vllm_omni.entrypoints.async_omni_llm import AsyncOmniStageLLM  # noqa: WPS433
+    from vllm_omni.entrypoints.async_omni_llm import AsyncOmniStageLLM
-    from vllm_omni.entrypoints.log_utils import (  # noqa: WPS433
+    from vllm_omni.entrypoints.log_utils import (
        compute_and_log_stage_request_stats,
        count_tokens_from_outputs,
        log_stage_batch_stats,
@@ -3,10 +3,10 @@ import base64
 import json
 import time
 import uuid
-from collections.abc import AsyncGenerator, AsyncIterator
+from collections.abc import AsyncGenerator, AsyncIterator, Sequence
 from datetime import datetime, timedelta, timezone
 from io import BytesIO
-from typing import Optional, Union
+from typing import Any, Callable, Optional, Union
 import jinja2
 from fastapi import Request
@@ -18,7 +18,16 @@ except ImportError:
    soundfile = None
 from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio
-from vllm.entrypoints.chat_utils import ConversationMessage, get_history_tool_calls_cnt, make_tool_call_id
+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.harmony_utils import parse_chat_output
 from vllm.entrypoints.openai.protocol import (
    ChatCompletionNamedToolChoiceParam,
@@ -35,7 +44,16 @@ from vllm.entrypoints.openai.protocol import (
    UsageInfo,
 )
 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
-from vllm.entrypoints.openai.serving_engine import RequestPrompt, clamp_prompt_logprobs
+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.tool_parsers.mistral_tool_parser import MistralToolCall
 from vllm.inputs.data import PromptType
 from vllm.logger import init_logger
@@ -50,6 +68,7 @@ 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__)
@@ -215,6 +234,123 @@ 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:
@@ -7,6 +7,7 @@ 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,
@@ -14,13 +15,7 @@ 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 SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.interfaces import SupportsMRoPE, 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
@@ -31,7 +26,14 @@ 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.utils import add_prefix_to_loaded_weights
+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
 TALKER_CODEC_EOS_TOKEN_ID = 8294
 TALKER_CODEC_BOS_TOKEN_ID = 8293
@@ -54,7 +56,7 @@ logger = init_logger(__name__)
    dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
 class Qwen2_5OmniForConditionalGeneration(
-    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin, SupportsMRoPE
 ):
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
@@ -459,6 +461,170 @@ 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):
@@ -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 Any, Optional, Union
+from typing import Annotated, Any, Callable, Literal, 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, SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMRoPE, SupportsMultiModal, SupportsPP
 from vllm.model_executor.models.qwen2_5_vl import (
    Qwen2_5_VisionTransformer,
    Qwen2_5_VLImageEmbeddingInputs,
@@ -24,7 +24,6 @@ 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,
 )
@@ -42,7 +41,7 @@ from vllm.multimodal.inputs import (
    ModalityData,
    MultiModalDataDict,
    MultiModalFieldConfig,
-    MultiModalKwargs,
+    MultiModalKwargsItems,
    NestedTensors,
 )
 from vllm.multimodal.parse import (
@@ -52,10 +51,19 @@ from vllm.multimodal.parse import (
    MultiModalDataItems,
    MultiModalDataParser,
 )
-from vllm.multimodal.processing import BaseMultiModalProcessor, PlaceholderFeaturesInfo, PromptReplacement, PromptUpdate
+from vllm.multimodal.processing import (
    BaseMultiModalProcessor,
    MultiModalPromptUpdates,
    PlaceholderFeaturesInfo,
    PromptReplacement,
    PromptUpdate,
    PromptUpdateDetails,
    ResolvedPromptUpdate,
 )
 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
@@ -67,6 +75,27 @@ 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,)))
@@ -90,6 +119,39 @@ 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,
@@ -166,10 +228,7 @@ 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,
+                width=target_width, height=target_height, num_frames=target_num_frames, num_videos=num_videos
                height=target_height,
                num_frames=target_num_frames,
                num_videos=num_videos,
            ),
        }
@@ -179,7 +238,10 @@ 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(target_sr=feature_extractor.sampling_rate)
+        return Qwen2_5OmniThinkerMultiModalDataParser(
            spatial_merge_size=self.info.get_hf_config().vision_config.spatial_merge_size,
            target_sr=feature_extractor.sampling_rate,
        )
    def _call_hf_processor(
        self,
@@ -214,6 +276,14 @@ 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(
@@ -221,66 +291,177 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
-        return _qwen2_5_omni_thinker_field_config(hf_inputs)
+        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
    def _maybe_apply_prompt_updates(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
        prompt_ids: list[int],
-        mm_kwargs: MultiModalKwargs,
+        mm_kwargs: MultiModalKwargsItems,
        mm_prompt_updates: MultiModalPromptUpdates,
        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 = hf_processor_mm_kwargs.get("use_audio_in_video", False)
+        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)
        if is_update_applied:
            mm_placeholders = self._find_mm_placeholders(
                prompt_ids,
                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)
+            self._validate_mm_placeholders(
                mm_placeholders,
                mm_item_counts,
            )
            tokenizer = self.info.get_tokenizer()
            prompt = decode_tokens(tokenizer, prompt_ids)
        else:
-            (
+            # When use_audio_in_video=True, audio tokens are not in the prompt,
-                prompt_ids,
+            # so we need to filter out audio updates before applying replacements
-                prompt,
+            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,
            ) = self._apply_prompt_updates(
                prompt_ids,
                mm_prompt_updates,
                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: MultiModalKwargs,
+        out_mm_kwargs: MultiModalKwargsItems,
    ) -> Sequence[PromptUpdate]:
        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
        tokenizer = self.info.get_tokenizer()
@@ -294,8 +475,9 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
        image_token_id = vocab[image_token]
        video_token_id = vocab[video_token]
-        audio_feature_lengths = out_mm_kwargs.get("audio_feature_lengths")
+        out_mm_data = out_mm_kwargs.get_data()
-        feature_attention_mask = out_mm_kwargs.get("feature_attention_mask")
+        audio_feature_lengths = out_mm_data.get("audio_feature_lengths")
        feature_attention_mask = out_mm_data.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:
@@ -323,7 +505,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_kwargs[f"{modality}_grid_thw"][item_idx]
+            grid_thw = out_mm_data[f"{modality}_grid_thw"][item_idx]
            assert isinstance(grid_thw, torch.Tensor)
            merge_length = image_processor.merge_size**2
@@ -337,7 +519,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_kwargs["video_grid_thw"][item_idx]
+            video_grid_thw = out_mm_data["video_grid_thw"][item_idx]
            audio_in_video_item_idx += 1
@@ -347,13 +529,15 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
            else:
                video_second_per_grid_t = 1.0
-            return MRotaryEmbedding.omni_get_updates_use_audio_in_video(
+            placeholder = 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
@@ -378,6 +562,50 @@ 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]],
@@ -436,30 +664,19 @@ 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[Qwen2AudioInputs]:
+    def _parse_and_validate_audio_input(self, **kwargs: object) -> Optional[Qwen2_5OmniAudioFeatureInputs]:
        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)
@@ -472,7 +689,8 @@ 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 Qwen2AudioInputs(
+        return Qwen2_5OmniAudioFeatureInputs(
            type="audio_features",
            input_features=input_audio_features,
            audio_feature_lengths=audio_feature_lengths,
            feature_attention_mask=feature_attention_mask,
@@ -497,9 +715,7 @@ class Qwen2_5OmniConditionalGenerationMixin:
                raise ValueError(f"Incorrect type of image pixel values. Got type: {type(pixel_values)}")
            return Qwen2_5_VLImagePixelInputs(
-                type="pixel_values",
+                type="pixel_values", pixel_values=pixel_values, image_grid_thw=image_grid_thw
                pixel_values=pixel_values,
                image_grid_thw=image_grid_thw,
            )
        if image_embeds is not None:
@@ -509,9 +725,7 @@ 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",
+                type="image_embeds", image_embeds=image_embeds, image_grid_thw=image_grid_thw
                image_embeds=image_embeds,
                image_grid_thw=image_grid_thw,
            )
    def _parse_and_validate_video_input(
@@ -542,14 +756,12 @@ 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",
+                type="video_embeds", video_embeds=video_embeds, video_grid_thw=video_grid_thw
                video_embeds=video_embeds,
                video_grid_thw=video_grid_thw,
            )
    def _process_audio_input(
        self,
-        audio_input: Qwen2AudioInputs,
+        audio_input: Qwen2_5OmniAudioFeatureInputs,
        audio_hashes: list[str] = None,
        cached_audio_features: torch.Tensor = None,
    ) -> torch.Tensor:
@@ -574,8 +786,7 @@ class Qwen2_5OmniConditionalGenerationMixin:
            feature_lens=audio_feature_lengths,
            aftercnn_lens=audio_feat_lengths,
        )
-        audio_features = audio_outputs.last_hidden_state
+        return audio_outputs.last_hidden_state.split(audio_output_lengths.tolist())
        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":
@@ -619,7 +830,7 @@ class Qwen2_5OmniConditionalGenerationMixin:
    dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
 class Qwen2_5OmniThinkerForConditionalGeneration(
-    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin, SupportsMRoPE
 ):
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
@@ -660,13 +871,21 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
                "exactly same result as the transformers implementation in the audio tower part."
            )
-        self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
+        if multimodal_config.get_limit_per_prompt("audio"):
-        self.visual = Qwen2_5_VisionTransformer(
+            self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
-            vision_config=thinker_config.vision_config,
+        else:
-            norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
+            self.audio_tower = None
-            quant_config=quant_config,
+
-            prefix=maybe_prefix(prefix, "visual"),
+        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.quant_config = quant_config
        self.language_model = init_vllm_registered_model(
            vllm_config=vllm_config,
@@ -23,7 +23,6 @@ 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
@@ -34,8 +33,6 @@ 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):
@@ -572,8 +569,7 @@ class DiTAttention(nn.Module):
        cos, sin = position_embeddings
        query[:, :1], key[:, :1] = apply_rotary_pos_emb(query[:, :1], key[:, :1], cos, sin)
-        attn_impl = self.config._attn_implementation or "sdpa"
+        attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
        attention_interface = ALL_ATTENTION_FUNCTIONS[attn_impl]
        attention_weights, _ = attention_interface(
            self,
            query,
@@ -1018,19 +1014,100 @@ 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(nn.Module):
+class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+    config: Qwen2_5OmniDiTConfig
-        super().__init__()
+    _no_split_modules = ["DiTDecoderLayer"]
        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)
@@ -1058,10 +1135,6 @@ class Qwen2_5OmniToken2WavDiTModel(nn.Module):
        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]
@@ -1175,8 +1248,8 @@ class Qwen2_5OmniToken2WavDiTModel(nn.Module):
        if sway_coefficient is not None:
            time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
-        self.ode_solver.reset_ode_solver(function=ode_function, initial_value=initial_state)
+        ode_solver = RungeKutta4ODESolver(function=ode_function, initial_value=initial_state)
-        solution_trajectory = self.ode_solver.integrate(time_embedding)
+        solution_trajectory = ode_solver.integrate(time_embedding)
        generated_waveform = solution_trajectory[-1]
        generated_mel_spectrogram = generated_waveform.permute(0, 2, 1)
@@ -1247,8 +1320,8 @@ class Qwen2_5OmniToken2WavDiTModel(nn.Module):
        if sway_coefficient is not None:
            time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
-        self.ode_solver.reset_ode_solver(function=ode_function, initial_value=initial_state)
+        ode_solver = RungeKutta4ODESolver(function=ode_function, initial_value=initial_state)
-        solution_trajectory = self.ode_solver.integrate(time_embedding)
+        solution_trajectory = ode_solver.integrate(time_embedding)
        generated_waveform = solution_trajectory[-1]
        generated_mel_spectrogram = generated_waveform.permute(0, 2, 1)
@@ -1262,13 +1335,16 @@ class Qwen2_5OmniToken2WavDiTModel(nn.Module):
    mel spectrogram as input and predict waveform.
    """
 )
-class Qwen2_5OmniToken2WavModel(nn.Module):
+class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
-    def __init__(self, *, vllm_config: VllmConfig, prefix=""):
+    config: Qwen2_5OmniToken2WavConfig
-        super().__init__()
+    base_model_prefix = "model"
-
+    _no_split_modules = [
-        self.model_config = vllm_config.model_config
+        "Qwen2_5OmniToken2WavDiTModel",
-        config: Qwen2_5OmniToken2WavConfig = self.model_config.hf_config
+        "Qwen2_5OmniToken2WavBigVGANModel",
    ]
    def __init__(self, config: Qwen2_5OmniToken2WavConfig):
        super().__init__(config)
        attn_impl = config._attn_implementation
        if config._attn_implementation == "flash_attention_2":
            logger.warning_once(
@@ -1283,10 +1359,9 @@ class Qwen2_5OmniToken2WavModel(nn.Module):
                "Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa"
            )
            attn_impl = "sdpa"
-        with set_default_torch_dtype(torch.float32):
+        self.code2wav_dit_model = Qwen2_5OmniToken2WavDiTModel._from_config(
-            self.code2wav_dit_model = Qwen2_5OmniToken2WavDiTModel(
+            config.dit_config, attn_implementation=attn_impl
-                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
        )
@@ -1319,10 +1394,6 @@ class Qwen2_5OmniToken2WavModel(nn.Module):
            # fallback to commonly used value
            self.vocoder_hop = 240
    @property
    def dtype(self):
        return next(self.parameters()).dtype
    def forward(
        self,
        code,
@@ -2,26 +2,31 @@ 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_old", "Qwen2ForCausalLM"),  # need to discuss
+    "Qwen2ForCausalLM_old": ("qwen2_5_omni", "qwen2_old", "Qwen2ForCausalLM"),  # need to discuss
 }
 _VLLM_OMNI_MODELS = {
@@ -41,10 +46,10 @@ OmniModelRegistry = _ModelRegistry(
        },
        **{
            model_arch: _LazyRegisteredModel(
-                module_name=f"vllm_omni.model_executor.models.{mod_relname}",
+                module_name=f"vllm_omni.model_executor.models.{mod_folder}.{mod_relname}",
                class_name=cls_name,
            )
-            for model_arch, (mod_relname, cls_name) in _OMNI_MODELS.items()
+            for model_arch, (mod_folder, mod_relname, cls_name) in _OMNI_MODELS.items()
        },
    }
 )
@@ -1,3 +1,4 @@
 import torch
 from vllm.model_executor.models.utils import maybe_prefix
@@ -6,3 +7,24 @@ 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
@@ -0,0 +1,23 @@
 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
@@ -1,15 +0,0 @@
 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"]
@@ -1,89 +0,0 @@
 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
@@ -60,7 +60,6 @@ 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
@@ -10,7 +10,12 @@ TALKER_CODEC_START_TOKEN_ID = 8293
 TALKER_CODEC_END_TOKEN_ID = 8294
-def thinker2talker(stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None):
+def thinker2talker(
    stage_list,
    engine_input_source,
    prompt: Union[OmniTokensPrompt, TextPrompt] = None,
    requires_multimodal_data: bool = False,
 ):
    if not engine_input_source:
        raise ValueError("engine_input_source cannot be empty")
    source_stage_id = engine_input_source[0]
@@ -48,7 +53,9 @@ def thinker2talker(stage_list, engine_input_source, prompt: Union[OmniTokensProm
                + [TALKER_CODEC_END_TOKEN_ID],
                additional_information=additional_information,
                multi_modal_data=(
-                    multi_modal_data[thinker_output.request_id] if multi_modal_data is not None else None
+                    multi_modal_data[thinker_output.request_id]
                    if requires_multimodal_data and multi_modal_data is not None
                    else None
                ),
                mm_processor_kwargs=None,
            )
@@ -6,6 +6,7 @@ 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
@@ -29,6 +30,42 @@ 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.
Author	SHA1	Message	Date
Roger Wang	27c6917149	update Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-27 01:24:32 +00:00
Roger Wang	a2dd5feb80	test Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-27 00:59:38 +00:00
Roger Wang	779aa534ba	add back Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-27 00:27:17 +00:00
Roger Wang	effe5b347f	Merge branch 'main' into buildkite-2	2025-11-27 00:24:54 +00:00
Roger Wang	99c46e1adc	[CI] Buildkite setup (#83 ) Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-26 16:20:00 -08:00
Roger Wang	42426e7c10	update Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-27 00:14:41 +00:00
Roger Wang	1a6879b8a4	dummify Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-27 00:10:27 +00:00
Roger Wang	3a3aebde0c	test Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-27 00:05:04 +00:00
Samit	b2d845742a	[Feature] Add Gradio Demo for Qwen2.5Omni (#60 ) Signed-off-by: SamitHuang <285365963@qq.com> Signed-off-by: samithuang <285365963@qq.com>	2025-11-26 09:23:46 +08:00
Gao Han	598b4ddb51	[Feature] support multimodal inputs with multiple requests (#76 ) Signed-off-by: Gaohan123 <gaohan19@huawei.com>	2025-11-24 16:02:44 +08:00
Alicia	08faf16f81	[CI] Add script for building wheel. (#75 ) Signed-off-by: WANG Cong <115451386+congw729@users.noreply.github.com>	2025-11-24 14:57:05 +08:00
Hongsheng LiuandRoger Wang	1fb46ddd62	update readme for v0.11.0rc1 release (#69 ) Signed-off-by: hsliu <liuhongsheng4@huawei.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Roger Wang <hey@rogerw.io>	2025-11-24 14:28:25 +08:00
Jiangyun Zhu	9ccff6c710	[CI] disable unnecessary ci and update pre-commit (#80 ) Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2025-11-24 18:14:08 +13:00
Roger Wang	b5bb75a032	update Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-23 21:09:58 -08:00
Roger Wang	0e48b214da	update Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-23 21:03:46 -08:00
Roger Wang	dc45f1b5cd	update Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-23 20:59:30 -08:00
Roger Wang	c0a264452c	update Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-23 19:55:56 -08:00
Roger Wang	35e063e102	update Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-23 19:24:02 -08:00
Roger Wang	89b1bdbb0e	add Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-23 19:12:58 -08:00