[cpu] remove verbose print

Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com>

.buildkite/pipeline.yml

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

.buildkite/scripts/run_qwen_offline_example.sh

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

.github/CODEOWNERS

@@ -1,17 +0,0 @@
-# CODEOWNERS is used to automatically assign PR reviewers
-# Format: path pattern @username or @team
-# More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
-
-# Default owners (fallback reviewers for all files)
-# later, we can add more reviewers here
-* @hsliuustc0106
-
-# Assign reviewers by directory
-/.github/ @owner1
-/tests/ @owner1
-/docs/ @owner1
-
-# Example: reviewers for specific directories
-# /examples/offline_inference/ @offline-inference-team
-# /vllm_omni/core/ @core-team
-# /vllm_omni/model_executor/ @model-team

.github/ISSUE_TEMPLATE/200-installation.yml

@@ -7,14 +7,14 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
+      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/hsliuustc0106/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
 - type: textarea
   attributes:
     label: Your current environment
     description: |
       Please run the following and paste the output below.
       ```sh
-      wget https://raw.githubusercontent.com/vllm-project/vllm-omni/main/vllm_omni/collect_env.py
+      wget https://raw.githubusercontent.com/hsliuustc0106/vllm-omni/main/vllm_omni/collect_env.py
       # For security purposes, please feel free to check the contents of collect_env.py before running it.
       python collect_env.py
       ```

.github/ISSUE_TEMPLATE/400-bug-report.yml

@@ -7,7 +7,7 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
+      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/hsliuustc0106/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
 - type: markdown
   attributes:
     value: |
@@ -24,7 +24,7 @@ body:
     description: |
       Please run the following and paste the output below.
       ```sh
-      wget https://raw.githubusercontent.com/vllm-project/vllm-omni/main/vllm_omni/collect_env.py
+      wget https://raw.githubusercontent.com/hsliuustc0106/vllm-omni/main/vllm_omni/collect_env.py
       # For security purposes, please feel free to check the contents of collect_env.py before running it.
       python collect_env.py
       ```

.github/ISSUE_TEMPLATE/500-feature-request.yml

@@ -7,7 +7,7 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
+      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/hsliuustc0106/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
 - type: textarea
   attributes:
     label: 🚀 The feature, motivation and pitch

.github/ISSUE_TEMPLATE/600-new-model.yml

@@ -7,7 +7,7 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
+      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/hsliuustc0106/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
 
       #### We also highly recommend you read the contributing guidelines first to understand how to add a new model.
 - type: textarea

.github/ISSUE_TEMPLATE/700-performance-discussion.yml

@@ -7,7 +7,7 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
+      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/hsliuustc0106/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
 - type: textarea
   attributes:
     label: Proposal to improve performance
@@ -19,7 +19,7 @@ body:
   attributes:
     label: Report of performance regression
     description: >
-      Please provide detailed description of performance comparison to confirm the regression. You may want to run the benchmark script at https://github.com/vllm-project/vllm-omni/tree/main/tests/benchmarks .
+      Please provide detailed description of performance comparison to confirm the regression. You may want to run the benchmark script at https://github.com/hsliuustc0106/vllm-omni/tree/main/tests/benchmarks .
   validations:
     required: false
 - type: textarea
@@ -35,7 +35,7 @@ body:
     description: |
       Please run the following and paste the output below.
       ```sh
-      wget https://raw.githubusercontent.com/vllm-project/vllm-omni/main/vllm_omni/collect_env.py
+      wget https://raw.githubusercontent.com/hsliuustc0106/vllm-omni/main/vllm_omni/collect_env.py
       # For security purposes, please feel free to check the contents of collect_env.py before running it.
       python collect_env.py
       ```

.github/ISSUE_TEMPLATE/750-RFC.yml

@@ -7,7 +7,7 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Please take a look at previous [RFCs](https://github.com/vllm-project/vllm-omni/issues?q=label%3ARFC+sort%3Aupdated-desc) for reference.
+      #### Please take a look at previous [RFCs](https://github.com/hsliuustc0106/vllm-omni/issues?q=label%3ARFC+sort%3Aupdated-desc) for reference.
 - type: textarea
   attributes:
     label: Motivation.

.github/ISSUE_TEMPLATE/config.yml

@@ -1,5 +1,5 @@
 blank_issues_enabled: false
 contact_links:
   - name: Questions
-    url: https://github.com/vllm-project/vllm-omni/discussions
+    url: https://github.com/hsliuustc0106/vllm-omni/discussions
     about: Ask questions and discuss with other vLLM-omni community members

.github/PULL_REQUEST_TEMPLATE.md

@@ -18,4 +18,4 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
 - [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft.
 </details>
 
-**BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)
+**BEFORE SUBMITTING, PLEASE READ <https://github.com/hsliuustc0106/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)

.github/scripts/detect_changed_tests.sh

@@ -1,89 +0,0 @@
-#!/bin/bash
-# Script to detect changed files and find corresponding test files
-# Usage: ./detect_changed_tests.sh [base_sha] [head_sha] [event_type]
-#
-# Outputs to GITHUB_OUTPUT:
-#   test_files: Space-separated list of test files to run
-#   has_tests: 'true' if test files found, 'false' otherwise
-
-set -e
-
-BASE_SHA="${1:-}"
-HEAD_SHA="${2:-}"
-EVENT_TYPE="${3:-push}"
-
-# Determine base and head SHA based on event type
-if [ "$EVENT_TYPE" = "pull_request" ]; then
-    BASE_SHA="${BASE_SHA:-$GITHUB_BASE_REF}"
-    HEAD_SHA="${HEAD_SHA:-$GITHUB_SHA}"
-elif [ -z "$BASE_SHA" ] || [ -z "$HEAD_SHA" ]; then
-    # For push events, use git to find the previous commit
-    BASE_SHA="${BASE_SHA:-$(git rev-parse HEAD~1 2>/dev/null || echo HEAD)}"
-    HEAD_SHA="${HEAD_SHA:-HEAD}"
-fi
-
-echo "Event type: $EVENT_TYPE"
-echo "Base SHA: $BASE_SHA"
-echo "Head SHA: $HEAD_SHA"
-
-# Get changed Python files
-CHANGED_FILES=$(git diff --name-only --diff-filter=ACMRT "$BASE_SHA" "$HEAD_SHA" | grep -E '\.(py)$' || true)
-
-if [ -z "$CHANGED_FILES" ]; then
-    echo "No Python files changed"
-    echo "test_files=" >> "$GITHUB_OUTPUT"
-    echo "has_tests=false" >> "$GITHUB_OUTPUT"
-    exit 0
-fi
-
-echo "Changed Python files:"
-echo "$CHANGED_FILES"
-
-# Extract related test files
-TEST_FILES=""
-for file in $CHANGED_FILES; do
-    # If changed file is a test file, add it directly
-    if [[ "$file" == tests/* ]]; then
-        TEST_FILES="$TEST_FILES $file"
-    # If changed file is source code, find corresponding test file
-    elif [[ "$file" == vllm_omni/* ]]; then
-        # Convert vllm_omni/path/to/module.py to possible test file paths
-        REL_PATH=$(echo "$file" | sed 's|^vllm_omni/||' | sed 's|\.py$||')
-        MODULE_NAME=$(basename "$REL_PATH")
-        DIR_PATH=$(dirname "$REL_PATH")
-
-        # Try multiple possible test file naming patterns
-        POSSIBLE_TESTS=(
-            "tests/test_${REL_PATH//\//_}.py"
-            "tests/${REL_PATH//\//_}_test.py"
-        )
-
-        # If multi-level directory, also try directory structure
-        if [ "$DIR_PATH" != "." ]; then
-            POSSIBLE_TESTS+=(
-                "tests/${DIR_PATH}/test_${MODULE_NAME}.py"
-                "tests/${DIR_PATH}/${MODULE_NAME}_test.py"
-            )
-        fi
-
-        for test_file in "${POSSIBLE_TESTS[@]}"; do
-            if [ -f "$test_file" ]; then
-                TEST_FILES="$TEST_FILES $test_file"
-                break
-            fi
-        done
-    fi
-done
-
-# Deduplicate and format
-TEST_FILES=$(echo "$TEST_FILES" | tr ' ' '\n' | grep -v '^$' | sort -u | tr '\n' ' ' | sed 's/^ *//;s/ *$//')
-
-if [ -z "$TEST_FILES" ]; then
-    echo "No related test files found, will run all tests"
-    echo "test_files=" >> "$GITHUB_OUTPUT"
-    echo "has_tests=false" >> "$GITHUB_OUTPUT"
-else
-    echo "Found test files: $TEST_FILES"
-    echo "test_files=$TEST_FILES" >> "$GITHUB_OUTPUT"
-    echo "has_tests=true" >> "$GITHUB_OUTPUT"
-fi

.github/workflows/pre-commit.yml

@@ -1,26 +0,0 @@
-name: pre-commit
-
-on:
-  pull_request:
-    branches: [main]
-  push:
-    branches: [main]
-  workflow_dispatch: # for manual trigger
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
-
-permissions:
-  contents: read
-
-jobs:
-  pre-commit:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-      - uses: actions/setup-python@v5
-        with:
-          python-version: "3.12"
-          cache: 'pip'
-      - uses: pre-commit/action@v3.0.1

.gitignore

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

.pre-commit-config.yaml

@@ -1,6 +1,3 @@
-default_install_hook_types:
-  - pre-commit
-  - commit-msg
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.6.0
@@ -15,46 +12,18 @@ repos:
       - id: trailing-whitespace
         args: ["--markdown-linebreak-ext=md"]
 
-  - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.14.0
+  - repo: https://github.com/PyCQA/isort
+    rev: 5.13.2
     hooks:
-    - id: ruff-check
-      args: [--output-format, github, --fix]
-    - id: ruff-format
+      - id: isort
 
-  - repo: https://github.com/crate-ci/typos
-    rev: v1.35.5
+  - repo: https://github.com/psf/black
+    rev: 24.10.0
     hooks:
-      - id: typos
-        # only for staged files
+      - id: black
 
-  - repo: https://github.com/rhysd/actionlint
-    rev: v1.7.7
+  - repo: https://github.com/PyCQA/flake8
+    rev: 7.1.1
     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
+      - id: flake8
+        args: ["--max-line-length=88", "--extend-ignore=E203,W503"]

.readthedocs.yml

@@ -1,24 +0,0 @@
-# Read the Docs configuration file
-# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
-
-version: 2
-
-build:
-  os: ubuntu-22.04
-  tools:
-    python: "3.12"
-  jobs:
-    post_checkout:
-      - git fetch --unshallow || true
-
-mkdocs:
-  configuration: mkdocs.yml
-  fail_on_warning: true
-
-# Optionally declare the Python requirements required to build your docs
-python:
-  install:
-    - method: pip
-      path: .
-      extra_requirements:
-        - docs

README.md

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

docs/DOCS_GUIDE.md

@@ -1,139 +0,0 @@
-# Documentation Build Guide
-
-This directory contains the source files for the vLLM-omni documentation.
-
-## Building Documentation Locally
-
-### Prerequisites
-
-Install documentation dependencies:
-
-```bash
-pip install -e ".[docs]"
-```
-
-### Build and Serve Documentation
-
-From the project root:
-
-```bash
-# Serve documentation locally (auto-reload on changes)
-# This starts a local web server at http://127.0.0.1:8000
-mkdocs serve
-
-# Build static site (generates HTML files in site/ directory)
-mkdocs build
-```
-
-When using `mkdocs serve`, the documentation will be automatically available at `http://127.0.0.1:8000`. The server will automatically reload when you make changes to the documentation files.
-
-## Auto-generating API Documentation
-
-The documentation automatically extracts docstrings from the code using mkdocstrings. To ensure your code is documented:
-
-1. Add docstrings to all public classes, functions, and methods
-2. Use Google or NumPy style docstrings (both are supported)
-3. Rebuild the documentation to see changes
-
-Example docstring:
-
-```python
-class OmniLLM:
-    """Main entry point for vLLM-omni inference.
-
-    This class provides a high-level interface for running multi-modal
-    inference with non-autoregressive models.
-
-    Args:
-        model: Model name or path
-        stage_configs: Optional stage configurations
-        **kwargs: Additional arguments passed to the engine
-
-    Example:
-        >>> llm = OmniLLM(model="Qwen/Qwen2.5-Omni")
-        >>> outputs = llm.generate(prompts="Hello")
-    """
-```
-
-## Documentation Structure
-
-```
-docs/
-├── index.md              # Main documentation page
-├── getting_started/      # Getting started guides
-├── architecture/        # Architecture documentation
-├── api/                 # API reference (auto-generated from code)
-├── examples/            # Code examples
-└── stylesheets/         # Custom CSS
-```
-
-## Publishing Documentation
-
-### GitHub Pages (Recommended)
-
-The documentation is automatically deployed to GitHub Pages using GitHub Actions.
-
-1. **Enable GitHub Pages**:
-   - Go to repository `Settings` → `Pages`
-   - Set `Source` to `GitHub Actions`
-   - Save settings
-
-2. **Push changes**:
-   ```bash
-   git push origin main
-   ```
-
-3. **Documentation will be available at**:
-   - `https://vllm-project.github.io/vllm-omni/`
-
-The GitHub Actions workflow (`.github/workflows/docs.yml`) will automatically:
-- Build the documentation when you push to `main` branch
-- Deploy it to GitHub Pages
-- Update the documentation whenever you make changes
-
-
-### Read the Docs (Alternative)
-
-You can also use Read the Docs for hosting:
-
-1. Sign up at https://readthedocs.org/
-2. Import the `vllm-project/vllm-omni` repository
-3. Read the Docs will automatically build using `.readthedocs.yml`
-4. Documentation will be available at: `https://vllm-omni.readthedocs.io/`
-
-## Configuration
-
-The documentation configuration is in `mkdocs.yml` at the project root.
-
-## Tips
-
-- **API Documentation**: API docs are automatically generated using `mkdocs-api-autonav` and `mkdocstrings`
-  - No need to manually create API pages - they're generated automatically
-  - Use `[module.name.ClassName][]` syntax for cross-references in Summary pages
-- **Code Snippets**: Use `--8<-- "path/to/file.py"` for including code snippets
-- **Markdown**: Use Markdown for all documentation (no need for RST)
-- **Material Theme**: Use Material theme features like:
-  - Admonitions: `!!! note`, `!!! warning`, etc.
-  - Code blocks with syntax highlighting
-  - Tabs for organizing content
-  - Math formulas using `pymdownx.arithmatex`
-
-## Troubleshooting
-
-### Documentation not updating
-
-- Make sure you've saved all files
-- If using `mkdocs serve`, it should auto-reload
-- Check for syntax errors in `mkdocs.yml`
-
-### API links not working
-
-- Ensure class names match exactly (case-sensitive)
-- Check that the module is imported correctly
-- Run `mkdocs build --strict` to check for errors
-
-### Build errors
-
-- Check Python version (requires 3.9+)
-- Ensure all dependencies are installed: `pip install -e ".[docs]"`
-- Check `mkdocs.yml` syntax with `mkdocs build --strict`

docs/api/README.md

@@ -1,10 +1,13 @@
-# API Design Documentation
+# vLLM-omni API Documentation
 
-This document contains comprehensive API design documentation for all core modules in vLLM-omni. These templates provide a standardized structure for designing and implementing the core, engine, executor, and worker modules.
+This directory contains comprehensive API design documentation for all core modules in vLLM-omni. These templates provide a standardized structure for designing and implementing the core, engine, executor, and worker modules.
 
-## 📋 Module API
+## 📋 API Design Templates
 
-### Core Module API
+### 1. [API Design Template](API_DESIGN_TEMPLATE.md)
+**Master template** for creating API documentation for any module in vLLM-omni. Use this as a starting point for new modules.
+
+### 2. [Core Module API](core_module_api.md)
 **Core module** provides fundamental scheduling, caching, and resource management functionality.
 
 **Key Components:**
@@ -13,7 +16,7 @@ This document contains comprehensive API design documentation for all core modul
 - Resource allocation and coordination
 - Inter-module communication
 
-### Engine Module API
+### 3. [Engine Module API](engine_module_api.md)
 **Engine module** handles model loading, inference execution, and output processing.
 
 **Key Components:**
@@ -22,7 +25,7 @@ This document contains comprehensive API design documentation for all core modul
 - Input preprocessing for various modalities
 - Output postprocessing and formatting
 
-### Executor Module API
+### 4. [Executor Module API](executor_module_api.md)
 **Executor module** coordinates and manages request execution across different engines and workers.
 
 **Key Components:**
@@ -31,7 +34,7 @@ This document contains comprehensive API design documentation for all core modul
 - Worker management and task distribution
 - Error handling and recovery
 
-### Worker Module API
+### 5. [Worker Module API](worker_module_api.md)
 **Worker module** provides the actual execution environment for model inference.
 
 **Key Components:**
@@ -40,6 +43,79 @@ This document contains comprehensive API design documentation for all core modul
 - Request batching and processing
 - Performance optimization
 
+## 🎯 How to Use These Templates
+
+### For New Module Development
+
+1. **Start with the Master Template**
+   - Copy `API_DESIGN_TEMPLATE.md`
+   - Rename to `[module_name]_api.md`
+   - Follow the structure exactly
+
+2. **Fill in Module-Specific Details**
+   - Update the module overview
+   - Define core classes and interfaces
+   - Specify public API methods
+   - Add configuration options
+   - Define error handling strategy
+   - Provide usage examples
+
+3. **Review and Validate**
+   - Ensure all sections are complete
+   - Verify code examples work
+   - Check for consistency with other modules
+   - Validate configuration options
+
+### For Existing Module Updates
+
+1. **Update API Documentation**
+   - Modify existing API files
+   - Add new methods and classes
+   - Update examples and configuration
+   - Maintain backward compatibility notes
+
+2. **Version Control**
+   - Track changes in git
+   - Use clear commit messages
+   - Tag major API changes
+
+## 📚 Template Structure
+
+Each API template follows this standardized structure:
+
+### 1. Module Overview
+- **Purpose**: What the module does
+- **Responsibilities**: Key responsibilities
+- **Dependencies**: Required modules
+- **Integration Points**: How it connects with other modules
+
+### 2. Core Classes/Interfaces
+- **Base Classes**: Abstract base classes
+- **Implementation Classes**: Concrete implementations
+- **Data Structures**: Key data models
+
+### 3. Public API Methods
+- **Initialization**: Constructor and setup
+- **Core Operations**: Main functionality
+- **Configuration**: Settings management
+- **Lifecycle Management**: Start/stop/cleanup
+- **Monitoring**: Status and metrics
+
+### 4. Configuration
+- **Configuration Classes**: Dataclasses for settings
+- **Required Parameters**: Must-have settings
+- **Optional Parameters**: Optional settings with defaults
+- **Validation**: Parameter validation rules
+
+### 5. Error Handling
+- **Custom Exceptions**: Module-specific errors
+- **Error Codes**: Standardized error codes
+- **Recovery Strategies**: Error recovery approaches
+
+### 6. Examples
+- **Basic Usage**: Simple usage examples
+- **Advanced Usage**: Complex scenarios
+- **Integration Examples**: Multi-module usage
 
 ## 🔧 Implementation Guidelines
 
@@ -79,11 +155,11 @@ This document contains comprehensive API design documentation for all core modul
 2. **Set Up Development Environment**
    ```bash
    # Clone the repository
-   git clone https://github.com/vllm-project/vllm-omni.git
+   git clone https://github.com/hsliuustc0106/vllm-omni.git
    cd vllm-omni
 
    # Install dependencies
-   pip install -e ".[dev]"
+   pip install -r requirements-dev.txt
 
    # Set up pre-commit hooks
    pre-commit install
@@ -113,12 +189,28 @@ This document contains comprehensive API design documentation for all core modul
    - Update documentation
    - Submit pull requests
 
+## 📖 Additional Resources
+
+### Architecture Overview
+- [vLLM-omni Architecture Design](../architecture/vLLM-omni%20arch%20design%20doc.md)
+- [Component Categorization](../../ARCHITECTURE_CATEGORIZATION.md)
+
+### Development Guides
+- [Development Setup](../../README.md#development)
+- [Testing Guidelines](../../tests/README.md)
+- [Contributing Guidelines](../../CONTRIBUTING.md)
+
+### Examples
+- [Basic Examples](../../examples/basic/)
+- [Advanced Examples](../../examples/advanced/)
+- [Multimodal Examples](../../examples/multimodal/)
+
 ## 🤝 Contributing
 
 We welcome contributions to improve the API documentation and implementation. Please:
 
 1. **Follow the Template Structure**
-   - Use the [api template](https://github.com/vllm-project/vllm-omni/docs/contributing/design_documents/api_design_template.md) as a guide
+   - Use the master template as a guide
    - Maintain consistency across modules
    - Include all required sections
 

docs/api/overview.md

@@ -1,68 +0,0 @@
-# Summary
-
-## Entry Points
-
-Main entry points for vLLM-omni inference and serving.
-
-- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
-- [vllm_omni.entrypoints.omni_llm.OmniStageLLM][]
-- [vllm_omni.entrypoints.omni_stage.OmniStage][]
-- [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
-
-## Inputs
-
-Input data structures for multi-modal inputs.
-
-- [vllm_omni.inputs.data.OmniTokenInputs][]
-- [vllm_omni.inputs.data.OmniTokensPrompt][]
-- [vllm_omni.inputs.data.OmniEmbedsPrompt][]
-- [vllm_omni.inputs.parse.parse_singleton_prompt_omni][]
-- [vllm_omni.inputs.preprocess.OmniInputPreprocessor][]
-
-## Engine
-
-Engine classes for offline and online inference.
-
-- [vllm_omni.engine.arg_utils.OmniEngineArgs][]
-- [vllm_omni.engine.processor.OmniProcessor][]
-- [vllm_omni.engine.output_processor.OmniRequestState][]
-- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
-- [vllm_omni.engine.PromptEmbedsPayload][]
-- [vllm_omni.engine.AdditionalInformationPayload][]
-- [vllm_omni.engine.OmniEngineCoreRequest][]
-
-## Core
-
-Core scheduling and caching components.
-
-- [vllm_omni.core.sched.scheduler.OmniScheduler][]
-- [vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler][]
-- [vllm_omni.core.sched.output.OmniNewRequestData][]
-- [vllm_omni.core.dit_cache_manager.DiTCacheManager][]
-
-## Model Executor
-
-Model execution components.
-
-- [vllm_omni.model_executor.models.qwen2_5_omni.Qwen2_5OmniForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen2_5_omni_talker.Qwen2_5OmniTalkerForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen2_5_omni_thinker.Qwen2_5OmniThinkerForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavForConditionalGenerationVLLM][]
-- [vllm_omni.model_executor.layers.mrope.MRotaryEmbedding][]
-- [vllm_omni.model_executor.models.registry.OmniModelRegistry][]
-
-## Configuration
-
-Configuration classes.
-
-- [vllm_omni.config.OmniModelConfig][]
-
-## Workers
-
-Worker classes and model runners for distributed inference.
-
-- [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
-- [vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker][]
-- [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
-- [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
-- [vllm_omni.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]

docs/architecture/high_level_arch_design.md

@@ -1,4 +1,4 @@
-# Architecture Overview
+# vLLM-omni Architecture Design Document
 
 # Introduction
 
@@ -114,7 +114,7 @@ vLLM-omni will import core modules from existing vLLM framework and introduce se
 
 For any questions or feedback regarding this design document, please contact  
 Han Gao: [hgaoaf@connect.ust.hk](mailto:hgaoaf@connect.ust.hk),  
-Hongsheng Liu: [vllm-project@gmail.com](mailto:vllm-project@gmail.com)
+Hongsheng Liu: [hsliuustc@gmail.com](mailto:hsliuustc@gmail.com)
 
 [image1]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAALYAAAFwCAYAAAAVNoEZAABxw0lEQVR4XuydBVRcWb6ve90709OStnTScel4x93dIUaEuBJ3dzciRPDgEtzdCRpCBCIEEkIgQHCCBHe+t6uSvNvNnZ55r2/fng5zvrX+69TZp4qkan3nz2/vKup88sm/OWqXL3HxwnnOnj7JubOnuaB6jqtXLqJ28TxXRWmrX8Ha0gIzE0OuX1bFQFsTcxMDLMyMMDXS44bWVRxsLLG5aYKuxlVsLMxwc3bAWF8bNzdnHGwt8fb2wNXZHlcnW3HMDjsbC4Ju3cLO2hwvT3ecnexwsLcmIMAHHy83+X18PVxxtLMUt+1xcbTBw8URRwdr+f1tLM1wd3PB08NZlAs24ue4uTrKf0bD5yfxb8qWDWvZtnE9W0Vt37yBvTu3cWDPTlE7OHviCBfPnsTEwECIfg7Vk0e5fkkVoxtamBneQF1NlSsXTmMrZJaN6elcF7LexFaIZ6SngbOjvfy2l5DUSwhoa2mK9U1jfL09ibwTKaS3IPx2BA4ONnh7uuHr4yUX1d/Pi0B/X1xd7PHz8cTT3UUc85SfIMHBQUJ4a3EfH/mJ4OsjThoXB/njPNydJbEl3rFisTKrli5i9fLFrF21jC0bVdi2aR07tqzn+KF9nD91HF1NDU4eO8CZIwe5fO4UBjoacpEvnDnG5fMnsDQ1Rk/rOje0rwmxLeQd3chAC0d7W6zMjfBwc8JHiGlmpIu5oS6hISHcibgtOrE99+4+wE50dR8hraeHm1zOW4G+BAb4yjtyYICf/FiAED0wMICw0DDxODvR8QNEBeLv7yO6txPeotN7e7lLYku8Y8akccycOhGlGVNRVprJsoXzWLlEmbUrl7Jv5xbOHD8sBD/ApnWrOLZvN2ePH0FDRJVronMf3rdDiH1aiK4p9s+JE0AWS6xE59bA2EAHO6ub3DTRw1UI7Ovjg/oVEWV01Llz567ouH6i+wbz8FEMtiKaBPr7yzuvn6+nED9I3rH9/bwJDgqUSy3r1OHh4YSEBOPu7io6/W1CQ0PFSeAvxHeTd3B/X19JbIl3TB4zjOnjRzJz8jiUFCazYPZ0li1QYrXo4nt3bObE4QPs3LIRlRWLObh7h9jfj5rqGdGtj7N/11Yh9hk0r14W+yfQ0bgiZLZC8/pljPR1ROwwwdxYF2d7GyGfN1cunkJfiB0h4oe3lweRkfeIjn4iF9vf1wcnkbX9fL0ICb4ll1kutrgtkzdEyB4ZGUlQ0C35Y+/cuSPv3oHimOxkkHX2W6KjN3x+Ev+mjBzQmzGD+zJxxCAmjxrC9AmjUBZyL1dWEkJvEF17G6tEB1++aB47Nq0Xsm+Vy71PbHdv2yTP4LL9Q3t3onVdTYhsyCXV02LyqCuyuS762texlk0oXV05e/IQetrqBInuK5v8RUc/JiLiDvYiisgmf7Y2N0UssZJv1a9fFQJ7yiUPkXXrsBAePHiAn+jMsghy//590dlDRRTxk8cW2VhoWLAktsQ7Bv7UleH9egi5+8kFl3Xw+TOnsmTBLDGpXMsOMaFcunCuiChKbF27ml1CbtnEcvt6FbnY504eE517uxB+CxpXLwlxNbhwXjbhvCGXWFOMmZsZibztwMmj+7khsniA6Kwurs4iSoi87OKCuuj4urqa3DQ3xdHBVt6pXZwd0NJQJ1AmbGiw6NC3xYkQLbK6lxA6WC55yP8VWxZrbhEWHiKJLfGOPl06MKhnVwb06Cy2XZg0aqiIJJNYNHcG61cvl2dtZXF78byZrF22iHUrl7B53WqRwxcI6dfLO/WW9Wvk3f3i2VOonj7BeRFT9HW0ULtwRoh/FIMb2hgZ6otMLuvql9klos3COYri3xrOiEH96dujK907d+Snrp3o37c3Gzdu4MKFCxjqG7Bn9y7uRt6Vd+gHUQ9xFsKHiMmnbEy29ReRRbaaEhYWJo8qDZ+fxL8pfbt2pF+3TqJ+ZHCvLiKODGX2tAly8VRWLJFPJBcoKcpFX6E8j9WL300slyrPkXd0WdfeIE4A2UqKbHlQtpJy7tQxtNWvckZILYspOprX0dFSZ8+2DezcqMKsCaOZOGwQvTp1oHPbVnRo2YL2rZrz/TdN+O6rJrRp8YMQvA+q5y9w9NARbgUFEf3wEffuR2FrZyUXOiIiQnTyEPmkUSa2LNLcu/dAElviHb06taF357bIBB/cqzNjh/YXOXs0cxWnsWjebPlKyVwRTebNmILyLAWWzJ3NMuW5LJijIO/o61YvY/WShWxUWcmB3dvZLgQ/emAvZ04cZc/2zULmTaiKuLJv22ZmjxvJqH496dK6pagW9Gnfhp7tWtH5h+9p//23tPruG1qI+v6br/n2qy/47K9/odn337FGRQUHJ2fCbt/GyFhPxI4Q+eRRJrivr7d8AvngQZTo6tGS2BLv6P6jELtLO7nY/bv/yPCBPZkiBJylMIV5QuS5M6czR3EKStMnMU8m+5yZLFGazdwZU1mzfAmrly2UTyzXr1omjyMb1qxgn8jch/fvYdsGFXZsXMMusV0i7j+6V1eGdm3PGBF5pvXtztxBPZjetyujxL8/VPw/erRuzo9C8pZNheBNv+b7r7/iy88/o2PHjixcspTbtyO4IbK4bMnvzp13KySyzB0cIptYRktiS/wXnzdpzpdNmtGtQyt+6tiK/j06MXbEYCaPG8WMKROYOWUc0yaOY6ro4orittKMacxWmMysaeNZtmAei+fPEZ19FquWLmSdkHv5ovnyzL1dxJQ1Qvr1ohaOG8rxJTO4ukYJ/a1LODprDEsGdmdkux8Y2KopHb/+knZffk6LJp/xvRD5e3G7+Zdf0PKbr2gtJG8jZP/uqy9Zu24tqmdOEn77jrxjy96wkb1BExERyd2794m4c1cSW+Idn33ZlL99IaJA6x/o1KaZ6N7tGTmkP+NGDpXLPG38KLnkk8eNYPqkMXKpZ0ydIKQfK2LJHHkHV1aaIYReIM/jMtFlgss691KRzeeKn6G7YwVup3fgr3YAw83KLO/VkbX9OrNhYBfmd23JrC6tmNT+B7p89zVdRafu1uwbOjf9ivbffUUHsS/r4h3F/69Fs2YsWbhAZOn78owdGBgol1vWyWWTyYg7EZLYEu9o/cVXrO71E6aL53Nx1mT6iLw9sE9Phg/qx/jhQ+Q1doTYjhryXu6xTJNtx4scrjCVOSKiKImoIsvispJ19MULlEQGn8GsYf05vWI2+aGOZHqb4LJvNY+v7iP2+hEeXdpL2NF1uG9V5ti4nuwb1pUlPVoz88cWTGzblHGihrVuypA23zOsfXOGCPE7t/qBZt9+g7qmFmFh4XKxZXlb9kbPbdG1ZZ87afj8JP5NGd2mFfuHD8BCeQ43F8/hzMxxzB3al5H9fmJo/z4MH9CXEYMHMEpMKscNFxFl9AjGjxjERNHRFSbIJB/N9MljmDV9MjOnTUJh8nhmT5/CxNHDMNu/luI77pRGB5Dhb0t1zB2qnoRREuxEnpsx989sI3T/apzXKXBz6WRMl01Bf9F4VKcP5uj4Pmwd1p21g7ui3Lsj83p1oEerZrQV3fv7pt9y4uRxbgUHERIWRlh4BMFC9FvBoZLYEu/YObgfp0YO4ZLowtemjhViKXB1wVT2KIxBccRAxorOPVgIPnRQX8YMHSjv4GPeSz5p7EgmjhkhjywyoadPkmXxMfJMPmHoAOJstKhLjKLiaTiVCQ+pSY6jJj6K0jA3CjxNSdI/S4L2MR5c2sW9CzsIOqKC156lGC2bivr8MRye0I+dI3uyelBXlvXvxMA2zWnb7Ds+//QvjBkziiAxaQwJDRWTx3D5kqC/9Ja6xAcU27Rkcgvx67/59yzo2JrDo/thvVYZj51rCT20leCDm7DfvALN5XNRHj2UcQP7MKRPV9HJe4suLuvkQvZRw+SCTxg9XH57pDgZji6bSXVcJHWZidS+SaMuN43alGdysasfBlER6UNpsCNv/azJctQh1eoacdrHuX9hF67bFmK+SpFz04dwZHw/to3sxfqh3ZgssniXZl/LlwT/+p//ia+vn5g8+uPt7SX/YJSrq4sktsQ7un39Fd2/+ZoB337NpJZNWdezI/pKY3FSmU+o2mVCjc3xd/LCw9kHbSNrDl/QQFFpAaNHDGNov3dZfIzo3rLJ5tiR7/L4zJGDCBQ5ui4tntr8TOpF1WWnUJsUQ+2z+1RHBVF1L4DSUFeKAuzJdjEkzUaT+BtniLq8F+ftypitVOCsEPvoxH7sHt2XzUN7MLtHGwa0aEpLIfYXf/sMS0sLPDw8cHR0xN7ORl4Nn5/Evymd2rXgxzY/0EF0wj5dO3NE2xTj5CLMC+rQL6zDtKQe+/I6rEvrMS6owfBNFQY5leimlnLI0pOl2/aLDN6X0UPede9TG1dxS/sMiO5cHO5DcYgHed4OZNkZURbqSfWDYKrvBlN+25/iIDcK/Rx542JEpp0uTzWPEXl+Jw6b52O8fDrnFYZxcupg9o7pK7p2Dxb3aYdil5a0+eYrvvnicw4e3I/lTQsM9G+gr6eL3g1dSWyJd3T44Rt+bN2C+bsOc9w7ApMiuFkMN/JrsSyuwbm8FuuiWkzF/s3Caizf1mCZX4VpbgV6yW+5/iCRpVv3MkFMJBdMn0BehDf50SGQLaLH6xfUy+JHVCiVYZ6UeltRGexO9T2xHxFI0S2RtX0cyHMxI9P2Bo9Flw8/vQ3bjfMwWDqVc0LsE1MGsXd0H7aP6M6SPu2Z3bUVbb9pwjdffsnKlcvQ1rjO1SuXuHZVVpclsf8gNomSvdgfKuaXh//1qD3O4fLrKjRya9HJrUE3uxqr4mrcy2vwLKvDMKsKSyG0s9i3Lxad+20txqJj3xRlJeQ2yyhBNy6Dq0GPWLN8Kbn3g6jNEdEjL5uHphZcO3QRNyNHEnyCqY2NpfbFC2rjnlBxL0x0clkMuUmSmQZPdS9wR3U/t4Line truncated
 

docs/architecture/vllm_omni_design.md

@@ -1,4 +1,4 @@
-# Architecture Design Principles
+# vLLM-omni Software Design Document
 
 ## Overview
 
@@ -206,7 +206,7 @@ dit_stage_config.executor_class = DiffusersPipelineExecutor
 
 Similar to OmniLLM in offline inference, add some asynchronous processing, referring to AsyncLLM
 ```python
-from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.engine.async_llm improt AsyncLLM
 
 
 class AsyncOmniLLM(AsyncLLM):
@@ -317,13 +317,13 @@ class DiffusersPipelineExecutor(Executor):
 
         # Generate image
         prompt_embeds = self._get_and_process_prompt_embeds(scheduler_output, positive_magic)
-        negative_prompt_embeds = self.pipe.embed_prompt(" ")
+        negtive_prompt_embeds = self.pipe.embed_prompt(" ")
 
         width, height = aspect_ratios["16:9"]
 
         image = pipe(
             prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
+            negtive_prompt_embeds=negtive_prompt_embeds,
             width=width,
             height=height,
             num_inference_steps=50,

docs/contributing/README.md

@@ -1,170 +0,0 @@
-# Contributing to vLLM-omni
-
-Thank you for your interest in contributing to vLLM-omni! This document provides guidelines and instructions for contributing.
-
-## Getting Started
-
-### Development Environment
-
-vLLM-omni uses `uv` as the environment manager, to create and manage Python environments. Please follow the documentation to install `uv`. After installing `uv`, you can create a new Python environment using the following commands:
-
-```bash
-uv venv --python 3.12 --seed
-source .venv/bin/activate
-```
-
-Install vLLM from source:
-
-```bash
-git clone https://github.com/vllm-project/vllm.git
-cd vllm
-git checkout v0.11.0
-VLLM_USE_PRECOMPILED=1 uv pip install -e .
-```
-
-Install vLLM-omni from source:
-
-```bash
-git clone https://github.com/vllm-project/vllm-omni.git
-cd vllm_omni
-uv pip install -e .
-```
-
-!!! tip
-    vLLM-omni is compatible with Python versions 3.10 to 3.12. However, we recommend developing with Python 3.12 to minimize the chance of your local environment clashing with our CI environment.
-
-### Linting
-
-vLLM-omni uses `pre-commit` to lint and format the codebase. See [pre-commit documentation](https://pre-commit.com/#usage) if `pre-commit` is new to you. Setting up `pre-commit` is as easy as:
-
-```bash
-uv pip install pre-commit
-pre-commit install
-```
-
-vLLM-omni's `pre-commit` hooks will now run automatically every time you commit.
-
-!!! tip
-    You can manually run the `pre-commit` hooks using:
-
-    ```bash
-    pre-commit run     # runs on staged files
-    pre-commit run -a  # runs on all files (short for --all-files)
-    ```
-
-### Documentation
-
-MkDocs is a fast, simple and downright gorgeous static site generator that's geared towards building project documentation. Documentation source files are written in Markdown, and configured with a single YAML configuration file, `mkdocs.yml`.
-
-Get started with:
-
-```bash
-uv pip install -e ".[docs]"
-```
-
-MkDocs comes with a built-in dev-server that lets you preview your documentation as you work on it. From the root of the repository, run:
-
-```bash
-mkdocs serve                           # with API ref (~10 minutes)
-API_AUTONAV_EXCLUDE=vllm_omni mkdocs serve  # API ref off (~15 seconds)
-```
-
-Once you see `Serving on http://127.0.0.1:8000/` in the logs, the live preview is ready! Open <http://127.0.0.1:8000/> in your browser to see it.
-
-For additional features and advanced configurations, refer to the:
-
-- [MkDocs documentation](https://www.mkdocs.org/)
-- [Material for MkDocs documentation](https://squidfunk.github.io/mkdocs-material/) (the MkDocs theme we use)
-
-### Testing
-
-vLLM-omni uses `pytest` to test the codebase.
-
-```bash
-# Install the test dependencies used in CI (CUDA only)
-uv pip install -e ".[dev]" --torch-backend=auto
-
-# Install some common test dependencies (hardware agnostic)
-uv pip install pytest pytest-asyncio
-
-# Run all tests
-pytest tests/
-
-# Run tests for a single test file with detailed output
-pytest -s -v tests/test_logger.py
-```
-
-!!! warning
-    Currently, not all unit tests pass when run on CPU platforms. If you don't have access to a GPU platform to run unit tests locally, rely on the continuous integration system to run the tests for now.
-
-## Issues
-
-If you encounter a bug or have a feature request, please search existing issues first to see if it has already been reported. If not, please file a new issue, providing as much relevant information as possible.
-
-!!! important
-    If you discover a security vulnerability, please report it by creating a GitHub issue with the `security` label.
-
-## Pull Requests & Code Reviews
-
-Thank you for your contribution to vLLM-omni! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM-omni maintain the code quality and improve the efficiency of the review process.
-
-### DCO and Signed-off-by
-
-When contributing changes to this project, you must agree to the [DCO](https://developercertificate.org/). Commits must include a `Signed-off-by:` header which certifies agreement with the terms of the DCO.
-
-Using `-s` with `git commit` will automatically add this header.
-
-!!! tip
-    You can enable automatic sign-off via your IDE:
-
-    - **PyCharm**: Click on the `Show Commit Options` icon to the right of the `Commit and Push...` button in the `Commit` window. It will bring up a `git` window where you can modify the `Author` and enable `Sign-off commit`.
-    - **VSCode**: Open the Settings editor and enable the `Git: Always Sign Off` (`git.alwaysSignOff`) field.
-
-### PR Title and Classification
-
-Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:
-
-- `[Bugfix]` for bug fixes.
-- `[CI/Build]` for build or continuous integration improvements.
-- `[Doc]` for documentation fixes and improvements.
-- `[Model]` for adding a new model or improving an existing model. Model name should appear in the title.
-- `[Frontend]` For changes on the vLLM-omni frontend (e.g., OpenAI API server, `OmniLLM` class, etc.)
-- `[Kernel]` for changes affecting CUDA kernels or other compute kernels.
-- `[Core]` for changes in the core vLLM-omni logic (e.g., `OmniProcessor`, `OmniScheduler`, etc.)
-- `[Hardware][Vendor]` for hardware-specific changes. Vendor name should appear in the prefix, such as [Ascend] for Ascend NPUs.
-- `[Misc]` for PRs that do not fit the above categories. Please use this sparingly.
-
-!!! note
-    If the PR spans more than one category, please include all relevant prefixes.
-
-### Code Quality
-
-The PR needs to meet the following code quality standards:
-
-- We adhere to Google Python style guide and Google C++ style guide.
-- Pass all linter checks.
-- The code needs to be well-documented to ensure future contributors can easily understand the code.
-- Include sufficient tests to ensure the project stays correct and robust. This includes both unit tests and integration tests.
-- Please add documentation to `docs/` if the PR modifies the user-facing behaviors of vLLM-omni. It helps vLLM-omni users understand and utilize the new features or changes.
-
-### Notes for Large Changes
-
-Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with `rfc-required` and might not go through the PR.
-
-### What to Expect for the Reviews
-
-The goal of the vLLM-omni team is to be a _transparent reviewing machine_. We would like to make the review process transparent and efficient and make sure no contributor feels confused or frustrated. However, the vLLM-omni team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:
-
-- After the PR is submitted, the PR will be assigned to a reviewer. Every reviewer will pick up the PRs based on their expertise and availability.
-- After the PR is assigned, the reviewer will provide status updates every 2-3 days. If the PR is not reviewed within 7 days, please feel free to ping the reviewer or the vLLM-omni team.
-- After the review, the reviewer will put an `action-required` label on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR.
-- Please respond to all comments within a reasonable time frame. If a comment isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion.
-- Note that not all CI checks will be executed due to limited computational resources. The reviewer will add `ready` label to the PR when the PR is ready to merge or a full CI run is needed.
-
-## Additional Resources
-
-- [Design Documents](design_documents/index.md) - Architecture and design documentation
-
-## Thank You
-
-Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM-omni. All of your contributions help make vLLM-omni a great tool and community for everyone!

docs/contributing/design_documents/index.md

@@ -1,12 +0,0 @@
-# Design Documents
-
-This section contains design documents and architecture specifications for vLLM-omni.
-
-## Architecture Documents
-
-- [vLLM-omni Software Design Document](vllm_omni_design.md)
-- [Architecture Overview](architecture_overview.md)
-
-## API Design Documentation
-
-- [vLLM-omni API Documentation](api_design_doc.md)

docs/index.md

@@ -1,32 +0,0 @@
-# Welcome to vLLM-omni
-
-**vLLM-omni** extends vLLM capabilities to support multi-modality model inference and serving, particularly focusing on non-autoregressive structures and non-textual outputs. It aims to support:
-
-- **Multi-modal Models**: Text, image, video, audio, and sensor data processing
-- **Non-autoregressive Architectures**: Diffusion Transformers (DiT) and other parallel generation models
-- **Heterogeneous Outputs**: Beyond traditional text generation to multimodal outputs
-
-## <span class="twemoji">🎯</span> Key Features
-
-### Multi-Engine Support
-
-- **Autoregressive Engine**: Traditional text generation with enhanced KV-caching
-- **Diffusion Engine**: Support for DiT models and iterative generation
-- **Hybrid Engine**: Combined AR+DiT processing pipelines
-
-### Modality Processing
-
-- **Text**: Advanced tokenization and embedding generation
-- **Image**: Vision encoder integration (CLIP, etc.)
-- **Audio**: Speech processing and audio embedding
-
-## <span class="twemoji">🤖</span> Supported Models
-
-- Qwen-omni (Thinker-Talker-Codec structure)
-- HunyunaImage 3.0 (Ongoing)
-- Qwen-Image (Ongoing)
-
-## <span class="twemoji">📚</span> Documentation Navigation
-
-- To run open-source models on vLLM-Omni, we recommend starting with the [:material-code-tags: User Quide](user_guide/getting_started/quickstart.md)
-- To develop and contribute to vLLM-Omni, we recommend starting with the [:material-tools: Developer Guide](contributing/README.md)

docs/mkdocs/hooks/generate_examples.py

@@ -1,233 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-omni project
-import itertools
-import logging
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Literal
-
-import regex as re
-
-logger = logging.getLogger("mkdocs")
-
-ROOT_DIR = Path(__file__).parent.parent.parent.parent
-ROOT_DIR_RELATIVE = "../../../../.."
-EXAMPLE_DIR = ROOT_DIR / "examples"
-EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
-
-
-def fix_case(text: str) -> str:
-    subs = {
-        "api": "API",
-        "cli": "CLI",
-        "cpu": "CPU",
-        "llm": "LLM",
-        "mae": "MAE",
-        "tpu": "TPU",
-        "gguf": "GGUF",
-        "lora": "LoRA",
-        "rlhf": "RLHF",
-        "vllm": "vLLM",
-        "omni": "Omni",
-        "openai": "OpenAI",
-        "lmcache": "LMCache",
-        "multilora": "MultiLoRA",
-        "mlpspeculator": "MLPSpeculator",
-        r"fp\d+": lambda x: x.group(0).upper(),  # e.g. fp16, fp32
-        r"int\d+": lambda x: x.group(0).upper(),  # e.g. int8, int16
-    }
-    for pattern, repl in subs.items():
-        text = re.sub(rf"\b{pattern}\b", repl, text, flags=re.IGNORECASE)
-    return text
-
-
-@dataclass
-class Example:
-    """
-    Example class for generating documentation content from a given path.
-
-    Attributes:
-        path (Path): The path to the main directory or file.
-        category (str): The category of the document.
-        main_file (Path): The main file in the directory.
-        other_files (list[Path]): list of other files in the directory.
-        title (str): The title of the document.
-
-    Methods:
-        __post_init__(): Initializes the main_file, other_files, and title attributes.
-        determine_main_file() -> Path: Determines the main file in the given path.
-        determine_other_files() -> list[Path]: Determines other files in the directory excluding the main file.
-        determine_title() -> str: Determines the title of the document.
-        generate() -> str: Generates the documentation content.
-    """  # noqa: E501
-
-    path: Path
-    category: str = None
-    main_file: Path = field(init=False)
-    other_files: list[Path] = field(init=False)
-    title: str = field(init=False)
-
-    def __post_init__(self):
-        self.main_file = self.determine_main_file()
-        self.other_files = self.determine_other_files()
-        self.title = self.determine_title()
-
-    @property
-    def is_code(self) -> bool:
-        return self.main_file.suffix != ".md"
-
-    def determine_main_file(self) -> Path:
-        """
-        Determines the main file in the given path.
-        If the path is a file, it returns the path itself. Otherwise, it searches
-        for Markdown files (*.md) in the directory and returns the first one found.
-        Returns:
-            Path: The main file path, either the original path if it's a file or the first
-            Markdown file found in the directory.
-        Raises:
-            IndexError: If no Markdown files are found in the directory.
-        """  # noqa: E501
-        return self.path if self.path.is_file() else list(self.path.glob("*.md")).pop()
-
-    def determine_other_files(self) -> list[Path]:
-        """
-        Determine other files in the directory excluding the main file.
-
-        This method checks if the given path is a file. If it is, it returns an empty list.
-        Otherwise, it recursively searches through the directory and returns a list of all
-        files that are not the main file.
-
-        Returns:
-            list[Path]: A list of Path objects representing the other files in the directory.
-        """  # noqa: E501
-        if self.path.is_file():
-            return []
-        # Text file extensions that can be safely included
-        text_extensions = {
-            ".py",
-            ".md",
-            ".sh",
-            ".yaml",
-            ".yml",
-            ".json",
-            ".txt",
-            ".toml",
-            ".cfg",
-            ".ini",
-        }
-        is_other_file = lambda file: (file.is_file() and file != self.main_file and file.suffix in text_extensions)  # noqa: E731
-        return [file for file in self.path.rglob("*") if is_other_file(file)]
-
-    def determine_title(self) -> str:
-        if not self.is_code:
-            # Specify encoding for building on Windows
-            with open(self.main_file, encoding="utf-8") as f:
-                first_line = f.readline().strip()
-            match = re.match(r"^#\s+(?P<title>.+)$", first_line)
-            if match:
-                return match.group("title")
-        return fix_case(self.path.stem.replace("_", " ").title())
-
-    def fix_relative_links(self, content: str) -> str:
-        """
-        Fix relative links in markdown content by converting them to gh-file
-        format.
-
-        Args:
-            content (str): The markdown content to process
-
-        Returns:
-            str: Content with relative links converted to gh-file format
-        """
-        # Regex to match markdown links [text](relative_path)
-        # This matches links that don't start with http, https, ftp, or #
-        link_pattern = r"\[([^\]]*)\]\((?!(?:https?|ftp)://|#)([^)]+)\)"
-
-        def replace_link(match):
-            link_text = match.group(1)
-            relative_path = match.group(2)
-
-            # Make relative to repo root
-            gh_file = (self.main_file.parent / relative_path).resolve()
-            gh_file = gh_file.relative_to(ROOT_DIR)
-
-            # Make GitHub URL
-            url = "https://github.com/vllm-project/vllm-omni/"
-            url += "tree/main" if self.path.is_dir() else "blob/main"
-            gh_url = f"{url}/{gh_file}"
-
-            return f"[{link_text}]({gh_url})"
-
-        return re.sub(link_pattern, replace_link, content)
-
-    def generate(self) -> str:
-        content = f"# {self.title}\n\n"
-        url = "https://github.com/vllm-project/vllm-omni/"
-        url += "tree/main" if self.path.is_dir() else "blob/main"
-        content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"
-
-        # Use long code fence to avoid issues with
-        # included files containing code fences too
-        code_fence = "``````"
-
-        if self.is_code:
-            # Use relative path from ROOT_DIR for snippet inclusion
-            main_file_rel_path = self.main_file.relative_to(ROOT_DIR)
-            content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel_path}"\n{code_fence}\n'
-        else:
-            with open(self.main_file, encoding="utf-8") as f:
-                # Skip the title from md snippets as it's been included above
-                main_content = f.readlines()[1:]
-            content += self.fix_relative_links("".join(main_content))
-        content += "\n"
-
-        if not self.other_files:
-            return content
-
-        content += "## Example materials\n\n"
-        for file in sorted(self.other_files):
-            # Use relative path from ROOT_DIR for snippet inclusion
-            file_rel_path = file.relative_to(ROOT_DIR)
-            content += f'??? abstract "{file.relative_to(self.path)}"\n'
-            if file.suffix != ".md":
-                content += f"    {code_fence}{file.suffix[1:]}\n"
-            content += f'    --8<-- "{file_rel_path}"\n'
-            if file.suffix != ".md":
-                content += f"    {code_fence}\n"
-
-        return content
-
-
-def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
-    logger.info("Generating example documentation")
-    logger.debug("Root directory: %s", ROOT_DIR.resolve())
-    logger.debug("Example directory: %s", EXAMPLE_DIR.resolve())
-    logger.debug("Example document directory: %s", EXAMPLE_DOC_DIR.resolve())
-
-    # Create the EXAMPLE_DOC_DIR if it doesn't exist
-    if not EXAMPLE_DOC_DIR.exists():
-        EXAMPLE_DOC_DIR.mkdir(parents=True)
-
-    categories = sorted(p for p in EXAMPLE_DIR.iterdir() if p.is_dir())
-
-    examples = []
-    glob_patterns = ["*.py", "*.md"]
-    # Find categorised examples
-    for category in categories:
-        globs = [category.glob(pattern) for pattern in glob_patterns]
-        for path in itertools.chain(*globs):
-            examples.append(Example(path, category.stem))
-        # Find examples in subdirectories
-        for path in category.glob("*/*.md"):
-            examples.append(Example(path.parent, category.stem))
-
-    # Generate the example documentation
-    for example in sorted(examples, key=lambda e: e.path.stem):
-        example_name = f"{example.path.stem}.md"
-        doc_path = EXAMPLE_DOC_DIR / example.category / example_name
-        if not doc_path.parent.exists():
-            doc_path.parent.mkdir(parents=True)
-        # Specify encoding for building on Windows
-        with open(doc_path, "w+", encoding="utf-8") as f:
-            f.write(example.generate())
-        logger.debug("Example generated: %s", doc_path.relative_to(ROOT_DIR))

docs/mkdocs/stylesheets/extra.css

@@ -1,55 +0,0 @@
-/* Custom CSS for vLLM-omni documentation - based on vLLM's extra.css */
-
-/* Improve code block readability */
-.highlight {
-    border-radius: 4px;
-}
-
-/* Style API documentation */
-.md-typeset .doc-contents {
-    margin-bottom: 1.5em;
-}
-
-/* Improve table readability */
-.md-typeset table {
-    display: table;
-    width: 100%;
-}
-
-/* Light mode: darker section titles (from vLLM) */
-body[data-md-color-scheme="default"] .md-nav__item--section > label.md-nav__link .md-ellipsis {
-  color: rgba(0, 0, 0, 0.7) !important;
-  font-weight: 700;
-}
-
-/* Dark mode: lighter gray section titles (from vLLM) */
-body[data-md-color-scheme="slate"] .md-nav__item--section > label.md-nav__link .md-ellipsis {
-  color: rgba(255, 255, 255, 0.75) !important;
-  font-weight: 700;
-}
-
-/* Navigation: Expand first level of API Reference by default, collapse deeper levels */
-/* Show direct children of API Reference section (vllm_omni's children: config, core, entrypoints, etc.) */
-.md-nav__item--nested[data-md-state] > .md-nav {
-    display: none;
-}
-
-/* Expand active items */
-.md-nav__item--nested.md-nav__item--active > .md-nav {
-    display: block;
-}
-
-/* For API Reference: show first level children (config, core, entrypoints, etc.) by default */
-.md-nav__item--section:has([href*="api/vllm_omni/index"]) > .md-nav > .md-nav__list > .md-nav__item--nested > .md-nav {
-    display: block;
-}
-
-/* But keep deeper levels collapsed */
-.md-nav__item--section:has([href*="api/vllm_omni/index"]) > .md-nav > .md-nav__list > .md-nav__item--nested > .md-nav > .md-nav__list > .md-nav__item--nested > .md-nav {
-    display: none;
-}
-
-/* Expand deeper levels only when active */
-.md-nav__item--section:has([href*="api/vllm_omni/index"]) > .md-nav > .md-nav__list > .md-nav__item--nested > .md-nav > .md-nav__list > .md-nav__item--nested.md-nav__item--active > .md-nav {
-    display: block;
-}

docs/user_guide/examples/index.md

@@ -1,8 +0,0 @@
-# Examples
-
-vLLM-omni's examples are split into two categories:
-
-- If you are using vLLM-omni from within Python code, see the *Offline Inference* section.
-- If you are using vLLM-omni from an HTTP application or client, see the *Online Serving* section.
-
-For detailed example documentation, check the [examples directory](https://github.com/vllm-project/vllm-omni/tree/main/examples) in the repository.

docs/user_guide/examples/offline_inference/qwen2_5_omni.md

@@ -1,64 +0,0 @@
-# Offline Example of vLLM-omni for Qwen2.5-omni
-
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples\offline_inference\qwen2_5_omni>.
-
-
-## 🛠️ Installation
-
-Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
-
-## Run examples (Qwen2.5-omni)
-### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
-```bash
-tar -xf <Your Download Path>/seedtts_testset.tar
-cp seedtts_testset/en/meta.lst examples/offline_inference/qwen2_5_omni/meta.lst
-python3 examples/offline_inference/qwen2_5_omni/extract_prompts.py \
-  --input examples/offline_inference/qwen2_5_omni/meta.lst \
-  --output examples/offline_inference/qwen2_5_omni/top100.txt \
-  --topk 100
-```
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen2_5_omni
-```
-Then run the command below.
-```bash
-bash run_multiple_prompts.sh
-```
-### Single Prompts
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen2_5_omni
-```
-Then run the command below.
-```bash
-bash run_single_prompt.sh
-```
-
-## Example materials
-
-??? abstract "end2end.py"
-    ``````py
-    --8<-- "examples\offline_inference\qwen2_5_omni\end2end.py"
-    ``````
-??? abstract "extract_prompts.py"
-    ``````py
-    --8<-- "examples\offline_inference\qwen2_5_omni\extract_prompts.py"
-    ``````
-??? abstract "processing_omni.py"
-    ``````py
-    --8<-- "examples\offline_inference\qwen2_5_omni\processing_omni.py"
-    ``````
-??? abstract "run_multiple_prompts.sh"
-    ``````sh
-    --8<-- "examples\offline_inference\qwen2_5_omni\run_multiple_prompts.sh"
-    ``````
-??? abstract "run_single_prompt.sh"
-    ``````sh
-    --8<-- "examples\offline_inference\qwen2_5_omni\run_single_prompt.sh"
-    ``````
-??? abstract "utils.py"
-    ``````py
-    --8<-- "examples\offline_inference\qwen2_5_omni\utils.py"
-    ``````

docs/user_guide/examples/online_serving/README.md

@@ -1,35 +0,0 @@
-# Online serving Example of vLLM-omni for Qwen2.5-omni
-
-Source <https://github.com/vllm-project/vllm-omni/blob/main/examples\online_serving\README.md>.
-
-
-## 🛠️ Installation
-
-Please refer to [README.md](https://github.com/vllm-project/vllm-omni/blob/main/README.md)
-
-## Run examples (Qwen2.5-omni)
-
-Launch the server
-```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
-```
-
-If you have custom stage configs file, launch the server with command below
-```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
-```
-
-Get into the example folder
-```bash
-cd examples/online_serving
-```
-
-Send request via python
-```bash
-python openai_chat_completion_client_for_multimodal_generation.py
-```
-
-Send request via curl
-```bash
-bash run_curl_multimodal_generation.sh
-```

docs/user_guide/examples/online_serving/index.md

@@ -1,5 +0,0 @@
-# Online Serving
-
-## Available Examples
-
-Online serving is under active development. Coming soon.

docs/user_guide/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.md

@@ -1,7 +0,0 @@
-# OpenAI Chat Completion Client For Multimodal Generation
-
-Source <https://github.com/vllm-project/vllm-omni/blob/main/examples\online_serving\openai_chat_completion_client_for_multimodal_generation.py>.
-
-``````py
---8<-- "examples\online_serving\openai_chat_completion_client_for_multimodal_generation.py"
-``````

docs/user_guide/getting_started/quickstart.md

@@ -1,55 +0,0 @@
-# Quickstart
-
-This guide will help you get started with vLLM-omni.
-
-## Installation
-
-Use Docker to keep consistent basic environment (Optional, Recommended)
-
-```bash
-docker run --gpus all --ipc=host --network=host -v $source_dir:$container_dir --rm --name $container_name -it nvcr.io/nvidia/pytorch:25.01-py3 bash
-```
-
-Set up basic uv environment
-
-```bash
-pip install uv
-uv venv --python 3.12 --seed
-source .venv/bin/activate
-```
-
-Install vLLM
-
-```bash
-uv pip install vllm==0.11.0 --torch-backend=auto
-```
-
-Install vLLM-omni
-
-```bash
-git clone https://github.com/vllm-project/vllm-omni.git
-cd vllm_omni
-uv pip install -e .
-```
-
-## Run examples (Qwen2.5-omni)
-
-Get into the example folder
-
-```bash
-cd examples/offline_inference/qwen_2_5_omni
-```
-
-Then run.
-
-```bash
-bash run.sh
-```
-
-The output audio is saved in ./output_audio
-
-## Next Steps
-
-- Read the [architecture documentation](../../contributing/design_documents/vllm_omni_design.md)
-- Check out the [API reference](../../api/overview.md)
-- Explore the [examples](../examples/index.md)

examples/offline_inference/qwen2_5_omni/README.md

@@ -1,42 +0,0 @@
-# Offline Example of vLLM-omni for Qwen2.5-omni
-
-## 🛠️ Installation
-
-Please refer to [README.md](../../../README.md)
-
-## Run examples (Qwen2.5-omni)
-### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
-```bash
-tar -xf <Your Download Path>/seedtts_testset.tar
-cp seedtts_testset/en/meta.lst examples/offline_inference/qwen2_5_omni/meta.lst
-python3 examples/offline_inference/qwen2_5_omni/extract_prompts.py \
-  --input examples/offline_inference/qwen2_5_omni/meta.lst \
-  --output examples/offline_inference/qwen2_5_omni/top100.txt \
-  --topk 100
-```
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen2_5_omni
-```
-Then run the command below.
-```bash
-bash run_multiple_prompts.sh
-```
-### Single Prompts
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen2_5_omni
-```
-Then run the command below.
-```bash
-bash run_single_prompt.sh
-```
-
-### FAQ
-
-If you encounter error about backend of librosa, try to install ffmpeg with command below.
-```
-sudo apt update
-sudo apt install ffmpeg
-```

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -1,307 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use vLLM-omni for running offline inference
-with the correct prompt format on Qwen2.5-Omni
-"""
-
-import argparse
-import os
-from typing import NamedTuple
-
-import soundfile as sf
-from vllm.assets.audio import AudioAsset
-from vllm.assets.image import ImageAsset
-from vllm.assets.video import VideoAsset
-from vllm.multimodal.image import convert_image_mode
-from vllm.sampling_params import SamplingParams
-from vllm.utils import FlexibleArgumentParser
-
-from vllm_omni import OmniLLM
-
-SEED = 42
-
-
-class QueryResult(NamedTuple):
-    inputs: dict
-    limit_mm_per_prompt: dict[str, int]
-
-
-# NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on
-# lower-end GPUs.
-# Unless specified, these settings have been tested to work on a single L4.
-
-default_system = (
-    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
-    "Group, capable of perceiving auditory and visual inputs, as well as "
-    "generating text and speech."
-)
-
-
-def get_text_query(question: str = None) -> QueryResult:
-    if question is None:
-        question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
-    prompt = (
-        f"<|im_start|>system\n{default_system}<|im_end|>\n"
-        "<|im_start|>user\n"
-        f"{question}<|im_end|>\n"
-        f"<|im_start|>assistant\n"
-    )
-    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]()
-
-    omni_llm = OmniLLM(
-        model=model_name,
-        log_stats=args.enable_stats,
-        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
-        init_sleep_seconds=args.init_sleep_seconds,
-        batch_timeout=args.batch_timeout,
-        init_timeout=args.init_timeout,
-        shm_threshold_bytes=args.shm_threshold_bytes,
-    )
-    thinker_sampling_params = SamplingParams(
-        temperature=0.0,  # Deterministic - no randomness
-        top_p=1.0,  # Disable nucleus sampling
-        top_k=-1,  # Disable top-k sampling
-        max_tokens=2048,
-        seed=SEED,  # Fixed seed for sampling
-        detokenize=True,
-        repetition_penalty=1.1,
-    )
-    talker_sampling_params = SamplingParams(
-        temperature=0.9,
-        top_p=0.8,
-        top_k=40,
-        max_tokens=2048,
-        seed=SEED,  # Fixed seed for sampling
-        detokenize=True,
-        repetition_penalty=1.05,
-        stop_token_ids=[8294],
-    )
-    code2wav_sampling_params = SamplingParams(
-        temperature=0.0,  # Deterministic - no randomness
-        top_p=1.0,  # Disable nucleus sampling
-        top_k=-1,  # Disable top-k sampling
-        max_tokens=2048,
-        seed=SEED,  # Fixed seed for sampling
-        detokenize=True,
-        repetition_penalty=1.1,
-    )
-
-    sampling_params_list = [
-        thinker_sampling_params,
-        talker_sampling_params,
-        code2wav_sampling_params,
-    ]
-
-    if args.txt_prompts is None:
-        prompts = [query_result.inputs for _ in range(args.num_prompts)]
-    else:
-        assert args.query_type == "text", "txt-prompts is only supported for text query type"
-        with open(args.txt_prompts, encoding="utf-8") as f:
-            lines = [ln.strip() for ln in f.readlines()]
-            prompts = [get_text_query(ln).inputs for ln in lines if ln != ""]
-            print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
-
-    omni_outputs = omni_llm.generate(prompts, sampling_params_list)
-
-    output_dir = None
-    if args.save_results:
-        # Determine output directory: prefer --output-dir; fallback to --output-wav
-        output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
-        os.makedirs(output_dir, exist_ok=True)
-    for stage_outputs in omni_outputs:
-        if stage_outputs.final_output_type == "text":
-            for output in stage_outputs.request_output:
-                request_id = int(output.request_id)
-                text_output = output.outputs[0].text
-                if args.save_results:
-                    # Save aligned text file per request
-                    prompt_text = prompts[request_id]["prompt"]
-                    out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
-                    lines = []
-                    lines.append("Prompt:\n")
-                    lines.append(str(prompt_text) + "\n")
-                    lines.append("vllm_text_output:\n")
-                    lines.append(str(text_output).strip() + "\n")
-                    try:
-                        with open(out_txt, "w", encoding="utf-8") as f:
-                            f.writelines(lines)
-                    except Exception as e:
-                        print(f"[Warn] Failed writing text file {out_txt}: {e}")
-                    print(f"Request ID: {request_id}, Text saved to {out_txt}")
-                else:
-                    print(f"Request ID: {request_id}, Text output:\n{text_output}")
-        elif stage_outputs.final_output_type == "audio":
-            for output in stage_outputs.request_output:
-                request_id = int(output.request_id)
-                audio_tensor = output.multimodal_output["audio"]
-                if args.save_results:
-                    output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
-                    sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
-                    print(f"Request ID: {request_id}, Saved audio to {output_wav}")
-                else:
-                    print(f"Request ID: {request_id}, Received audio output (not saved).")
-
-
-def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
-    parser.add_argument(
-        "--query-type",
-        "-q",
-        type=str,
-        default="mixed_modalities",
-        choices=query_map.keys(),
-        help="Query type.",
-    )
-    parser.add_argument(
-        "--enable-stats",
-        action="store_true",
-        default=False,
-        help="Enable writing detailed statistics (default: disabled)",
-    )
-    parser.add_argument(
-        "--init-sleep-seconds",
-        type=int,
-        default=20,
-        help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
-    )
-    parser.add_argument(
-        "--batch-timeout",
-        type=int,
-        default=5,
-        help="Timeout for batching in seconds (default: 5)",
-    )
-    parser.add_argument(
-        "--init-timeout",
-        type=int,
-        default=300,
-        help="Timeout for initializing stages in seconds (default: 300)",
-    )
-    parser.add_argument(
-        "--shm-threshold-bytes",
-        type=int,
-        default=65536,
-        help="Threshold for using shared memory in bytes (default: 65536)",
-    )
-    parser.add_argument(
-        "--save-results",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help="Save text/audio outputs to files (default: enabled). Use --no-save-results to disable.",
-    )
-    parser.add_argument(
-        "--output-wav",
-        default="output_audio",
-        help="[Deprecated] Output wav directory (use --output-dir).",
-    )
-    parser.add_argument(
-        "--num-prompts",
-        type=int,
-        default=1,
-        help="Number of prompts to generate.",
-    )
-    parser.add_argument(
-        "--txt-prompts",
-        type=str,
-        default=None,
-        help="Path to a .txt file with one prompt per line (preferred).",
-    )
-
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    main(args)

examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh

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

examples/offline_inference/qwen2_5_omni/run_single_prompt.sh

@@ -1,2 +0,0 @@
-python end2end.py --output-wav output_audio \
-                  --query-type use_audio_in_video

examples/offline_inference/qwen_2_5_omni/README.md

@@ -1,8 +1,5 @@
 # Offline Example of vLLM-omni for Qwen2.5-omni
 
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen_2_5_omni>.
-
-
 ## 🛠️ Installation
 
 ### Environment setup
@@ -52,7 +49,7 @@ 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.
+Just run the command below. If no error, it demonstrates that the installation is successfull.
 ```bash
 python -c "import vllm._C"
 ```
@@ -62,36 +59,34 @@ Install additional requirements for vllm-omni
 ```bash
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm_omni
-uv pip install -e .
+uv pip install -r requirements.txt
 ```
 
 ## Run examples (Qwen2.5-omni)
-
+### Multiple Prompts
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
+```bash
+tar -xf <Your Download Path>/seedtts_testset.tar
+cp seedtts_testset/en/meta.lst examples/offline_inference/qwen_2_5_omni/meta.lst
+python3 examples/offline_inference/qwen_2_5_omni/extract_prompts.py \
+  --input examples/offline_inference/qwen_2_5_omni/meta.lst \
+  --output examples/offline_inference/qwen_2_5_omni/top100.txt \
+  --topk 100
+```
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen_2_5_omni
 ```
-Then run
+Modify PYTHONPATH in run_multiple_prompts.sh as your path of vllm_omni. Then run.
 ```bash
-bash run.sh
+bash run_multiple_prompts.sh
+```
+### Single Prompts
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen_2_5_omni
+```
+Modify PYTHONPATH in run_single_prompt.sh as your path of vllm_omni. Then run.
+```bash
+bash run_single_prompt.sh
 ```
-The output audio is saved in ./output_audio
-
-## Example materials
-
-??? abstract "end2end.py"
-    ``````py
-    --8<-- "examples/offline_inference/qwen_2_5_omni/end2end.py"
-    ``````
-??? abstract "processing_omni.py"
-    ``````py
-    --8<-- "examples/offline_inference/qwen_2_5_omni/processing_omni.py"
-    ``````
-??? abstract "run.sh"
-    ``````sh
-    --8<-- "examples/offline_inference/qwen_2_5_omni/run.sh"
-    ``````
-??? abstract "utils.py"
-    ``````py
-    --8<-- "examples/offline_inference/qwen_2_5_omni/utils.py"
-    ``````

examples/offline_inference/qwen_2_5_omni/end2end.py

@@ -0,0 +1,253 @@
+import argparse
+import os
+import os as _os_env_toggle
+import random
+
+import numpy as np
+import soundfile as sf
+import torch
+from utils import make_omni_prompt
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+_os_env_toggle.environ["VLLM_USE_V1"] = "1"
+
+SEED = 42
+# Set all random seeds
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+torch.cuda.manual_seed(SEED)
+torch.cuda.manual_seed_all(SEED)
+
+# Make PyTorch deterministic
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+
+# Set environment variables for deterministic behavior
+os.environ["PYTHONHASHSEED"] = str(SEED)
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--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():
+    args = parse_args()
+    model_name = args.model
+    try:
+        # Preferred: load from txt file (one prompt per line)
+        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
+            with open(args.txt_prompts, "r", encoding="utf-8") as f:
+                lines = [ln.strip() for ln in f.readlines()]
+            args.prompts = [ln for ln in lines if ln != ""]
+            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
+    except Exception as e:
+        print(f"[Error] Failed to load prompts: {e}")
+        raise
+
+    if args.prompts is None:
+        raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts <file.txt> (with --prompt_type text)")
+    omni_llm = OmniLLM(
+        model=model_name,
+        log_stats=args.enable_stats,
+        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
+        init_sleep_seconds=args.init_sleep_seconds,
+        batch_timeout=args.batch_timeout,
+        init_timeout=args.init_timeout,
+        shm_threshold_bytes=args.shm_threshold_bytes,
+    )
+    thinker_sampling_params = SamplingParams(
+        temperature=0.0,  # Deterministic - no randomness
+        top_p=1.0,  # Disable nucleus sampling
+        top_k=-1,  # Disable top-k sampling
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.1,
+    )
+    talker_sampling_params = SamplingParams(
+        temperature=0.9,
+        top_p=0.8,
+        top_k=40,
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.05,
+        stop_token_ids=[8294],
+    )
+    code2wav_sampling_params = SamplingParams(
+        temperature=0.0,  # Deterministic - no randomness
+        top_p=1.0,  # Disable nucleus sampling
+        top_k=-1,  # Disable top-k sampling
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.1,
+    )
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        talker_sampling_params,
+        code2wav_sampling_params,
+    ]
+
+    prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
+    omni_outputs = omni_llm.generate(prompt, sampling_params_list)
+
+    # Determine output directory: prefer --output-dir; fallback to --output-wav
+    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    os.makedirs(output_dir, exist_ok=True)
+    for stage_outputs in omni_outputs:
+        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
+                prompt_text = args.prompts[request_id]
+                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
+                lines = []
+                lines.append("Prompt:\n")
+                lines.append(str(prompt_text) + "\n")
+                lines.append("vllm_text_output:\n")
+                lines.append(str(text_output).strip() + "\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        f.writelines(lines)
+                except Exception as e:
+                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
+                print(f"Request ID: {request_id}, Text saved to {out_txt}")
+        elif stage_outputs.final_output_type == "audio":
+            for output in stage_outputs.request_output:
+                request_id = int(output.request_id)
+                audio_tensor = output.multimodal_output["audio"]
+                output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
+                sf.write(
+                    output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000
+                )
+                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/qwen_2_5_omni/extract_prompts.py

@@ -1,34 +1,30 @@
 #!/usr/bin/env python3
 import argparse
 from typing import Optional
+import os
+import torch
 
 
 def extract_prompt(line: str) -> Optional[str]:
     # 提取第一个 '|' 与第二个 '|' 之间的内容
-    i = line.find("|")
+    i = line.find('|')
     if i == -1:
         return None
-    j = line.find("|", i + 1)
+    j = line.find('|', i + 1)
     if j == -1:
         return None
-    return line[i + 1 : j].strip()
+    return line[i + 1:j].strip()
 
 
 def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--input", "-i", required=True, help="Input .lst file path")
     parser.add_argument("--output", "-o", required=True, help="Output file path")
-    parser.add_argument(
-        "--topk",
-        "-k",
-        type=int,
-        default=100,
-        help="Extract the top K prompts (default: 100)",
-    )
+    parser.add_argument("--topk", "-k", type=int, default=100, help="Extract the top K prompts (default: 100)")
     args = parser.parse_args()
 
     prompts = []
-    with open(args.input, encoding="utf-8", errors="ignore") as f:
+    with open(args.input, "r", encoding="utf-8", errors="ignore") as f:
         for line in f:
             if len(prompts) >= args.topk:
                 break
@@ -43,3 +39,5 @@ def main() -> None:
 
 if __name__ == "__main__":
     main()
+
+

examples/offline_inference/qwen_2_5_omni/processing_omni.py

@@ -0,0 +1,391 @@
+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}, "
+            f"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, "
+            f"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}], "
+            f"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=}, "
+        f"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=}, "
+        f"time={time.time() - st:.3f}s"
+    )
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
+    video = vr.get_batch(idx).asnumpy()
+    video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
+    return video, total_duration, nframes
+
+
+VIDEO_READER_BACKENDS = {
+    "decord": _read_video_decord,
+    "torchvision": _read_video_torchvision,
+}
+
+FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
+
+
+@lru_cache(maxsize=1)
+def get_video_reader_backend() -> str:
+    if FORCE_QWENVL_VIDEO_READER is not None:
+        video_reader_backend = FORCE_QWENVL_VIDEO_READER
+    elif is_decord_available():
+        video_reader_backend = "decord"
+    else:
+        video_reader_backend = "torchvision"
+    # print(f"qwen-vl-utils using {video_reader_backend} to read video.",
+    #       file=sys.stderr)
+    return video_reader_backend
+
+
+def fetch_video(
+    ele: dict, image_factor: int = IMAGE_FACTOR
+) -> torch.Tensor | list[Image.Image]:
+    if isinstance(ele["video"], str):
+        video_reader_backend = get_video_reader_backend()
+        video, total_dur, nframes = VIDEO_READER_BACKENDS[video_reader_backend](ele)
+        frame_timestamps = total_dur * torch.arange(1, nframes + 1) / nframes
+        grid_timestamps = frame_timestamps[::FRAME_FACTOR]
+        second_per_grid = grid_timestamps[1] - grid_timestamps[0]
+        nframes, _, height, width = video.shape
+        min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
+        total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
+        max_pixels = max(
+            min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR),
+            int(min_pixels * 1.05),
+        )
+        max_pixels = ele.get("max_pixels", max_pixels)
+        # min_pixels = (factor ** 2) * 52
+        # max_pixels = (factor ** 2) * min(768, (16384 / nframes * temporal_patch_size))
+        if "resized_height" in ele and "resized_width" in ele:
+            resized_height, resized_width = smart_resize(
+                ele["resized_height"],
+                ele["resized_width"],
+                factor=image_factor,
+            )
+        else:
+            resized_height, resized_width = smart_resize(
+                height,
+                width,
+                factor=image_factor,
+                min_pixels=min_pixels,
+                max_pixels=max_pixels,
+            )
+        video = transforms.functional.resize(
+            video,
+            [resized_height, resized_width],
+            interpolation=InterpolationMode.BICUBIC,
+            antialias=True,
+        ).float()
+        return video, total_dur, nframes, second_per_grid
+    else:
+        assert isinstance(ele["video"], (list, tuple))
+        process_info = ele.copy()
+        process_info.pop("type", None)
+        process_info.pop("video", None)
+        images = [
+            fetch_image(
+                {"image": video_element, **process_info}, size_factor=image_factor
+            )
+            for video_element in ele["video"]
+        ]
+        nframes = ceil_by_factor(len(images), FRAME_FACTOR)
+        if len(images) < nframes:
+            images.extend([images[-1]] * (nframes - len(images)))
+        return images, None, None, None
+
+
+def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]:
+    vision_infos = []
+    if isinstance(conversations[0], dict):
+        conversations = [conversations]
+    for conversation in conversations:
+        for message in conversation:
+            if isinstance(message["content"], list):
+                for ele in message["content"]:
+                    if (
+                        "image" in ele
+                        or "image_url" in ele
+                        or "video" in ele
+                        or ele["type"] in ("image", "image_url", "video")
+                    ):
+                        vision_infos.append(ele)
+    return vision_infos
+
+
+def process_vision_info(
+    conversations: list[dict] | list[list[dict]],
+) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None]:
+    vision_infos = extract_vision_info(conversations)
+    # Read images or videos
+    image_inputs = []
+    video_inputs = []
+    for vision_info in vision_infos:
+        if "image" in vision_info or "image_url" in vision_info:
+            image_inputs.append(fetch_image(vision_info))
+        elif "video" in vision_info:
+            video_inputs.append(fetch_video(vision_info))
+        else:
+            raise ValueError("image, image_url or video should in content.")
+    if len(image_inputs) == 0:
+        image_inputs = None
+    if len(video_inputs) == 0:
+        video_inputs = None
+    return image_inputs, video_inputs

examples/offline_inference/qwen_2_5_omni/run_multiple_prompts.sh

@@ -0,0 +1,9 @@
+export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH
+python end2end.py --model Qwen/Qwen2.5-Omni-7B \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --txt-prompts top100.txt

examples/offline_inference/qwen_2_5_omni/run_single_prompt.sh

@@ -0,0 +1,9 @@
+export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH
+python end2end.py --model Qwen/Qwen2.5-Omni-7B \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."

examples/offline_inference/qwen_2_5_omni/utils.py

@@ -0,0 +1,324 @@
+import tempfile
+from typing import Dict, List, 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 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}, "
+                f"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 = not ("Qwen2_5OmniModel" 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("Unknown ele {}".format(ele))
+                # insert a audio after the video
+                message["content"].insert(
+                    index + 1,
+                    {
+                        "type": "audio",
+                        "audio": ele[audio_key],
+                    },
+                )
+                # no need to load the added audio again
+                index += 1
+            elif ele["type"] == "video" or ele["type"] == "video_url":
+                if "video_url" in ele:
+                    video_key = "video_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
+                        temp_video_file.write(urlopen(ele["video_url"]).read())
+                        temp_video_file_path = temp_video_file.name
+                        videos.append(fetch_and_read_video(args, temp_video_file))
+                        ele["video"] = temp_video_file_path
+                else:
+                    video_key = "video"
+                    videos.append(fetch_and_read_video(args, ele[video_key]))
+            elif ele["type"] == "image" or ele["type"] == "image_url":
+                images.append(fetch_image(ele))
+
+            # move to the next content
+            index += 1
+
+    prompt = processor.apply_chat_template(
+        messages,
+        tokenize=tokenize,
+        add_generation_prompt=True,
+        add_vision_id=True,
+    )
+
+    audios = audios if len(audios) > 0 else None
+    images = images if len(images) > 0 else None
+    videos = videos if len(videos) > 0 else None
+
+    multi_modal_data = {}
+    if audios:
+        multi_modal_data["audio"] = audios
+    if images:
+        multi_modal_data["image"] = images
+    if videos:
+        multi_modal_data["video"] = videos
+
+    if isinstance(prompt, list) and isinstance(prompt[0], (list, str)):
+        prompt = prompt[0]
+
+    if tokenize:
+        return OmniTokensPrompt(
+            prompt_token_ids=prompt,
+            multi_modal_data=multi_modal_data,
+        )
+    else:
+        return TextPrompt(
+            prompt=prompt,
+            multi_modal_data=multi_modal_data,
+        )
+
+
+def make_text_prompt(args, prompt):
+    messages = [
+        get_system_prompt(),
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": prompt},
+            ],
+        },
+    ]
+
+    prompt = make_inputs_qwen2_omni(args, messages, tokenize=args.tokenize)
+    return prompt
+
+
+def make_audio_in_video_v2_prompt(args):
+    messages = [
+        {
+            "role": "system",
+            "content": [
+                {
+                    "type": "text",
+                    "text": (
+                        "You are Qwen, a virtual human developed by the Qwen Team, "
+                        "Alibaba Group, capable of perceiving auditory and visual "
+                        "inputs, as well as generating text and speech."
+                    ),
+                }
+            ],
+        },
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "video_url",
+                    "video_url": (
+                        "https://qianwen-res.oss-cn-beijing.aliyuncs.com/"
+                        "Qwen2.5-Omni/draw_small.mp4"
+                    ),
+                },
+            ],
+        },
+    ]
+    prompt = make_inputs_qwen2_omni(
+        args,
+        messages,
+        use_audio_in_video=True,
+        tokenize=args.tokenize,
+    )
+    return prompt
+
+
+def make_omni_prompt(
+    args, prompt=None
+) -> Union[OmniTokensPrompt, List[OmniTokensPrompt]]:
+    if args.prompt_type == "text":
+        prompt = make_text_prompt(args, prompt)
+    elif args.prompt_type == "audio-in-video-v2":
+        prompt = make_audio_in_video_v2_prompt(args)
+    else:
+        raise ValueError(f"Unsupported prompt type: {args.prompt_type}")
+    return prompt

examples/online_serving/README.md

@@ -1,77 +0,0 @@
-# Online serving Example of vLLM-omni for Qwen2.5-omni
-
-## 🛠️ Installation
-
-Please refer to [README.md](../../README.md)
-
-## Run examples (Qwen2.5-omni)
-
-Launch the server
-```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
-```
-
-If you have custom stage configs file, launch the server with command below
-```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
-```
-
-Get into the example folder
-```bash
-cd examples/online_serving
-```
-
-Send request via python
-```bash
-python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
-```
-
-Send request via curl
-```bash
-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)

examples/online_serving/gradio_demo.py

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

examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py

@@ -1,228 +0,0 @@
-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"
-openai_api_base = "http://localhost:8091/v1"
-
-client = OpenAI(
-    # defaults to os.environ.get("OPENAI_API_KEY")
-    api_key=openai_api_key,
-    base_url=openai_api_base,
-)
-
-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",
-        "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 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
-        "top_k": -1,  # Disable top-k sampling
-        "max_tokens": 2048,
-        "seed": SEED,  # Fixed seed for sampling
-        "detokenize": True,
-        "repetition_penalty": 1.1,
-    }
-    talker_sampling_params = {
-        "temperature": 0.9,
-        "top_p": 0.8,
-        "top_k": 40,
-        "max_tokens": 2048,
-        "seed": SEED,  # Fixed seed for sampling
-        "detokenize": True,
-        "repetition_penalty": 1.05,
-        "stop_token_ids": [8294],
-    }
-    code2wav_sampling_params = {
-        "temperature": 0.0,  # Deterministic - no randomness
-        "top_p": 1.0,  # Disable nucleus sampling
-        "top_k": -1,  # Disable top-k sampling
-        "max_tokens": 2048,
-        "seed": SEED,  # Fixed seed for sampling
-        "detokenize": True,
-        "repetition_penalty": 1.1,
-    }
-
-    sampling_params_list = [
-        thinker_sampling_params,
-        talker_sampling_params,
-        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,
-        ],
-        model=model_name,
-        extra_body=extra_body,
-    )
-
-    count = 0
-    for choice in chat_completion.choices:
-        if choice.message.audio:
-            audio_data = base64.b64decode(choice.message.audio.data)
-            audio_file_path = f"audio_{count}.wav"
-            with open(audio_file_path, "wb") as f:
-                f.write(audio_data)
-            print(f"Audio saved to {audio_file_path}")
-            count += 1
-        elif choice.message.content:
-            print("Chat completion output from text:", choice.message.content)
-
-
-def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
-    parser.add_argument(
-        "--query-type",
-        "-q",
-        type=str,
-        default="mixed_modalities",
-        choices=query_map.keys(),
-        help="Query type.",
-    )
-
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    run_multimodal_generation(args)

examples/online_serving/run_curl_multimodal_generation.sh

@@ -1,187 +0,0 @@
-#!/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='{
-  "temperature": 0.0,
-  "top_p": 1.0,
-  "top_k": -1,
-  "max_tokens": 2048,
-  "seed": 42,
-  "detokenize": true,
-  "repetition_penalty": 1.1
-}'
-
-talker_sampling_params='{
-  "temperature": 0.9,
-  "top_p": 0.8,
-  "top_k": 40,
-  "max_tokens": 2048,
-  "seed": 42,
-  "detokenize": true,
-  "repetition_penalty": 1.05,
-  "stop_token_ids": [8294]
-}'
-
-code2wav_sampling_params='{
-  "temperature": 0.0,
-  "top_p": 1.0,
-  "top_k": -1,
-  "max_tokens": 2048,
-  "seed": 42,
-  "detokenize": true,
-  "repetition_penalty": 1.1
-}'
-# Above is optional, it has a default setting in stage_configs of the corresponding model.
-
-# Define URLs for assets
-MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
-WINNING_CALL_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/winning_call.ogg"
-CHERRY_BLOSSOM_IMAGE_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
-SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
-
-# Build user content and extra fields based on query type
-case "$QUERY_TYPE" in
-  text)
-    user_content='[
-      {
-        "type": "text",
-        "text": "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
-      }
-    ]'
-    sampling_params_list='[
-      '"$thinker_sampling_params"',
-      '"$talker_sampling_params"',
-      '"$code2wav_sampling_params"'
-    ]'
-    mm_processor_kwargs="{}"
-    ;;
-  mixed_modalities)
-    user_content='[
-        {
-          "type": "audio_url",
-          "audio_url": {
-            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
-          }
-        },
-        {
-          "type": "image_url",
-          "image_url": {
-            "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
-          }
-        },
-        {
-          "type": "video_url",
-          "video_url": {
-            "url": "'"$SAMPLE_VIDEO_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "What is recited in the audio? What is the content of this image? Why is this video funny?"
-        }
-      ]'
-    sampling_params_list='[
-      '"$thinker_sampling_params"',
-      '"$talker_sampling_params"',
-      '"$code2wav_sampling_params"'
-    ]'
-    mm_processor_kwargs="{}"
-    ;;
-  use_audio_in_video)
-    user_content='[
-        {
-          "type": "video_url",
-          "video_url": {
-            "url": "'"$SAMPLE_VIDEO_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "Describe the content of the video, then convert what the baby say into text."
-        }
-      ]'
-    sampling_params_list='[
-      '"$thinker_sampling_params"',
-      '"$talker_sampling_params"',
-      '"$code2wav_sampling_params"'
-    ]'
-    mm_processor_kwargs='{
-      "use_audio_in_video": true
-    }'
-    ;;
-  multi_audios)
-    user_content='[
-        {
-          "type": "audio_url",
-          "audio_url": {
-            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
-          }
-        },
-        {
-          "type": "audio_url",
-          "audio_url": {
-            "url": "'"$WINNING_CALL_AUDIO_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "Are these two audio clips the same?"
-        }
-      ]'
-    sampling_params_list='[
-      '"$thinker_sampling_params"',
-      '"$talker_sampling_params"',
-      '"$code2wav_sampling_params"'
-    ]'
-    mm_processor_kwargs="{}"
-    ;;
-esac
-
-echo "Running query type: $QUERY_TYPE"
-echo ""
-
-
-output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d @- <<EOF
-{
-  "model": "Qwen/Qwen2.5-Omni-7B",
-  "sampling_params_list": $sampling_params_list,
-  "mm_processor_kwargs": $mm_processor_kwargs,
-  "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": $user_content
-    }
-  ]
-}
-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')"

mkdocs.yml

@@ -1,180 +0,0 @@
-site_name: vLLM-omni
-site_description: Multi-modality models inference and serving with non-autoregressive structures
-site_author: vLLM-omni Team
-site_url: https://vllm-project.github.io/vllm-omni/
-
-repo_name: vllm-project/vllm-omni
-repo_url: https://github.com/vllm-project/vllm-omni
-edit_uri: edit/main/docs/
-
-# Copyright
-copyright: Copyright © 2025 vLLM-omni Team
-
-# Theme
-theme:
-  name: material
-  palette:
-    # Light mode
-    - scheme: default
-      primary: blue
-      accent: blue
-      toggle:
-        icon: material/brightness-7
-        name: Switch to dark mode
-    # Dark mode
-    - scheme: slate
-      primary: blue
-      accent: blue
-      toggle:
-        icon: material/brightness-4
-        name: Switch to light mode
-  features:
-    - content.action.edit
-    - content.code.copy
-    - content.tabs.link
-    - navigation.instant
-    - navigation.instant.progress
-    - navigation.tracking
-    - navigation.tabs
-    - navigation.tabs.sticky
-    - navigation.sections
-    - navigation.indexes
-    - navigation.top
-    - search.suggest
-    - search.highlight
-    - search.share
-    - content.code.annotate
-    - content.tabs
-    - content.tooltips
-    - toc.follow
-  icon:
-    repo: fontawesome/brands/github
-
-hooks:
-  - docs/mkdocs/hooks/generate_examples.py
-
-# Plugins
-plugins:
-  - meta
-  - search
-  # - awesome-nav  # Uncomment if installed
-  - glightbox
-  - git-revision-date-localized:
-      # exclude files
-      exclude:
-        - api/*
-        - user_guide/examples/**
-        - contributing/design_documents/api_design_template.md
-        - DOCS_GUIDE.md
-  - minify:
-      minify_html: true
-      minify_js: true
-      minify_css: true
-      cache_safe: true
-  - api-autonav:
-      modules: ["vllm_omni"]
-      api_root_uri: "api"
-      nav_item_prefix: ""  # No prefix in navigation tree (clean names)
-      show_full_namespace: false  # Show only module name, not full path
-      exclude:
-        - "re:vllm_omni\\._.*"  # Internal modules
-  - mkdocstrings:
-      handlers:
-        python:
-          paths: [vllm_omni]
-          options:
-            show_symbol_type_heading: true
-            show_symbol_type_toc: true
-            show_source: true
-            show_root_heading: true
-            show_root_toc_entry: true
-            show_root_full_path: false
-            show_object_full_path: false
-            show_if_no_docstring: true
-            show_signature_annotations: true
-            separate_signature: true
-            show_overloads: true
-            signature_crossrefs: true
-            filters: []
-            summary:
-              modules: true
-          inventories:
-            - https://docs.python.org/3/objects.inv
-            - https://typing-extensions.readthedocs.io/en/latest/objects.inv
-            - https://docs.aiohttp.org/en/stable/objects.inv
-            - https://pillow.readthedocs.io/en/stable/objects.inv
-            - https://numpy.org/doc/stable/objects.inv
-            - https://pytorch.org/docs/stable/objects.inv
-            - https://psutil.readthedocs.io/en/stable/objects.inv
-
-# Markdown extensions
-markdown_extensions:
-  - attr_list
-  - md_in_html
-  - admonition
-  - pymdownx.details
-  # For content tabs
-  - pymdownx.superfences
-  - pymdownx.tabbed:
-      slugify: !!python/object/apply:pymdownx.slugs.slugify
-        kwds:
-          case: lower
-      alternate_style: true
-  # For code highlighting
-  - pymdownx.highlight:
-      anchor_linenums: true
-      line_spans: __span
-      pygments_lang_class: true
-  - pymdownx.inlinehilite
-  - pymdownx.snippets
-  # For emoji and icons
-  - pymdownx.emoji:
-      emoji_index: !!python/name:material.extensions.emoji.twemoji
-      emoji_generator: !!python/name:material.extensions.emoji.to_svg
-  # For in page [TOC] (not sidebar)
-  - toc:
-      permalink: true
-  # For math rendering
-  - pymdownx.arithmatex:
-      generic: true
-  # For task lists
-  - pymdownx.tasklist:
-      custom_checkbox: true
-  # For tables
-  - tables
-
-# Navigation
-nav:
-  - Home: index.md
-  - User Guide:
-    - Getting Started:
-      - user_guide/getting_started/quickstart.md
-    - Examples:
-      - user_guide/examples/index.md
-      - Offline Inference:
-        - user_guide/examples/offline_inference/qwen_2_5_omni.md
-      - Online Serving:
-        - user_guide/examples/online_serving/index.md
-  - API Reference:
-      - Summary: api/overview.md
-  - Developer Guide:
-    - contributing/README.md
-    - Design Documents:
-      - contributing/design_documents/index.md
-      - contributing/design_documents/vllm_omni_design.md
-      - contributing/design_documents/architecture_overview.md
-      - contributing/design_documents/api_design_doc.md
-
-# Extra
-extra:
-  version:
-    provider: mike
-
-# Extra CSS
-extra_css:
-  - mkdocs/stylesheets/extra.css
-
-# Makes the url format end in .html rather than act as a dir
-# So index.md generates as index.html and is available under URL /index.html
-# https://www.mkdocs.org/user-guide/configuration/#use_directory_urls
-use_directory_urls: false

pyproject.toml

@@ -4,13 +4,13 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "vllm-omni"
-version = "0.11.0rc1"
+version = "0.1.0"
 description = "vLLM-omni: Multi-modality models inference and serving with non-autoregressive structures"
 readme = "README.md"
-requires-python = ">=3.9,<3.14"  # Align with vLLM v0.11 supported Python versions
+requires-python = ">=3.8"
 license = {text = "Apache-2.0"}
 authors = [
-    {name = "vLLM-omni Team", email = "vllm-project@gmail.com"}
+    {name = "vLLM-omni Team", email = "hsliuustc@gmail.com"}
 ]
 keywords = ["vllm", "multimodal", "diffusion", "transformer", "inference", "serving"]
 classifiers = [
@@ -19,108 +19,96 @@ classifiers = [
     "Intended Audience :: Science/Research",
     "License :: OSI Approved :: Apache Software License",
     "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
-    "Programming Language :: Python :: 3.12",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 
 dependencies = [
-    # Core runtime dependencies (required for actual usage)
-    "omegaconf>=2.3.0",
-    # "decord>=0.6.0", # install manually
-    "librosa>=0.11.0",
-    "resampy>=0.4.3",
-    # "vllm==0.11.0",  # TODO: fix the entrypoints overwrite problem
+    "vllm>=0.2.0",
+    "torch>=2.0.0",
+    "transformers>=4.30.0",
+    "numpy>=1.21.0",
+    "pillow>=8.0.0",
+    "opencv-python>=4.5.0",
+    "PyYAML>=6.0",
+    "gradio>=3.0.0",
+    "fastapi>=0.100.0",
+    "uvicorn>=0.20.0",
+    "pydantic>=2.0.0",
+    "asyncio-mqtt>=0.11.0",
+    "ray>=2.0.0",
 ]
 
 [project.optional-dependencies]
-
 dev = [
     "pytest>=7.0.0",
     "pytest-asyncio>=0.21.0",
     "pytest-cov>=4.0.0",
     "black>=23.0.0",
     "isort>=5.12.0",
+    "flake8>=6.0.0",
     "mypy>=1.0.0",
     "pre-commit>=3.0.0",
-    "gradio>=5.49.1",
-    "soundfile>=0.13.1"
 ]
-
 docs = [
-    "mkdocs>=1.5.0",
-    "mkdocs-api-autonav",
-    "mkdocs-material",
-    "mkdocstrings-python",
-    "mkdocs-gen-files",
-    # "mkdocs-awesome-nav",  # Not available on PyPI
-    "mkdocs-glightbox",
-    "mkdocs-git-revision-date-localized-plugin",
-    "mkdocs-minify-plugin",
-    "regex",
-    "pydantic",
+    "sphinx>=6.0.0",
+    "sphinx-rtd-theme>=1.2.0",
+    "myst-parser>=1.0.0",
+]
+examples = [
+    "jupyter>=1.0.0",
+    "matplotlib>=3.5.0",
+    "seaborn>=0.11.0",
 ]
-
 
 [project.urls]
-Homepage = "https://github.com/vllm-project/vllm-omni"
-Repository = "https://github.com/vllm-project/vllm-omni"
+Homepage = "https://github.com/hsliuustc0106/vllm-omni"
+Repository = "https://github.com/hsliuustc0106/vllm-omni"
 Documentation = "https://vllm-omni.readthedocs.io"
-"Bug Tracker" = "https://github.com/vllm-project/vllm-omni/issues"
+"Bug Tracker" = "https://github.com/hsliuustc0106/vllm-omni/issues"
 
 [project.scripts]
 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.black]
+line-length = 88
+target-version = ['py38', 'py39', 'py310', 'py311']
+include = '\.pyi?$'
+extend-exclude = '''
+/(
+  # directories
+  \.eggs
+  | \.git
+  | \.hg
+  | \.mypy_cache
+  | \.tox
+  | \.venv
+  | build
+  | dist
+)/
+'''
 
-[tool.ruff]
-line-length = 120
-exclude = [
-    ".eggs",
-    ".git",
-    ".hg",
-    ".mypy_cache",
-    ".tox",
-    ".venv",
-    "build",
-    "dist",
-    "vllm_omni.egg-info",
-]
+[tool.isort]
+profile = "black"
+multi_line_output = 3
+line_length = 88
+known_first_party = ["vllm_omni"]
 
-[tool.ruff.lint]
-select = [
-    "E",  # pycodestyle errors
-    "W",  # pycodestyle warnings
-    "F",  # pyflakes
-    "I",  # isort (handled separately, but included for compatibility)
-    "N",  # pep8-naming
-    "UP", # pyupgrade
-]
-ignore = [
-    "E203",  # whitespace before ':' (conflicts with black)
-    # W503 is not needed in ruff as it's compatible with black by default
-    "N801",  # class names should use CapWords convention
-    "N802",  # function name should follow snake_case
-    "N806",  # variable in function should follow snake_case
-    "N812",  # lowercase imported as non-lowercase: functional as F
-]
-
-[tool.ruff.lint.per-file-ignores]
-"examples/**" = ["E501"]  # Allow long lines in examples
-"tests/**" = ["E501"]  # Allow long lines in tests
+[tool.flake8]
+max-line-length = 88
+extend-ignore = ["E203", "W503"]
 
 [tool.mypy]
-python_version = "3.12, 3.13"
+python_version = "3.8"
 warn_return_any = true
 warn_unused_configs = true
 disallow_untyped_defs = true
@@ -153,11 +141,3 @@ markers = [
     "benchmark: Benchmark tests",
     "slow: Slow tests",
 ]
-
-[tool.typos.default]
-extend-ignore-identifiers-re = [
-    ".*_thw",
-    ".*thw",
-    "ein",
-    ".*arange"
-]

requirements.txt

@@ -0,0 +1,4 @@
+omegaconf==2.3.0
+decord==0.6.0
+librosa==0.11.0
+resampy==0.4.3

scripts/README.md

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

scripts/build_wheel.sh

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

scripts/test_serving.sh

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

vllm_omni/__init__.py

@@ -12,28 +12,16 @@ Architecture:
   processing
 """
 
-try:
-    from . import patch  # noqa: F401
-except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
-    if exc.name != "vllm":
-        raise
-    # Allow importing vllm_omni without vllm (e.g., documentation builds)
-    patch = None  # type: ignore
+from . import patch  # noqa: F401
 from .config import OmniModelConfig
-from .entrypoints.async_omni_llm import AsyncOmniLLM
 
 # Main entry points
 from .entrypoints.omni_llm import OmniLLM
 
-# from .version import __version__, __version_tuple__  # isort:skip
-
-
 __all__ = [
-    # "__version__",  # TODO: add __version__ back when it is implemented
-    # "__version_tuple__",  # TODO: add __version_tuple__ back when it is implemented
     # Main components
     "OmniLLM",
-    "AsyncOmniLLM",
+    # "AsyncOmniLLM", # TODO: add AsyncOmniLLM back when it is implemented
     # Configuration
     "OmniModelConfig",
     # All other components are available through their respective modules

vllm_omni/config/model.py

@@ -2,9 +2,9 @@ from typing import Optional
 
 from pydantic import ConfigDict
 from pydantic.dataclasses import dataclass
-from vllm.config import ModelConfig, config
 
 import vllm_omni.model_executor.models as me_models
+from vllm.config import ModelConfig, config
 
 
 @config

vllm_omni/core/sched/diffusion_scheduler.py

@@ -4,10 +4,15 @@ from typing import Optional
 
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.request_queue import create_request_queue
-from vllm.v1.core.sched.scheduler import EngineCoreOutputs, Request, RequestStatus, SchedulerOutput, SpecDecodingStats
+from vllm.v1.core.sched.scheduler import (
+    EngineCoreOutputs,
+    Request,
+    RequestStatus,
+    SchedulerOutput,
+    SpecDecodingStats,
+)
 from vllm.v1.core.sched.utils import remove_all
 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput
-
 from vllm_omni.core.sched.output import OmniNewRequestData
 from vllm_omni.core.sched.scheduler import OmniScheduler
 from vllm_omni.outputs import OmniModelRunnerOutput
@@ -92,13 +97,17 @@ class DiffusionScheduler(OmniScheduler):
         num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
         if self.running:
             any_request = self.running[0]
-            num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(
-                any_request, len(self.running)
+            num_common_prefix_blocks = (
+                self.kv_cache_manager.get_num_common_prefix_blocks(
+                    any_request, len(self.running)
+                )
             )
 
         # Assemble SchedulerOutput
         new_reqs_data = [
-            OmniNewRequestData.from_request(req, req_to_new_blocks[req.request_id].get_block_ids())
+            OmniNewRequestData.from_request(
+                req, req_to_new_blocks[req.request_id].get_block_ids()
+            )
             for req in scheduled_new_reqs
         ]
         cached_reqs_data = self._make_cached_request_data(
@@ -108,7 +117,9 @@ class DiffusionScheduler(OmniScheduler):
             scheduled_spec_decode_tokens,
             req_to_new_blocks,
         )
-        scheduled_requests = scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs
+        scheduled_requests = (
+            scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs
+        )
         structured_output_request_ids, grammar_bitmask = self.get_grammar_bitmask(
             scheduled_requests, scheduled_spec_decode_tokens
         )
@@ -166,7 +177,9 @@ class DiffusionScheduler(OmniScheduler):
 
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: Optional[SpecDecodingStats] = None
-        kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None
+        kv_connector_stats = (
+            kv_connector_output.kv_connector_stats if kv_connector_output else None
+        )
 
         # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
         # the below loop can be a performance bottleneck. We should do our best
@@ -183,9 +196,13 @@ class DiffusionScheduler(OmniScheduler):
                 continue
 
             req_index = model_runner_output.req_id_to_index[req_id]
-            generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
+            generated_token_ids = (
+                sampled_token_ids[req_index] if sampled_token_ids else []
+            )
 
-            scheduled_spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id)
+            scheduled_spec_token_ids = (
+                scheduler_output.scheduled_spec_decode_tokens.get(req_id)
+            )
             if scheduled_spec_token_ids:
                 num_draft_tokens = len(scheduled_spec_token_ids)
                 num_accepted = len(generated_token_ids) - 1
@@ -222,7 +239,11 @@ class DiffusionScheduler(OmniScheduler):
                 stopped_preempted_reqs.add(request)
 
             # Extract sample logprobs if needed.
-            if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
+            if (
+                request.sampling_params is not None
+                and request.sampling_params.logprobs is not None
+                and logprobs
+            ):
                 # NOTE: once we support N tokens per step (spec decode),
                 # the outer lists can be of length > 1.
                 new_logprobs = logprobs.slice(req_index, req_index + 1)
@@ -231,7 +252,9 @@ class DiffusionScheduler(OmniScheduler):
                 # NOTE: structured_output_request should not be None if
                 # use_structured_output, we have check above, so safe to ignore
                 # type warning
-                request.structured_output_request.grammar.accept_tokens(req_id, new_token_ids)  # type: ignore[union-attr]  # noqa: E501
+                request.structured_output_request.grammar.accept_tokens(  # type: ignore[union-attr]  # noqa: E501
+                    req_id, new_token_ids
+                )
 
             # spec_token_ids comes from the model runner output
             if num_nans_in_logits is not None and req_id in num_nans_in_logits:
@@ -240,6 +263,7 @@ class DiffusionScheduler(OmniScheduler):
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
             if new_token_ids or pooler_output is not None or kv_transfer_params:
+
                 # Add EngineCoreOutput for this Request.
                 outputs[request.client_index].append(
                     EngineCoreOutput(
@@ -274,7 +298,10 @@ class DiffusionScheduler(OmniScheduler):
 
         # Create EngineCoreOutputs for all clients that have requests with
         # outputs in this step.
-        engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()}
+        engine_core_outputs = {
+            client_index: EngineCoreOutputs(outputs=outs)
+            for client_index, outs in outputs.items()
+        }
 
         finished_req_ids = self.finished_req_ids_dict
         if finished_req_ids:
@@ -285,10 +312,14 @@ class DiffusionScheduler(OmniScheduler):
                 if (eco := engine_core_outputs.get(client_index)) is not None:
                     eco.finished_requests = finished_set
                 else:
-                    engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
+                    engine_core_outputs[client_index] = EngineCoreOutputs(
+                        finished_requests=finished_set
+                    )
             finished_req_ids.clear()
 
-        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None:
+        if (
+            stats := self.make_stats(spec_decoding_stats, kv_connector_stats)
+        ) is not None:
             # Return stats to only one of the front-ends.
             if (eco := next(iter(engine_core_outputs.values()), None)) is None:
                 # We must return the stats even if there are no request

vllm_omni/core/sched/output.py

@@ -3,7 +3,6 @@ from typing import Optional
 
 from vllm.v1.core.sched.output import NewRequestData
 from vllm.v1.request import Request
-
 from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
 
 

vllm_omni/core/sched/scheduler.py

@@ -39,14 +39,22 @@ class OmniScheduler(VLLMScheduler):
                     num_computed_tokens=nr.num_computed_tokens,
                     lora_request=nr.lora_request,
                     # Enrich with omni payloads from the live request object
-                    prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None),
-                    additional_information=(getattr(request, "additional_information", None) if request else None),
+                    prompt_embeds=(
+                        getattr(request, "prompt_embeds", None) if request else None
+                    ),
+                    additional_information=(
+                        getattr(request, "additional_information", None)
+                        if request
+                        else None
+                    ),
                 )
                 new_list.append(omni_nr)
 
             scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
         except Exception:
             # If anything goes wrong, leave the original output unchanged
-            init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
+            init_logger(__name__).exception(
+                "Failed to wrap scheduled_new_reqs with OmniNewRequestData"
+            )
 
         return scheduler_output

vllm_omni/engine/__init__.py

@@ -5,6 +5,7 @@ Engine components for vLLM-omni.
 from typing import Any, Optional
 
 import msgspec
+
 from vllm.v1.engine import EngineCoreRequest
 
 

vllm_omni/engine/arg_utils.py

@@ -2,8 +2,6 @@ from dataclasses import dataclass
 from typing import Optional
 
 from vllm.engine.arg_utils import EngineArgs
-from vllm.v1.engine.async_llm import AsyncEngineArgs
-
 from vllm_omni.config import OmniModelConfig
 
 
@@ -33,31 +31,3 @@ class OmniEngineArgs(EngineArgs):
         omni_config.hf_config.architectures = omni_config.architectures
 
         return omni_config
-
-
-@dataclass
-class AsyncOmniEngineArgs(AsyncEngineArgs):
-    stage_id: int = 0
-    model_stage: str = "thinker"
-    model_arch: str = "Qwen2_5OmniForConditionalGeneration"
-    engine_output_type: Optional[str] = None
-
-    def create_model_config(self) -> OmniModelConfig:
-        # First, get the base ModelConfig from the parent class
-        base_config = super().create_model_config()
-
-        # Create OmniModelConfig by copying all base config attributes
-        # and adding the new omni-specific fields
-        config_dict = base_config.__dict__.copy()
-
-        # Add the new omni-specific fields
-        config_dict["stage_id"] = self.stage_id
-        config_dict["model_stage"] = self.model_stage
-        config_dict["model_arch"] = self.model_arch
-        config_dict["engine_output_type"] = self.engine_output_type
-
-        # Create and return the OmniModelConfig instance
-        omni_config = OmniModelConfig(**config_dict)
-        omni_config.hf_config.architectures = omni_config.architectures
-
-        return omni_config

vllm_omni/engine/output_processor.py

@@ -1,6 +1,7 @@
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Dict, Optional, Union
 
 import torch
+
 from vllm.logger import init_logger
 from vllm.outputs import PoolingRequestOutput
 from vllm.sampling_params import RequestOutputKind
@@ -9,16 +10,20 @@ from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
 from vllm.v1.engine.detokenizer import IncrementalDetokenizer
 from vllm.v1.engine.logprobs import LogprobsProcessor
 from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
-from vllm.v1.engine.output_processor import OutputProcessorOutput, RequestOutputCollector, RequestState
+from vllm.v1.engine.output_processor import (
+    OutputProcessorOutput,
+    RequestOutputCollector,
+    RequestState,
+)
 from vllm.v1.engine.parallel_sampling import ParentRequest
 from vllm.v1.metrics.stats import IterationStats
-
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
 
 
 class OmniRequestState(RequestState):
+
     def __init__(
         self,
         *args,
@@ -39,6 +44,7 @@ class OmniRequestState(RequestState):
         queue: Optional[Any],
         log_stats: bool,
     ) -> "OmniRequestState":
+
         if sampling_params := request.sampling_params:
             if not sampling_params.detokenize:
                 tokenizer = None
@@ -69,7 +75,9 @@ class OmniRequestState(RequestState):
             request_id=request.request_id,
             parent_req=parent_req,
             request_index=request_index,
-            lora_name=(request.lora_request.name if request.lora_request is not None else None),
+            lora_name=(
+                request.lora_request.name if request.lora_request is not None else None
+            ),
             output_kind=output_kind,
             prompt=prompt,
             prompt_token_ids=request.prompt_token_ids,
@@ -85,7 +93,9 @@ class OmniRequestState(RequestState):
             log_stats=log_stats,
         )
 
-    def add_multimodal_tensor(self, tensor: Optional[torch.Tensor], mm_type: Optional[str]) -> None:
+    def add_multimodal_tensor(
+        self, tensor: Optional[torch.Tensor], mm_type: Optional[str]
+    ) -> None:
         if tensor is None:
             return
         try:
@@ -127,11 +137,15 @@ class OmniRequestState(RequestState):
         if self.parent_req is None:
             outputs = [output]
         else:
-            request_id, outputs, finished = self.parent_req.get_outputs(request_id, output)
+            request_id, outputs, finished = self.parent_req.get_outputs(
+                request_id, output
+            )
             if not outputs:
                 return None
 
-        return self._new_request_output(request_id, outputs, finished, kv_transfer_params)
+        return self._new_request_output(
+            request_id, outputs, finished, kv_transfer_params
+        )
 
     def _new_completion_output(
         self,
@@ -140,7 +154,9 @@ class OmniRequestState(RequestState):
         stop_reason: Optional[Union[int, str]],
     ) -> Any:
         # Reuse base text/logprobs logic, then annotate with pooling_result.
-        base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)
+        base_output = super()._new_completion_output(
+            token_ids, finish_reason, stop_reason
+        )
         try:
             if self.mm_accumulated is not None:
                 tensor = self.mm_accumulated
@@ -180,12 +196,14 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         engine_core_output_type: Optional[str] = None,
     ):
         super().__init__(tokenizer=tokenizer, log_stats=log_stats)
-        self.output_handlers: dict[str, Callable[[EngineCoreOutput], None]] = {}
-        self._reqid_to_mm_type: dict[str, str] = {}
+        self.output_handlers: Dict[str, Callable[[EngineCoreOutput], None]] = {}
+        self._reqid_to_mm_type: Dict[str, str] = {}
         self.request_states: dict[str, OmniRequestState] = {}
         self.engine_core_output_type = engine_core_output_type
 
-    def register_handler(self, modality: str, handler: Callable[[EngineCoreOutput], None]) -> None:
+    def register_handler(
+        self, modality: str, handler: Callable[[EngineCoreOutput], None]
+    ) -> None:
         self.output_handlers[modality.lower()] = handler
 
     def add_request(
@@ -237,7 +255,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                 continue
 
             # 1) Stats
-            self._update_stats_from_output(req_state, eco, engine_core_timestamp, iteration_stats)
+            self._update_stats_from_output(
+                req_state, eco, engine_core_timestamp, iteration_stats
+            )
 
             new_token_ids = eco.new_token_ids
             pooling_output = eco.pooling_output
@@ -250,7 +270,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             # 2) Detokenize and logprobs when text path
             assert req_state.detokenizer is not None
             assert req_state.logprobs_processor is not None
-            stop_string = req_state.detokenizer.update(new_token_ids, finish_reason == FinishReason.STOP)
+            stop_string = req_state.detokenizer.update(
+                new_token_ids, finish_reason == FinishReason.STOP
+            )
             if stop_string:
                 finish_reason = FinishReason.STOP
                 stop_reason = stop_string
@@ -258,8 +280,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
 
             # 2.5) Accumulate multimodal tensors in RequestState
             try:
-                mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
-                if pooling_output is not None and isinstance(req_state, OmniRequestState):
+                mm_type = (
+                    getattr(eco, "output_type", self.engine_core_output_type) or ""
+                ).lower()
+                if pooling_output is not None and isinstance(
+                    req_state, OmniRequestState
+                ):
                     req_state.add_multimodal_tensor(pooling_output, mm_type)
             except Exception:
                 logger.debug(
@@ -284,13 +310,18 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             if ro:
                 # Attach accumulated multimodal payload if any
                 try:
-                    if isinstance(req_state, OmniRequestState) and req_state.mm_accumulated is not None:
+                    if (
+                        isinstance(req_state, OmniRequestState)
+                        and req_state.mm_accumulated is not None
+                    ):
                         mm_key = req_state.mm_type or "latents"
                         if not hasattr(ro, "multimodal_output"):
                             setattr(ro, "multimodal_output", {})
                         ro.multimodal_output[mm_key] = req_state.mm_accumulated
                 except Exception:
-                    logger.exception("Error attaching multimodal payload in process_outputs")
+                    logger.exception(
+                        "Error attaching multimodal payload in process_outputs"
+                    )
                 if req_state.queue is not None:
                     req_state.queue.put(ro)
                 else:
@@ -304,7 +335,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                     self.parent_requests.pop(parent_req.request_id, None)
                 if not eco.finished:
                     reqs_to_abort.append(req_id)
-                self._update_stats_from_finished(req_state, finish_reason, iteration_stats)
+                self._update_stats_from_finished(
+                    req_state, finish_reason, iteration_stats
+                )
                 if self.tracer:
                     self.do_tracing(eco, req_state, iteration_stats)
                 # Cleanup per-request mm state
@@ -319,7 +352,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
 
     # ---- routing helpers ----
     def _route_and_normalize(self, eco: EngineCoreOutput) -> None:
-        output_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
+        output_type = (
+            getattr(eco, "output_type", self.engine_core_output_type) or ""
+        ).lower()
 
         # Custom handler first (if registered)
         if output_type in self.output_handlers:
@@ -350,7 +385,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
     def _process_image_output(self, eco: EngineCoreOutput) -> None:
         """Ensure image tensors are surfaced via pooling_output for vLLM."""
         if eco.pooling_output is None:
-            tensor = self._extract_from_multimodal_outputs(eco, keys=("image", "images", "pixel_values", "pixels"))
+            tensor = self._extract_from_multimodal_outputs(
+                eco, keys=("image", "images", "pixel_values", "pixels")
+            )
             if tensor is not None:
                 eco.pooling_output = tensor
 
@@ -377,7 +414,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
     def _process_latents_output(self, eco: EngineCoreOutput) -> None:
         """Ensure latent tensors are surfaced via pooling_output."""
         if eco.pooling_output is None:
-            tensor = self._extract_from_multimodal_outputs(eco, keys=("latent", "latents", "z", "posterior"))
+            tensor = self._extract_from_multimodal_outputs(
+                eco, keys=("latent", "latents", "z", "posterior")
+            )
             if tensor is not None:
                 eco.pooling_output = tensor
 
@@ -405,7 +444,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             except Exception:
                 pass
 
-    def _extract_from_multimodal_outputs(self, eco: EngineCoreOutput, keys: tuple[str, ...]) -> Optional[torch.Tensor]:
+    def _extract_from_multimodal_outputs(
+        self, eco: EngineCoreOutput, keys: tuple[str, ...]
+    ) -> Optional[torch.Tensor]:
         mm = getattr(eco, "multimodal_outputs", None)
         if not isinstance(mm, dict):
             return None

vllm_omni/engine/processor.py

@@ -3,10 +3,10 @@ from collections.abc import Mapping
 from typing import Any, Optional, Union
 
 import torch
+
 from vllm.config import VllmConfig
 from vllm.inputs import ProcessorInputs, PromptType
 from vllm.inputs.parse import split_enc_dec_inputs
-from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
 from vllm.multimodal.inputs import MultiModalFeatureSpec
@@ -17,7 +17,6 @@ from vllm.sampling_params import SamplingParams
 from vllm.transformers_utils.tokenizer import AnyTokenizer
 from vllm.utils import length_from_prompt_token_ids_or_embeds
 from vllm.v1.engine.processor import Processor
-
 from vllm_omni.engine import (
     AdditionalInformationEntry,
     AdditionalInformationPayload,
@@ -26,8 +25,6 @@ from vllm_omni.engine import (
 )
 from vllm_omni.inputs.preprocess import OmniInputPreprocessor
 
-logger = init_logger(__name__)
-
 
 class OmniProcessor(Processor):
     @staticmethod
@@ -78,14 +75,20 @@ class OmniProcessor(Processor):
         priority: int = 0,
         data_parallel_rank: Optional[int] = None,
     ) -> tuple[Optional[str], OmniEngineCoreRequest]:
+
         # TODO(woosuk): Support pooling models.
         # TODO(woosuk): Support encoder-decoder models.
         self._validate_lora(lora_request)
         self._validate_params(params)
 
         data_parallel_size = self.vllm_config.parallel_config.data_parallel_size
-        if data_parallel_rank is not None and not (0 <= data_parallel_rank < data_parallel_size):
-            raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {data_parallel_size}).")
+        if data_parallel_rank is not None and not (
+            0 <= data_parallel_rank < data_parallel_size
+        ):
+            raise ValueError(
+                f"data_parallel_rank {data_parallel_rank} "
+                f"is out of range [0, {data_parallel_size})."
+            )
 
         if arrival_time is None:
             arrival_time = time.time()
@@ -136,9 +139,19 @@ class OmniProcessor(Processor):
         # discriminated unions of TypedDicts, because of how it handles
         # inheritance of TypedDict. If we explicitly extract the items we want
         # we can avoid type errors from using `dict.get` later in the method.
-        prompt_str: Optional[str] = None if decoder_inputs["type"] == "embeds" else decoder_inputs.get("prompt")
-        prompt_token_ids = decoder_inputs["prompt_token_ids"] if decoder_inputs["type"] != "embeds" else None
-        prompt_embeds = decoder_inputs["prompt_embeds"] if decoder_inputs["type"] == "embeds" else None
+        prompt_str: Optional[str] = (
+            None if decoder_inputs["type"] == "embeds" else decoder_inputs.get("prompt")
+        )
+        prompt_token_ids = (
+            decoder_inputs["prompt_token_ids"]
+            if decoder_inputs["type"] != "embeds"
+            else None
+        )
+        prompt_embeds = (
+            decoder_inputs["prompt_embeds"]
+            if decoder_inputs["type"] == "embeds"
+            else None
+        )
 
         sampling_params = None
         pooling_params = None
@@ -147,9 +160,13 @@ class OmniProcessor(Processor):
             sampling_params = params.clone()
             # If unset max tokens, then generate up to the max_model_len.
             if sampling_params.max_tokens is None:
-                seq_len = length_from_prompt_token_ids_or_embeds(prompt_token_ids, prompt_embeds)
+                seq_len = length_from_prompt_token_ids_or_embeds(
+                    prompt_token_ids, prompt_embeds
+                )
                 sampling_params.max_tokens = self.model_config.max_model_len - seq_len
-            sampling_params.update_from_generation_config(self.generation_config_fields, eos_token_id)
+            sampling_params.update_from_generation_config(
+                self.generation_config_fields, eos_token_id
+            )
             if self.tokenizer is not None:
                 sampling_params.update_from_tokenizer(self.tokenizer)
         else:
@@ -184,9 +201,12 @@ class OmniProcessor(Processor):
         prompt_embeds_payload: Optional[PromptEmbedsPayload] = None
         additional_information_payload: Optional[AdditionalInformationPayload] = None
         if "prompt_embeds" in decoder_inputs:  # type: ignore[operator]
+
             pe: torch.Tensor = decoder_inputs["prompt_embeds"]  # type: ignore[index]
             if pe.ndim != 2:
-                raise ValueError("prompt_embeds must be of shape (seq_len, hidden_size)")
+                raise ValueError(
+                    "prompt_embeds must be of shape (seq_len, hidden_size)"
+                )
             # Move to CPU and ensure contiguous memory for stable serialization
             pe_cpu = pe.detach().to("cpu").contiguous()
             seq_len, hidden_size = pe_cpu.shape
@@ -213,9 +233,13 @@ class OmniProcessor(Processor):
                 elif isinstance(value, list):
                     entry = AdditionalInformationEntry(list_data=value)
                 else:
-                    raise ValueError("additional_information values must be Tensor or list")
+                    raise ValueError(
+                        "additional_information values must be Tensor or list"
+                    )
                 entries[key] = entry
-            additional_information_payload = AdditionalInformationPayload(entries=entries)
+            additional_information_payload = AdditionalInformationPayload(
+                entries=entries
+            )
 
         return prompt_str, OmniEngineCoreRequest(
             request_id=request_id,

vllm_omni/entrypoints/async_omni_llm.py

@@ -1,742 +0,0 @@
-import asyncio
-import multiprocessing as mp
-import os
-import socket
-import sys
-import time
-from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import Any, Optional, Union
-
-import torch
-
-# External library imports (vLLM)
-import vllm.envs as envs
-from vllm.config import VllmConfig
-from vllm.engine.protocol import EngineClient
-from vllm.inputs import PromptType
-from vllm.inputs.preprocess import InputPreprocessor
-from vllm.logger import init_logger
-from vllm.lora.request import LoRARequest
-from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
-from vllm.sampling_params import SamplingParams
-from vllm.tracing import init_tracer
-from vllm.transformers_utils.config import maybe_register_config_serialize_by_value
-from vllm.transformers_utils.tokenizer import AnyTokenizer, init_tokenizer_from_configs
-from vllm.usage.usage_lib import UsageContext
-from vllm.utils import Device, deprecate_kwargs
-from vllm.v1.engine.async_llm import AsyncLLM
-from vllm.v1.engine.core_client import EngineCoreClient
-from vllm.v1.engine.exceptions import EngineDeadError
-from vllm.v1.executor.abstract import Executor
-from vllm.v1.metrics.loggers import StatLoggerFactory, StatLoggerManager
-
-# Internal imports (our code)
-from vllm_omni.config import OmniModelConfig
-from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
-from vllm_omni.engine.output_processor import MultimodalOutputProcessor
-from vllm_omni.engine.processor import OmniProcessor
-from vllm_omni.entrypoints.log_utils import (
-    OrchestratorMetrics,
-    configure_orchestrator_logger,
-    init_stats_paths,
-    remove_old_logs,
-)
-from vllm_omni.entrypoints.omni_stage import OmniStage
-from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
-from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
-from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
-from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
-from vllm_omni.outputs import OmniRequestOutput
-
-logger = init_logger(__name__)
-
-
-class AsyncOmniLLM(EngineClient):
-    def __init__(
-        self,
-        model: str,
-        stage_configs_path: Optional[str] = None,
-        log_stats: bool = False,
-        log_file: Optional[str] = None,
-        init_sleep_seconds: int = 20,
-        shm_threshold_bytes: int = 65536,
-        batch_timeout: int = 10,
-        init_timeout: int = 300,
-        **kwargs,
-    ):
-        self.batch_timeout = batch_timeout
-        self._enable_stats: bool = bool(log_stats)
-
-        if stage_configs_path is None:
-            self.stage_configs = load_stage_configs_from_model(model)
-        else:
-            self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
-
-        self.stage_list: list[OmniStage] = []
-        self.default_sampling_params_list: list[SamplingParams] = []
-        # Optional file handler for orchestrator
-        self._log_file = log_file
-        if self._log_file:
-            remove_old_logs(self._log_file, len(self.stage_list))
-            configure_orchestrator_logger(logger, self._log_file)
-
-        self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
-        self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout)
-
-    def _initialize_stages(
-        self,
-        model: str,
-        init_sleep_seconds: int,
-        shm_threshold_bytes: int,
-        init_timeout: int,
-    ) -> None:
-        self.stage_list: list[OmniStage] = []
-
-        # Build OmniStage instances in parallel, preserve original order
-        def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]:
-            idx, cfg = idx_cfg
-            return idx, OmniStage(cfg)
-
-        with ThreadPoolExecutor(max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))) as executor:
-            futures = [executor.submit(_build_stage, (idx, cfg)) for idx, cfg in enumerate(self.stage_configs)]
-            results: list[tuple[int, OmniStage]] = []
-            for fut in as_completed(futures):
-                results.append(fut.result())
-        results.sort(key=lambda x: x[0])
-        self.stage_list = [st for _, st in results]
-        self.default_sampling_params_list = [st.default_sampling_params for st in self.stage_list]
-        logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
-
-        self._ctx = mp.get_context("spawn")
-        self._stage_in_queues: list[mp.Queue] = []
-        self._stage_out_queues: list[mp.Queue] = []
-        self._init_sleep_seconds = max(0, int(init_sleep_seconds))
-        self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
-        self._start_stage_processes(model)
-        # Wait for all stages to report readiness before seeding
-        self._stages_ready: set[int] = set()
-        self._wait_for_stages_ready(timeout=init_timeout)
-
-    def _start_stage_processes(self, model: str) -> None:
-        for stage_id, stage in enumerate(self.stage_list):
-            # Use unbounded queues to avoid deadlock when seeding many requests
-            in_q: mp.Queue = self._ctx.Queue(maxsize=0)
-            out_q: mp.Queue = self._ctx.Queue(maxsize=0)
-            self._stage_in_queues.append(in_q)
-            self._stage_out_queues.append(out_q)
-
-            # Attach queues and start Stage-owned worker process
-            stage.attach_queues(in_q, out_q)
-            stage.init_stage_worker(
-                model,
-                is_async=True,
-                log_file=self._log_file,
-                shm_threshold_bytes=self._shm_threshold_bytes,
-                ctx=self._ctx,
-                batch_timeout=self.batch_timeout,
-            )
-            logger.debug("[Orchestrator] Stage-%s process started", stage_id)
-            time.sleep(self._init_sleep_seconds)
-
-    def close(self) -> None:
-        for q in self._stage_in_queues:
-            try:
-                q.put_nowait(None)
-            except Exception as e:
-                logger.warning(
-                    "[Orchestrator] Failed to send shutdown signal to \
-                        stage input queue: %s",
-                    e,
-                )
-        for stage in self.stage_list:
-            try:
-                stage.stop_stage_worker()
-            except Exception as e:
-                logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
-
-    def __del__(self) -> None:  # best-effort
-        try:
-            self.close()
-        except Exception as e:
-            logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
-
-    def shutdown(self):
-        """Shutdown, cleaning up the background proc and IPC."""
-        try:
-            self.close()
-        except Exception as e:
-            logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
-
-    async def generate(
-        self,
-        prompt: PromptType,
-        request_id: str,
-        sampling_params_list: Optional[Union[SamplingParams, Sequence[SamplingParams]]] = None,
-        lora_request: Optional[LoRARequest] = None,
-        trace_headers: Optional[Mapping[str, str]] = None,
-        priority: int = 0,
-        data_parallel_rank: Optional[int] = None,
-    ) -> AsyncGenerator[OmniRequestOutput, None]:
-        logger.debug("[Orchestrator] generate() called")
-        if sampling_params_list is None:
-            sampling_params_list = self.default_sampling_params_list
-        if len(sampling_params_list) != len(self.stage_list):
-            raise ValueError(
-                f"Expected {len(self.stage_list)} sampling params, \
-                got {len(sampling_params_list)}"
-            )
-
-        # Orchestrator keeps stage objects for input derivation
-        num_stages = len(self.stage_list)
-        # Track per-request start time for end-to-end timing
-        _req_start_ts: dict[int, float] = {}
-        _wall_start_ts: float = time.time()
-        # _last_finish_ts: float = _wall_start_ts
-
-        # Determine the final stage for E2E stats (highest stage_id with
-        # final_output=True; fallback to last stage)
-        final_stage_id_for_e2e = -1
-        try:
-            for _sid, _st in enumerate(self.stage_list):
-                if getattr(_st, "final_output", False):
-                    final_stage_id_for_e2e = max(final_stage_id_for_e2e, _sid)
-            if final_stage_id_for_e2e < 0:
-                final_stage_id_for_e2e = len(self.stage_list) - 1
-        except Exception as e:
-            logger.debug(
-                "[Orchestrator] Failed to determine final stage for E2E; \
-                    falling back to last: %s",
-                e,
-                exc_info=True,
-            )
-            final_stage_id_for_e2e = len(self.stage_list) - 1
-        # Metrics/aggregation helper
-        metrics = OrchestratorMetrics(
-            num_stages,
-            self._enable_stats,
-            self._stats_file,
-            self._overall_stats_file,
-            _wall_start_ts,
-        )
-        # Seed stage-0 queue with all requests
-        logger.debug("[Orchestrator] Seeding request into stage-0")
-        # Mark first input time for stage-0
-        metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
-
-        sp0: SamplingParams = sampling_params_list[0]  # type: ignore[index]
-        task = {
-            "request_id": request_id,
-            "engine_inputs": prompt,
-            "sampling_params": sp0,
-        }
-        self.stage_list[0].submit(task)
-        _req_start_ts[request_id] = time.time()
-        logger.debug("[Orchestrator] Enqueued request %s to stage-0", request_id)
-
-        logger.debug("[Orchestrator] Entering scheduling loop: stages=%d", num_stages)
-        finished = False
-        while not finished:
-            made_progress = False
-            for stage_id, stage in enumerate(self.stage_list):
-                result = stage.try_collect()
-                if result is None:
-                    continue
-
-                made_progress = True
-                req_id = result.get("request_id")
-                if "error" in result:
-                    logger.error(
-                        "Stage %s error on request %s: %s",
-                        stage_id,
-                        req_id,
-                        result["error"],
-                    )
-                    continue
-
-                if result.get("type") == "stage_ready":
-                    # Only happens when stage is initialized slower than expected,
-                    # so we wait for a short time and try again
-                    time.sleep(0.05)
-                    continue
-
-                engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
-                # Mark last output time for this stage whenever we receive outputs
-                metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
-                try:
-                    _m = result.get("metrics")
-                    if _m is not None:
-                        metrics.on_stage_metrics(stage_id, req_id, _m)
-                except Exception as e:
-                    logger.exception(
-                        "[Orchestrator] Failed to process metrics for stage %s, \
-                            req %s: %s",
-                        stage_id,
-                        req_id,
-                        e,
-                    )
-                logger.debug(
-                    "[Orchestrator] Stage-%s completed request %s; \
-                        forwarding or finalizing",
-                    stage_id,
-                    req_id,
-                )
-                stage.set_engine_outputs(engine_outputs)
-
-                if getattr(stage, "final_output", False):
-                    logger.debug(
-                        "[Orchestrator] Request %s finalized at stage-%s",
-                        req_id,
-                        stage_id,
-                    )
-
-                    # End-to-end timing and time-per-token for final output
-                    # (only once per request at the designated final stage)
-                    try:
-                        rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
-                        if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
-                            metrics.on_finalize_request(
-                                stage_id,
-                                req_id,
-                                engine_outputs,
-                                _req_start_ts.get(req_id, _wall_start_ts),
-                            )
-                    except Exception as e:
-                        logger.exception(
-                            "[Orchestrator] Finalize request handling error for \
-                                req %s at stage %s: %s",
-                            req_id,
-                            stage_id,
-                            e,
-                        )
-
-                    if isinstance(engine_outputs, list):
-                        engine_outputs = engine_outputs[0]
-                    yield OmniRequestOutput(
-                        stage_id=stage_id,
-                        final_output_type=stage.final_output_type,
-                        request_output=engine_outputs,
-                    )
-
-                next_stage_id = stage_id + 1
-                if next_stage_id < num_stages:
-                    next_stage: OmniStage = self.stage_list[next_stage_id]
-                    next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
-                    sp_next: SamplingParams = sampling_params_list[next_stage_id]
-                    try:
-                        # Measure transfer size and time (encode + enqueue)
-                        size_bytes = 0
-                        try:
-                            size_bytes = len(_set(next_inputs))
-                        except Exception:
-                            size_bytes = 0
-                        t0 = time.time()
-                        ipc_payload = _encode(
-                            next_inputs,
-                            getattr(self, "_shm_threshold_bytes", 65536),
-                            obj_key="engine_inputs",
-                            shm_key="engine_inputs_shm",
-                        )
-                        ipc_payload.update(
-                            {
-                                "request_id": req_id,
-                                "sampling_params": sp_next,
-                                "sent_ts": time.time(),
-                            }
-                        )
-                        self.stage_list[next_stage_id].submit(ipc_payload)
-                        t1 = time.time()
-                        tx_ms = (t1 - t0) * 1000.0
-                        metrics.on_forward(
-                            stage_id,
-                            next_stage_id,
-                            req_id,
-                            int(size_bytes),
-                            float(tx_ms),
-                            bool("engine_inputs_shm" in ipc_payload),
-                        )
-                    except Exception as e:
-                        logger.warning(
-                            "[Orchestrator] IPC encode failed for req %s: %s; \
-                                falling back to inline payload",
-                            req_id,
-                            e,
-                        )
-                        self.stage_list[next_stage_id].submit(
-                            {
-                                "request_id": req_id,
-                                "engine_inputs": next_inputs,
-                                "sampling_params": sp_next,
-                            }
-                        )
-                    logger.debug(
-                        "[Orchestrator] Forwarded request %s to stage-%s",
-                        req_id,
-                        next_stage_id,
-                    )
-                else:
-                    finished = True
-                    logger.debug("[Orchestrator] Request %s fully completed", req_id)
-
-            if not made_progress:
-                time.sleep(0.005)
-
-        logger.debug("[Orchestrator] All requests completed")
-
-        # Summarize and print stats
-        try:
-            summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
-            logger.info("[Summary] %s", summary)
-        except Exception as e:
-            logger.exception("[Orchestrator] Failed to build/log summary: %s", e)
-
-    def _wait_for_stages_ready(self, timeout: int = 120) -> None:
-        deadline = time.time() + max(0, int(timeout))
-        num_stages = len(self.stage_list)
-        while len(self._stages_ready) < num_stages and time.time() < deadline:
-            progressed = False
-            for stage_id, stage in enumerate(self.stage_list):
-                if stage_id in self._stages_ready:
-                    continue
-                result = stage.try_collect()
-                if result is None:
-                    continue
-                progressed = True
-                if result.get("type") == "stage_ready":
-                    self._stages_ready.add(stage_id)
-                    # Store vllm_config received from worker process
-                    vllm_config = result.get("vllm_config")
-                    if vllm_config is not None:
-                        stage.set_vllm_config(vllm_config)
-                    tokenizer = result.get("tokenizer")
-                    if tokenizer is not None:
-                        stage.set_tokenizer(tokenizer)
-                    # input_preprocessor = result.get("input_preprocessor")
-                    # if input_preprocessor is not None:
-                    #     stage.set_input_preprocessor(input_preprocessor)
-                    is_tracing_enabled = result.get("is_tracing_enabled")
-                    if is_tracing_enabled is not None:
-                        stage.set_is_tracing_enabled(is_tracing_enabled)
-                    logger.debug("[Orchestrator] Stage-%s reported ready", stage_id)
-                else:
-                    # No user data should arrive before seeding; ignore other messages
-                    pass
-            if not progressed:
-                time.sleep(0.01)
-        if len(self._stages_ready) < num_stages:
-            not_ready = sorted(set(range(num_stages)) - set(self._stages_ready))
-            logger.warning(
-                "[Orchestrator] Initialization timeout: only %s/%s stages are \
-                    ready; not ready: %s",
-                len(self._stages_ready),
-                num_stages,
-                not_ready,
-            )
-            # Provide actionable suggestions before shutdown
-            try:
-                suggestions = [
-                    "Verify GPU/device assignment in config (runtime.devices) is \
-                        correct.",
-                    "Check GPU/host memory availability; reduce model or batch size if needed.",  # noqa: E501
-                    "Check model weights path and network reachability (if loading remotely).",  # noqa: E501
-                    "Increase initialization wait time (init_sleep_seconds or \
-                        call-site timeout).",
-                ]
-                if getattr(self, "_log_file", None):
-                    suggestions.append(
-                        f"Inspect per-stage log files for details: \
-                            {self._log_file}.stage<id>.log"
-                    )
-                logger.error(
-                    "[Orchestrator] Stage initialization failed, shutting down. \
-                        Suggestions:\n- %s",
-                    "\n- ".join(suggestions),
-                )
-            except Exception:
-                # Best-effort logging of suggestions
-                logger.error(
-                    "[Orchestrator] Stage initialization failed and an error \
-                        occurred while logging suggestions",
-                )
-
-            # Attempt graceful shutdown of all stages before exiting
-            try:
-                self.close()
-            except Exception:
-                pass
-
-            # Terminate the current process with non-zero exit code
-            try:
-                sys.exit(1)
-            except SystemExit:
-                raise
-            except Exception:
-                os._exit(1)
-
-    @property
-    def is_running(self) -> bool:
-        # Is None before the loop is started.
-        return len(self._stage_in_queues) > 0
-
-    @property
-    def is_stopped(self) -> bool:
-        return self.errored
-
-    @property
-    def errored(self) -> bool:
-        return not self.is_running
-
-    @property
-    def dead_error(self) -> BaseException:
-        return EngineDeadError()
-
-    async def abort(self, request_id: Union[str, Iterable[str]]) -> None:
-        pass
-
-    async def get_vllm_config(self) -> VllmConfig:
-        for stage in self.stage_list:
-            if stage.is_comprehension:
-                # Use the vllm_config received from worker process
-                if stage.vllm_config is not None:
-                    return stage.vllm_config
-        return None
-
-    async def get_model_config(self) -> OmniModelConfig:
-        for stage in self.stage_list:
-            if stage.is_comprehension:
-                # Use the vllm_config received from worker process
-                if stage.vllm_config is not None:
-                    return stage.vllm_config.model_config
-        return None
-
-    async def get_input_preprocessor(self) -> InputPreprocessor:
-        # for stage in self.stage_list:
-        #     if stage.is_comprehension:
-        #         return stage.input_preprocessor
-        return None
-
-    async def get_tokenizer(self) -> AnyTokenizer:
-        for stage in self.stage_list:
-            if stage.is_comprehension:
-                return stage.tokenizer
-        return None
-
-    async def is_tracing_enabled(self) -> bool:
-        for stage in self.stage_list:
-            if stage.is_comprehension:
-                return stage.is_tracing_enabled
-        return False
-
-    async def do_log_stats(self) -> None:
-        pass
-
-    async def check_health(self) -> None:
-        pass
-
-    async def reset_mm_cache(self) -> None:
-        pass
-
-    async def reset_prefix_cache(self, device: Optional[Device] = None) -> None:
-        pass
-
-    async def sleep(self, level: int = 1) -> None:
-        pass
-
-    async def wake_up(self, tags: Optional[list[str]] = None) -> None:
-        pass
-
-    async def is_sleeping(self) -> bool:
-        """Check whether the engine is sleeping"""
-        return False
-
-    async def add_lora(self, lora_request: LoRARequest) -> bool:
-        """Load a new LoRA adapter into the engine for future requests."""
-        return False
-
-    async def encode(
-        self,
-        *args,
-        **kwargs,
-    ):
-        """Generate outputs for a request from a pooling model."""
-        raise NotImplementedError("encode() is not implemented for AsyncOmniLLM")
-
-    async def start_profile(self) -> None:
-        raise NotImplementedError("start_profile() is not implemented for AsyncOmniLLM")
-
-    async def stop_profile(self) -> None:
-        raise NotImplementedError("stop_profile() is not implemented for AsyncOmniLLM")
-
-
-class AsyncOmniStageLLM(AsyncLLM):
-    def __init__(
-        self,
-        engine_args: AsyncOmniEngineArgs,
-        vllm_config: VllmConfig,
-        executor_class: type[Executor],
-        log_stats: bool,
-        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
-        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
-        use_cached_outputs: bool = False,
-        log_requests: bool = True,
-        start_engine_loop: bool = True,
-        stat_loggers: Optional[list[StatLoggerFactory]] = None,
-        client_addresses: Optional[dict[str, str]] = None,
-        client_count: int = 1,
-        client_index: int = 0,
-    ) -> None:
-        """
-        Create an AsyncLLM.
-
-        Args:
-            vllm_config: global configuration.
-            executor_class: an Executor impl, e.g. MultiprocExecutor.
-            log_stats: Whether to log stats.
-            usage_context: Usage context of the LLM.
-            mm_registry: Multi-modal registry.
-            use_cached_outputs: Whether to use cached outputs.
-            log_requests: Whether to log requests.
-            start_engine_loop: Whether to start the engine loop.
-            stat_loggers: customized stat loggers for the engine.
-                If not provided, default stat loggers will be used.
-                PLEASE BE AWARE THAT STAT LOGGER IS NOT STABLE
-                IN V1, AND ITS BASE CLASS INTERFACE MIGHT CHANGE.
-
-        Returns:
-            None
-        """
-        if not envs.VLLM_USE_V1:
-            raise ValueError(
-                "Using V1 AsyncLLMEngine, but envs.VLLM_USE_V1=False. "
-                "This should not happen. As a workaround, try using "
-                "AsyncLLMEngine.from_vllm_config(...) or explicitly set "
-                "VLLM_USE_V1=0 or 1 and report this issue on Github."
-            )
-
-        # Ensure we can serialize custom transformer configs
-        maybe_register_config_serialize_by_value()
-
-        self.model_config = vllm_config.model_config
-        self.vllm_config = vllm_config
-        self.observability_config = vllm_config.observability_config
-        self.log_requests = log_requests
-
-        self.log_stats = log_stats or (stat_loggers is not None)
-        if not log_stats and stat_loggers is not None:
-            logger.info(
-                "AsyncLLM created with log_stats=False and non-empty custom logger list; "
-                "enabling logging without default stat loggers"
-            )
-
-        if self.model_config.skip_tokenizer_init:
-            self.tokenizer = None
-        else:
-            # Tokenizer (+ ensure liveness if running in another process).
-            self.tokenizer = init_tokenizer_from_configs(model_config=vllm_config.model_config)
-
-        # Processor (converts Inputs --> EngineCoreRequests).
-        self.processor = OmniProcessor(
-            vllm_config=vllm_config,
-            tokenizer=self.tokenizer,
-            mm_registry=mm_registry,
-        )
-
-        # OutputProcessor (converts EngineCoreOutputs --> RequestOutput).
-        self.output_processor = MultimodalOutputProcessor(
-            tokenizer=self.tokenizer,
-            log_stats=self.log_stats,
-            engine_core_output_type=engine_args.engine_output_type,
-        )
-        if self.observability_config.otlp_traces_endpoint is not None:
-            tracer = init_tracer("vllm.llm_engine", self.observability_config.otlp_traces_endpoint)
-            self.output_processor.tracer = tracer
-
-        # EngineCore (starts the engine in background process).
-        self.engine_core = EngineCoreClient.make_async_mp_client(
-            vllm_config=vllm_config,
-            executor_class=executor_class,
-            log_stats=self.log_stats,
-            client_addresses=client_addresses,
-            client_count=client_count,
-            client_index=client_index,
-        )
-
-        # Loggers.
-        self.logger_manager: Optional[StatLoggerManager] = None
-        if self.log_stats:
-            self.logger_manager = StatLoggerManager(
-                vllm_config=vllm_config,
-                engine_idxs=self.engine_core.engine_ranks_managed,
-                custom_stat_loggers=stat_loggers,
-                enable_default_loggers=log_stats,
-                client_count=client_count,
-            )
-            self.logger_manager.log_engine_initialized()
-
-        self.output_handler: Optional[asyncio.Task] = None
-        try:
-            # Start output handler eagerly if we are in the asyncio eventloop.
-            asyncio.get_running_loop()
-            self._run_output_handler()
-        except RuntimeError:
-            pass
-
-        if envs.VLLM_TORCH_PROFILER_DIR:
-            logger.info(
-                "Torch profiler enabled. AsyncLLM CPU traces will be collected under %s",  # noqa: E501
-                envs.VLLM_TORCH_PROFILER_DIR,
-            )
-            worker_name = f"{socket.gethostname()}_{os.getpid()}.async_llm"
-            self.profiler = torch.profiler.profile(
-                activities=[
-                    torch.profiler.ProfilerActivity.CPU,
-                ],
-                with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
-                on_trace_ready=torch.profiler.tensorboard_trace_handler(
-                    envs.VLLM_TORCH_PROFILER_DIR, worker_name=worker_name, use_gzip=True
-                ),
-            )
-        else:
-            self.profiler = None
-
-    @classmethod
-    @deprecate_kwargs(
-        "disable_log_requests",
-        additional_message=("This argument will have no effect. Use `enable_log_requests` instead."),
-    )
-    def from_vllm_config(
-        cls,
-        vllm_config: VllmConfig,
-        engine_args: AsyncOmniEngineArgs,
-        start_engine_loop: bool = True,
-        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
-        stat_loggers: Optional[list[StatLoggerFactory]] = None,
-        enable_log_requests: bool = False,
-        disable_log_stats: bool = False,
-        client_addresses: Optional[dict[str, str]] = None,
-        client_count: int = 1,
-        client_index: int = 0,
-        disable_log_requests: bool = True,  # Deprecated, will be removed
-    ) -> "AsyncLLM":
-        if not envs.VLLM_USE_V1:
-            raise ValueError(
-                "Using V1 AsyncLLMEngine, but envs.VLLM_USE_V1=False. "
-                "This should not happen. As a workaround, try using "
-                "AsyncLLMEngine.from_vllm_config(...) or explicitly set "
-                "VLLM_USE_V1=0 or 1 and report this issue on Github."
-            )
-
-        # Create the LLMEngine.
-        return cls(
-            vllm_config=vllm_config,
-            executor_class=Executor.get_class(vllm_config),
-            start_engine_loop=start_engine_loop,
-            stat_loggers=stat_loggers,
-            log_requests=enable_log_requests,
-            log_stats=not disable_log_stats,
-            usage_context=usage_context,
-            client_addresses=client_addresses,
-            client_count=client_count,
-            client_index=client_index,
-            engine_args=engine_args,
-        )

vllm_omni/entrypoints/chat_utils.py

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

vllm_omni/entrypoints/cli/main.py

@@ -2,55 +2,24 @@
 CLI entry point for vLLM-omni that intercepts vLLM commands.
 """
 
-import importlib.metadata
 import sys
 
+from vllm_omni.entrypoints.cli.serve import OmniServeCommand
+
 
 def main():
     """Main CLI entry point that intercepts vLLM commands."""
     # Check if --omni flag is present
-    if "--omni" not in sys.argv:
+    if "--omni" in sys.argv:
+        # Remove --omni flag and "serve" command, process with vLLM-omni
+        omni_args = [arg for arg in sys.argv[1:] if arg not in ["--omni", "serve"]]
+        omni_serve = OmniServeCommand()
+        omni_serve.run(omni_args)
+    else:
+        # Forward to original vLLM CLI
         from vllm.entrypoints.cli.main import main as vllm_main
 
         vllm_main()
-        return
-    else:
-        from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup
-        from vllm.utils import FlexibleArgumentParser
-
-        import vllm_omni.entrypoints.cli.serve
-
-        CMD_MODULES = [
-            vllm_omni.entrypoints.cli.serve,
-        ]
-
-        cli_env_setup()
-
-        parser = FlexibleArgumentParser(
-            description="vLLM OMNI CLI",
-            epilog=VLLM_SUBCMD_PARSER_EPILOG.format(subcmd="[subcommand]"),
-        )
-        parser.add_argument(
-            "-v",
-            "--version",
-            action="version",
-            version=importlib.metadata.version("vllm_omni"),
-        )
-        subparsers = parser.add_subparsers(required=False, dest="subparser")
-        cmds = {}
-        for cmd_module in CMD_MODULES:
-            new_cmds = cmd_module.cmd_init()
-            for cmd in new_cmds:
-                cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
-                cmds[cmd.name] = cmd
-        args = parser.parse_args()
-        if args.subparser in cmds:
-            cmds[args.subparser].validate(args)
-
-        if hasattr(args, "dispatch_function"):
-            args.dispatch_function(args)
-        else:
-            parser.print_help()
 
 
 if __name__ == "__main__":

vllm_omni/entrypoints/cli/serve.py

@@ -3,65 +3,93 @@ Omni serve command for vLLM-omni.
 """
 
 import argparse
-
-import uvloop
-from vllm.entrypoints.cli.types import CLISubcommand
-from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
-from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
-from vllm.logger import init_logger
-from vllm.utils import FlexibleArgumentParser
-
-from vllm_omni.entrypoints.openai.api_server import omni_run_server
-
-logger = init_logger(__name__)
-
-DESCRIPTION = """Launch a local OpenAI-compatible API server to serve Omni LLM
-completions via HTTP. Defaults to Qwen/Qwen2.5-Omni-7B if no model is specified.
-
-Search by using: `--help=<ConfigGroup>` to explore options by section (e.g.,
---help=ModelConfig, --help=Frontend)
-  Use `--help=all` to show all available flags at once.
-"""
+from typing import List
 
 
-class OmniServeCommand(CLISubcommand):
-    """The `serve` subcommand for the vLLM CLI."""
+class OmniServeCommand:
+    """Command handler for vLLM-omni serve command."""
 
-    name = "serve"
+    def __init__(self):
+        self.parser = self._create_parser()
 
-    @staticmethod
-    def cmd(args: argparse.Namespace) -> None:
-        # If model is specified in CLI (as positional arg), it takes precedence
-        if hasattr(args, "model_tag") and args.model_tag is not None:
-            args.model = args.model_tag
-
-        uvloop.run(omni_run_server(args))
-
-    def validate(self, args: argparse.Namespace) -> None:
-        validate_parsed_serve_args(args)
-
-    def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
-        serve_parser = subparsers.add_parser(
-            self.name,
-            description=DESCRIPTION,
-            usage="vllm serve [model_tag] --omni [options]",
+    def _create_parser(self) -> argparse.ArgumentParser:
+        """Create argument parser for omni serve command."""
+        parser = argparse.ArgumentParser(
+            description="vLLM-omni: Multi-modality models inference and serving"
         )
 
-        serve_parser = make_arg_parser(serve_parser)
-        serve_parser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(subcmd=self.name)
-        serve_parser.add_argument(
-            "--omni",
-            action="store_true",
-            help="Enable vLLM-omni mode",
+        # Model arguments - make it optional with default
+        parser.add_argument(
+            "model",
+            nargs="?",
+            default="Qwen/Qwen3-0.6B",
+            help="Path to the model or model name (default: Qwen/Qwen3-0.6B)",
         )
-        serve_parser.add_argument(
-            "--stage-configs-path",
+
+        # Server arguments
+        parser.add_argument(
+            "--port",
+            type=int,
+            default=8000,
+            help="Port to run the server on",
+        )
+
+        parser.add_argument(
+            "--host",
             type=str,
-            default=None,
-            help="Path to the stage configs file. If not specified, the stage configs will be loaded from the model.",
+            default="0.0.0.0",
+            help="Host to run the server on",
         )
-        return serve_parser
 
+        # Stage configuration arguments
+        parser.add_argument("--ar-stage", type=str, help="AR stage model path")
 
-def cmd_init() -> list[CLISubcommand]:
-    return [OmniServeCommand()]
+        parser.add_argument("--dit-stage", type=str, help="DiT stage model path")
+
+        parser.add_argument(
+            "--dit-steps",
+            type=int,
+            default=50,
+            help="Number of DiT inference steps",
+        )
+
+        parser.add_argument(
+            "--dit-guidance-scale",
+            type=float,
+            default=7.5,
+            help="DiT guidance scale",
+        )
+
+        parser.add_argument(
+            "--use-diffusers",
+            action="store_true",
+            help="Use diffusers pipeline for DiT stage",
+        )
+
+        # Other arguments
+        parser.add_argument(
+            "--log-stats", action="store_true", help="Enable logging statistics"
+        )
+
+        return parser
+
+    def run(self, args: List[str]) -> None:
+        """Run the omni serve command."""
+        pass
+
+    def _create_stage_configs(self, args) -> List:
+        """Create stage configurations based on arguments."""
+        pass
+
+    async def _start_server(self, omni_llm, args) -> None:
+        """Start the API server."""
+        try:
+            # Import here to avoid circular imports
+            from vllm_omni.entrypoints.api_server import run_server
+
+            await run_server(omni_llm_instance=omni_llm, host=args.host, port=args.port)
+        except KeyboardInterrupt:
+            print("\nShutting down server...")
+        except Exception as e:
+            print(f"Error starting server: {e}")
+            raise

vllm_omni/entrypoints/log_utils.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
 
+from typing import Any, Dict, List, Optional
+
 import logging
 import os
 import time
-from typing import Any
 
 from vllm_omni.entrypoints.stage_utils import append_jsonl as _append_jsonl
 
 
-def remove_old_logs(log_file: str | None, num_stages: int) -> None:
+def remove_old_logs(log_file: Optional[str], num_stages: int) -> None:
     try:
         if not log_file:
             return
@@ -48,7 +49,7 @@ def remove_old_logs(log_file: str | None, num_stages: int) -> None:
         pass
 
 
-def configure_orchestrator_logger(logger: logging.Logger, log_file: str | None) -> None:
+def configure_orchestrator_logger(logger: logging.Logger, log_file: Optional[str]) -> None:
     try:
         if not log_file:
             return
@@ -63,9 +64,9 @@ def configure_orchestrator_logger(logger: logging.Logger, log_file: str | None)
         pass
 
 
-def init_stats_paths(enable_stats: bool, log_file: str | None) -> tuple[str | None, str | None]:
-    stats_file: str | None = None
-    overall_file: str | None = None
+def init_stats_paths(enable_stats: bool, log_file: Optional[str]) -> tuple[Optional[str], Optional[str]]:
+    stats_file: Optional[str] = None
+    overall_file: Optional[str] = None
     try:
         if enable_stats and log_file:
             stats_file = f"{log_file}.orchestrator.stats.jsonl"
@@ -76,7 +77,7 @@ def init_stats_paths(enable_stats: bool, log_file: str | None) -> tuple[str | No
     return stats_file, overall_file
 
 
-def _safe_append_jsonl(path: str | None, record: dict[str, Any]) -> None:
+def _safe_append_jsonl(path: Optional[str], record: Dict[str, Any]) -> None:
     if not path:
         return
     try:
@@ -85,204 +86,122 @@ def _safe_append_jsonl(path: str | None, record: dict[str, Any]) -> None:
         pass
 
 
-def log_transfer_tx(
-    stats_file: str | None,
-    from_stage: int,
-    to_stage: int,
-    request_id: Any,
-    size_bytes: int,
-    tx_time_ms: float,
-    used_shm: bool,
-) -> None:
-    _safe_append_jsonl(
-        stats_file,
-        {
-            "type": "transfer_stats",
-            "from_stage": from_stage,
-            "to_stage": to_stage,
-            "request_id": request_id,
-            "size_bytes": int(size_bytes),
-            "tx_time_ms": float(tx_time_ms),
-            "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0),
-            "used_shm": bool(used_shm),
-        },
-    )
+def log_transfer_tx(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
+                    size_bytes: int, tx_time_ms: float, used_shm: bool) -> None:
+    _safe_append_jsonl(stats_file, {
+        "type": "transfer_stats",
+        "from_stage": from_stage,
+        "to_stage": to_stage,
+        "request_id": request_id,
+        "size_bytes": int(size_bytes),
+        "tx_time_ms": float(tx_time_ms),
+        "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0),
+        "used_shm": bool(used_shm),
+    })
 
 
-def log_transfer_rx(
-    stats_file: str | None,
-    from_stage: int,
-    to_stage: int,
-    request_id: Any,
-    rx_bytes: int,
-    rx_decode_time_ms: float,
-    in_flight_time_ms: float,
-) -> None:
-    _safe_append_jsonl(
-        stats_file,
-        {
-            "type": "transfer_rx_stats",
-            "from_stage": from_stage,
-            "to_stage": to_stage,
-            "request_id": request_id,
-            "rx_bytes": int(rx_bytes),
-            "rx_decode_time_ms": float(rx_decode_time_ms),
-            "in_flight_time_ms": float(in_flight_time_ms),
-            "rx_time_per_kb_ms": (
-                (float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0
-            ),
-        },
-    )
+def log_transfer_rx(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
+                    rx_bytes: int, rx_decode_time_ms: float, in_flight_time_ms: float) -> None:
+    _safe_append_jsonl(stats_file, {
+        "type": "transfer_rx_stats",
+        "from_stage": from_stage,
+        "to_stage": to_stage,
+        "request_id": request_id,
+        "rx_bytes": int(rx_bytes),
+        "rx_decode_time_ms": float(rx_decode_time_ms),
+        "in_flight_time_ms": float(in_flight_time_ms),
+        "rx_time_per_kb_ms": (float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0,
+    })
 
 
-def log_transfer_total(
-    stats_file: str | None,
-    from_stage: int,
-    to_stage: int,
-    request_id: Any,
-    size_bytes: int,
-    tx_time_ms: float,
-    in_flight_time_ms: float,
-    rx_decode_time_ms: float,
-    total_time_ms: float,
-) -> None:
-    _safe_append_jsonl(
-        stats_file,
-        {
-            "type": "transfer_total_stats",
-            "from_stage": from_stage,
-            "to_stage": to_stage,
-            "request_id": request_id,
-            "size_bytes": int(size_bytes),
-            "tx_time_ms": float(tx_time_ms),
-            "in_flight_time_ms": float(in_flight_time_ms),
-            "rx_decode_time_ms": float(rx_decode_time_ms),
-            "total_time_ms": float(total_time_ms),
-            "total_time_per_kb_ms": (
-                float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0
-            ),
-        },
-    )
+def log_transfer_total(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
+                       size_bytes: int, tx_time_ms: float, in_flight_time_ms: float,
+                       rx_decode_time_ms: float, total_time_ms: float) -> None:
+    _safe_append_jsonl(stats_file, {
+        "type": "transfer_total_stats",
+        "from_stage": from_stage,
+        "to_stage": to_stage,
+        "request_id": request_id,
+        "size_bytes": int(size_bytes),
+        "tx_time_ms": float(tx_time_ms),
+        "in_flight_time_ms": float(in_flight_time_ms),
+        "rx_decode_time_ms": float(rx_decode_time_ms),
+        "total_time_ms": float(total_time_ms),
+        "total_time_per_kb_ms": float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0,
+    })
 
 
-def log_orchestrator_e2e(
-    stats_file: str | None,
-    request_id: Any,
-    final_stage_id: int,
-    e2e_time_ms: float,
-    num_tokens_out: int,
-    e2e_time_per_token_ms: float,
-) -> None:
-    _safe_append_jsonl(
-        stats_file,
-        {
-            "type": "orchestrator_request_e2e",
-            "request_id": request_id,
-            "final_stage_id": final_stage_id,
-            "e2e_time_ms": float(e2e_time_ms),
-            "num_tokens_out": int(num_tokens_out),
-            "e2e_time_per_token_ms": float(e2e_time_per_token_ms),
-        },
-    )
+def log_orchestrator_e2e(stats_file: Optional[str], request_id: Any, final_stage_id: int,
+                         e2e_time_ms: float, num_tokens_out: int, e2e_time_per_token_ms: float) -> None:
+    _safe_append_jsonl(stats_file, {
+        "type": "orchestrator_request_e2e",
+        "request_id": request_id,
+        "final_stage_id": final_stage_id,
+        "e2e_time_ms": float(e2e_time_ms),
+        "num_tokens_out": int(num_tokens_out),
+        "e2e_time_per_token_ms": float(e2e_time_per_token_ms),
+    })
 
 
-def log_orchestrator_summary(stats_file: str | None, summary: dict[str, Any]) -> None:
+def log_orchestrator_summary(stats_file: Optional[str], summary: Dict[str, Any]) -> None:
     _safe_append_jsonl(stats_file, {"type": "orchestrator_summary", **summary})
 
 
-def log_overall_summary(overall_stats_file: str | None, summary: dict[str, Any]) -> None:
+def log_overall_summary(overall_stats_file: Optional[str], summary: Dict[str, Any]) -> None:
     _safe_append_jsonl(overall_stats_file, {"type": "overall_summary", **summary})
 
 
-def log_overall_record(overall_stats_file: str | None, record: dict[str, Any]) -> None:
+def log_overall_record(overall_stats_file: Optional[str], record: Dict[str, Any]) -> None:
     _safe_append_jsonl(overall_stats_file, record)
 
 
-def log_stage_request_stats(
-    stats_file: str | None,
-    stage_id: int,
-    request_id: Any,
-    batch_size: int,
-    num_tokens_out: int,
-    stage_gen_time_ms: float,
-    tokens_per_s: float,
-    rx_transfer_bytes: int,
-    rx_decode_time_ms: float,
-    rx_mbps: float,
-) -> None:
-    _safe_append_jsonl(
-        stats_file,
-        {
-            "type": "stage_request_stats",
-            "stage_id": stage_id,
-            "request_id": request_id,
-            "batch_size": int(batch_size),
-            "num_tokens_out": int(num_tokens_out),
-            "stage_gen_time_ms": float(stage_gen_time_ms),
-            "tokens_per_s": float(tokens_per_s),
-            "rx_transfer_bytes": int(rx_transfer_bytes),
-            "rx_decode_time_ms": float(rx_decode_time_ms),
-            "rx_mbps": float(rx_mbps),
-        },
-    )
+def log_stage_request_stats(stats_file: Optional[str], stage_id: int, request_id: Any, batch_size: int,
+                            num_tokens_out: int, stage_gen_time_ms: float, tokens_per_s: float,
+                            rx_transfer_bytes: int, rx_decode_time_ms: float, rx_mbps: float) -> None:
+    _safe_append_jsonl(stats_file, {
+        "type": "stage_request_stats",
+        "stage_id": stage_id,
+        "request_id": request_id,
+        "batch_size": int(batch_size),
+        "num_tokens_out": int(num_tokens_out),
+        "stage_gen_time_ms": float(stage_gen_time_ms),
+        "tokens_per_s": float(tokens_per_s),
+        "rx_transfer_bytes": int(rx_transfer_bytes),
+        "rx_decode_time_ms": float(rx_decode_time_ms),
+        "rx_mbps": float(rx_mbps),
+    })
 
 
-def log_stage_running_avg(
-    stats_file: str | None,
-    stage_id: int,
-    total_tokens: int,
-    total_gen_time_ms: float,
-    avg_tokens_per_s: float,
-) -> None:
-    _safe_append_jsonl(
-        stats_file,
-        {
-            "type": "stage_running_avg",
-            "stage_id": stage_id,
-            "total_tokens": int(total_tokens),
-            "total_gen_time_ms": float(total_gen_time_ms),
-            "avg_tokens_per_s": float(avg_tokens_per_s),
-        },
-    )
+def log_stage_running_avg(stats_file: Optional[str], stage_id: int, total_tokens: int,
+                          total_gen_time_ms: float, avg_tokens_per_s: float) -> None:
+    _safe_append_jsonl(stats_file, {
+        "type": "stage_running_avg",
+        "stage_id": stage_id,
+        "total_tokens": int(total_tokens),
+        "total_gen_time_ms": float(total_gen_time_ms),
+        "avg_tokens_per_s": float(avg_tokens_per_s),
+    })
 
 
-def log_stage_batch_stats(
-    stats_file: str | None,
-    stage_id: int,
-    batch_size: int,
-    batch_gen_time_ms: float,
-    request_ids: list[Any],
-) -> None:
-    _safe_append_jsonl(
-        stats_file,
-        {
-            "type": "stage_batch_stats",
-            "stage_id": stage_id,
-            "batch_size": int(batch_size),
-            "batch_gen_time_ms": float(batch_gen_time_ms),
-            "request_ids": list(request_ids),
-        },
-    )
+def log_stage_batch_stats(stats_file: Optional[str], stage_id: int, batch_size: int,
+                          batch_gen_time_ms: float, request_ids: List[Any]) -> None:
+    _safe_append_jsonl(stats_file, {
+        "type": "stage_batch_stats",
+        "stage_id": stage_id,
+        "batch_size": int(batch_size),
+        "batch_gen_time_ms": float(batch_gen_time_ms),
+        "request_ids": list(request_ids),
+    })
 
 
-def compute_and_log_stage_request_stats(
-    stats_file: str | None,
-    stage_id: int,
-    request_id: Any,
-    batch_size: int,
-    engine_outputs: list[Any],
-    stage_gen_time_ms: float,
-    rx_transfer_bytes: int,
-    rx_decode_time_ms: float,
-) -> None:
+def compute_and_log_stage_request_stats(stats_file: Optional[str], stage_id: int, request_id: Any,
+                                       batch_size: int, engine_outputs: List[Any],
+                                       stage_gen_time_ms: float, rx_transfer_bytes: int,
+                                       rx_decode_time_ms: float) -> None:
     """Compute per-request metrics and log them in one call."""
     num_tokens = count_tokens_from_outputs(engine_outputs)
     tokens_per_s = (num_tokens * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0
-    rx_mbps = (
-        (float(rx_transfer_bytes) * 8.0) / (max(float(rx_decode_time_ms), 1e-6) * 1000.0)
-        if rx_transfer_bytes > 0
-        else 0.0
-    )
+    rx_mbps = (float(rx_transfer_bytes) * 8.0) / (max(float(rx_decode_time_ms), 1e-6) * 1000.0) if rx_transfer_bytes > 0 else 0.0
     log_stage_request_stats(
         stats_file,
         stage_id,
@@ -299,16 +218,13 @@ def compute_and_log_stage_request_stats(
 
 # ----------------- Aggregation helpers for orchestrator -----------------
 
-
-def record_stage_metrics(
-    per_request: dict[int, dict[str, Any]],
-    stage_req_counts: list[int],
-    stage_total_time_ms: list[float],
-    stage_total_tokens: list[int],
-    stage_id: int,
-    req_id: Any,
-    metrics: dict[str, Any],
-) -> None:
+def record_stage_metrics(per_request: Dict[int, Dict[str, Any]],
+                         stage_req_counts: List[int],
+                         stage_total_time_ms: List[float],
+                         stage_total_tokens: List[int],
+                         stage_id: int,
+                         req_id: Any,
+                         metrics: Dict[str, Any]) -> None:
     try:
         stage_req_counts[stage_id] += 1
         stage_total_tokens[stage_id] += int(metrics.get("num_tokens_out", 0))
@@ -323,32 +239,23 @@ def record_stage_metrics(
         pass
 
 
-def aggregate_rx_and_maybe_total(
-    transfer_edge_req: dict[tuple[int, int, int], dict[str, float]],
-    transfer_agg: dict[tuple[int, int], dict[str, float]],
-    per_request: dict[int, dict[str, Any]],
-    stage_id: int,
-    req_id: Any,
-    rx_bytes: float,
-    rx_ms: float,
-    in_flight_ms: float,
-) -> tuple[int, float, float] | None:
+def aggregate_rx_and_maybe_total(transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]],
+                                 transfer_agg: Dict[tuple[int, int], Dict[str, float]],
+                                 per_request: Dict[int, Dict[str, Any]],
+                                 stage_id: int,
+                                 req_id: Any,
+                                 rx_bytes: float,
+                                 rx_ms: float,
+                                 in_flight_ms: float) -> Optional[tuple[int, float, float]]:
     try:
         # Update RX aggregates for (stage_id-1 -> stage_id)
         if stage_id > 0:
             key = (stage_id - 1, stage_id)
             agg = transfer_agg.get(key)
             if agg is None:
-                agg = {
-                    "sum_bytes": 0.0,
-                    "sum_ms": 0.0,
-                    "count": 0.0,
-                    "sum_rx_bytes": 0.0,
-                    "sum_rx_ms": 0.0,
-                    "rx_count": 0.0,
-                    "sum_total_ms": 0.0,
-                    "total_count": 0.0,
-                }
+                agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0,
+                       "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0,
+                       "sum_total_ms": 0.0, "total_count": 0.0}
                 transfer_agg[key] = agg
             agg["sum_rx_bytes"] += float(rx_bytes)
             agg["sum_rx_ms"] += float(rx_ms)
@@ -380,29 +287,20 @@ def aggregate_rx_and_maybe_total(
         return None
 
 
-def record_sender_transfer_agg(
-    transfer_agg: dict[tuple[int, int], dict[str, float]],
-    transfer_edge_req: dict[tuple[int, int, int], dict[str, float]],
-    from_stage: int,
-    to_stage: int,
-    req_id: Any,
-    size_bytes: int,
-    tx_ms: float,
-) -> None:
+def record_sender_transfer_agg(transfer_agg: Dict[tuple[int, int], Dict[str, float]],
+                               transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]],
+                               from_stage: int,
+                               to_stage: int,
+                               req_id: Any,
+                               size_bytes: int,
+                               tx_ms: float) -> None:
     try:
         key = (from_stage, to_stage)
         agg = transfer_agg.get(key)
         if agg is None:
-            agg = {
-                "sum_bytes": 0.0,
-                "sum_ms": 0.0,
-                "count": 0.0,
-                "sum_rx_bytes": 0.0,
-                "sum_rx_ms": 0.0,
-                "rx_count": 0.0,
-                "sum_total_ms": 0.0,
-                "total_count": 0.0,
-            }
+            agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0,
+                   "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0,
+                   "sum_total_ms": 0.0, "total_count": 0.0}
             transfer_agg[key] = agg
         agg["sum_bytes"] += float(size_bytes)
         agg["sum_ms"] += float(tx_ms)
@@ -420,7 +318,7 @@ def record_sender_transfer_agg(
         pass
 
 
-def count_tokens_from_outputs(engine_outputs: list[Any]) -> int:
+def count_tokens_from_outputs(engine_outputs: List[Any]) -> int:
     total = 0
     for _ro in engine_outputs:
         try:
@@ -434,35 +332,29 @@ def count_tokens_from_outputs(engine_outputs: list[Any]) -> int:
     return total
 
 
-def build_stage_summary(
-    stage_req_counts: list[int],
-    stage_total_tokens: list[int],
-    stage_total_time_ms: list[float],
-) -> list[dict[str, Any]]:
-    summary: list[dict[str, Any]] = []
+def build_stage_summary(stage_req_counts: List[int],
+                        stage_total_tokens: List[int],
+                        stage_total_time_ms: List[float]) -> List[Dict[str, Any]]:
+    summary: List[Dict[str, Any]] = []
     for sid in range(len(stage_req_counts)):
         reqs = stage_req_counts[sid]
         tokens = stage_total_tokens[sid]
         total_ms = float(stage_total_time_ms[sid])
         avg_req = (total_ms / reqs) if reqs > 0 else 0.0
         avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0
-        summary.append(
-            {
-                "stage_id": sid,
-                "requests": int(reqs),
-                "tokens": int(tokens),
-                "total_time_ms": total_ms,
-                "avg_time_per_request_ms": avg_req,
-                "avg_tokens_per_s": avg_tok,
-            }
-        )
+        summary.append({
+            "stage_id": sid,
+            "requests": int(reqs),
+            "tokens": int(tokens),
+            "total_time_ms": total_ms,
+            "avg_time_per_request_ms": avg_req,
+            "avg_tokens_per_s": avg_tok,
+        })
     return summary
 
 
-def build_transfer_summary(
-    transfer_agg: dict[tuple[int, int], dict[str, float]],
-) -> list[dict[str, Any]]:
-    summary: list[dict[str, Any]] = []
+def build_transfer_summary(transfer_agg: Dict[tuple[int, int], Dict[str, float]]) -> List[Dict[str, Any]]:
+    summary: List[Dict[str, Any]] = []
     for (src, dst), agg in transfer_agg.items():
         sum_bytes = float(agg.get("sum_bytes", 0.0))
         sum_ms = float(agg.get("sum_ms", 0.0))
@@ -475,57 +367,50 @@ def build_transfer_summary(
         sum_total_ms = float(agg.get("sum_total_ms", 0.0))
         samples_total = int(agg.get("total_count", 0.0))
         total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
-        summary.append(
-            {
-                "from_stage": src,
-                "to_stage": dst,
-                "samples": samples,
-                "total_bytes": int(sum_bytes),
-                "total_time_ms": sum_ms,
-                "tx_mbps": tx_mbps,
-                "rx_samples": samples_rx,
-                "rx_total_bytes": int(sum_rx_bytes),
-                "rx_total_time_ms": sum_rx_ms,
-                "rx_mbps": rx_mbps,
-                "total_samples": samples_total,
-                "total_transfer_time_ms": sum_total_ms,
-                "total_mbps": total_mbps,
-            }
-        )
+        summary.append({
+            "from_stage": src,
+            "to_stage": dst,
+            "samples": samples,
+            "total_bytes": int(sum_bytes),
+            "total_time_ms": sum_ms,
+            "tx_mbps": tx_mbps,
+            "rx_samples": samples_rx,
+            "rx_total_bytes": int(sum_rx_bytes),
+            "rx_total_time_ms": sum_rx_ms,
+            "rx_mbps": rx_mbps,
+            "total_samples": samples_total,
+            "total_transfer_time_ms": sum_total_ms,
+            "total_mbps": total_mbps,
+        })
     return summary
 
 
 class OrchestratorMetrics:
-    def __init__(
-        self,
-        num_stages: int,
-        enable_stats: bool,
-        stats_file: str | None,
-        overall_stats_file: str | None,
-        wall_start_ts: float,
-    ) -> None:
+    def __init__(self, num_stages: int, enable_stats: bool,
+                 stats_file: Optional[str], overall_stats_file: Optional[str],
+                 wall_start_ts: float) -> None:
         self.num_stages = int(num_stages)
         self.enable_stats = bool(enable_stats)
         self.stats_file = stats_file
         self.overall_stats_file = overall_stats_file
-        self.stage_total_time_ms: list[float] = [0.0 for _ in range(self.num_stages)]
-        self.stage_total_tokens: list[int] = [0 for _ in range(self.num_stages)]
-        self.stage_req_counts: list[int] = [0 for _ in range(self.num_stages)]
-        self.transfer_agg: dict[tuple[int, int], dict[str, float]] = {}
-        self.transfer_edge_req: dict[tuple[int, int, int], dict[str, float]] = {}
+        self.stage_total_time_ms: List[float] = [0.0 for _ in range(self.num_stages)]
+        self.stage_total_tokens: List[int] = [0 for _ in range(self.num_stages)]
+        self.stage_req_counts: List[int] = [0 for _ in range(self.num_stages)]
+        self.transfer_agg: Dict[tuple[int, int], Dict[str, float]] = {}
+        self.transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]] = {}
         self.e2e_total_ms: float = 0.0
         self.e2e_total_tokens: int = 0
         self.e2e_count: int = 0
         self.e2e_done: set[int] = set()
-        self.per_request: dict[int, dict[str, Any]] = {}
+        self.per_request: Dict[int, Dict[str, Any]] = {}
         self.sum_per_request_transfer_ms: float = 0.0
         self.wall_start_ts: float = float(wall_start_ts)
         self.last_finish_ts: float = float(wall_start_ts)
-        self.stage_seen_batches: dict[int, set] = {sid: set() for sid in range(self.num_stages)}
-        self.stage_first_ts: list[float | None] = [None for _ in range(self.num_stages)]
-        self.stage_last_ts: list[float | None] = [None for _ in range(self.num_stages)]
+        self.stage_seen_batches: Dict[int, set] = {sid: set() for sid in range(self.num_stages)}
+        self.stage_first_ts: List[Optional[float]] = [None for _ in range(self.num_stages)]
+        self.stage_last_ts: List[Optional[float]] = [None for _ in range(self.num_stages)]
 
-    def on_stage_metrics(self, stage_id: int, req_id: Any, metrics: dict[str, Any]) -> None:
+    def on_stage_metrics(self, stage_id: int, req_id: Any, metrics: Dict[str, Any]) -> None:
         record_stage_metrics(
             self.per_request,
             self.stage_req_counts,
@@ -581,15 +466,8 @@ class OrchestratorMetrics:
                     float(total_ms_c),
                 )
 
-    def on_forward(
-        self,
-        from_stage: int,
-        to_stage: int,
-        req_id: Any,
-        size_bytes: int,
-        tx_ms: float,
-        used_shm: bool,
-    ) -> None:
+    def on_forward(self, from_stage: int, to_stage: int, req_id: Any,
+                   size_bytes: int, tx_ms: float, used_shm: bool) -> None:
         # Mark first input time for the destination stage if not set
         if self.stage_first_ts[to_stage] is None:
             self.stage_first_ts[to_stage] = time.time()
@@ -613,7 +491,8 @@ class OrchestratorMetrics:
             float(tx_ms),
         )
 
-    def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: list[Any], req_start_ts: float) -> None:
+    def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: List[Any],
+                            req_start_ts: float) -> None:
         try:
             rid_int = int(req_id)
         except Exception:
@@ -647,34 +526,30 @@ class OrchestratorMetrics:
             e2e_tpt = (e2e_ms / num_tokens) if num_tokens > 0 else 0.0
             log_orchestrator_e2e(self.stats_file, req_id, stage_id, e2e_ms, int(num_tokens), e2e_tpt)
 
-    def build_and_log_summary(self, final_stage_id_for_e2e: int) -> dict[str, Any]:
+    def build_and_log_summary(self, final_stage_id_for_e2e: int) -> Dict[str, Any]:
         # Compute stage summary using wall time between first input and last output per stage
-        stage_summary: list[dict[str, Any]] = []
+        stage_summary: List[Dict[str, Any]] = []
         for sid in range(self.num_stages):
             first_ts = self.stage_first_ts[sid]
             last_ts = self.stage_last_ts[sid]
-            total_ms = (
-                (max(0.0, (last_ts - first_ts)) * 1000.0) if (first_ts is not None and last_ts is not None) else 0.0
-            )
+            total_ms = (max(0.0, (last_ts - first_ts)) * 1000.0) if (first_ts is not None and last_ts is not None) else 0.0
             reqs = self.stage_req_counts[sid]
             tokens = self.stage_total_tokens[sid]
             avg_req = (total_ms / reqs) if reqs > 0 else 0.0
             avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0
-            stage_summary.append(
-                {
-                    "stage_id": sid,
-                    "requests": int(reqs),
-                    "tokens": int(tokens),
-                    "total_time_ms": float(total_ms),
-                    "avg_time_per_request_ms": float(avg_req),
-                    "avg_tokens_per_s": float(avg_tok),
-                }
-            )
+            stage_summary.append({
+                "stage_id": sid,
+                "requests": int(reqs),
+                "tokens": int(tokens),
+                "total_time_ms": float(total_ms),
+                "avg_time_per_request_ms": float(avg_req),
+                "avg_tokens_per_s": float(avg_tok),
+            })
         transfer_summary = build_transfer_summary(self.transfer_agg)
         e2e_avg_req = (self.e2e_total_ms / self.e2e_count) if self.e2e_count > 0 else 0.0
         e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / self.e2e_total_ms) if self.e2e_total_ms > 0 else 0.0
         wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0)
-        summary: dict[str, Any] = {
+        summary: Dict[str, Any] = {
             "e2e_requests": int(self.e2e_count),
             "e2e_total_time_ms": float(wall_time_ms),
             "e2e_sum_time_ms": float(self.e2e_total_ms),
@@ -691,3 +566,5 @@ class OrchestratorMetrics:
         if self.enable_stats and self.overall_stats_file:
             log_overall_summary(self.overall_stats_file, summary)
         return summary
+
+

vllm_omni/entrypoints/omni_llm.py

@@ -1,10 +1,11 @@
+from typing import Any, Optional, Sequence, Union
+import logging
 import multiprocessing as mp
+from concurrent.futures import ThreadPoolExecutor, as_completed
 import os
+import queue
 import sys
 import time
-from collections.abc import Sequence
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import Any, Optional, Union
 
 import cloudpickle
 from pydantic import ValidationError
@@ -25,17 +26,19 @@ from vllm.v1.engine.llm_engine import LLMEngine
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.engine.processor import OmniProcessor
+from vllm_omni.entrypoints.omni_stage import OmniStage
+from vllm_omni.entrypoints.utils import load_stage_configs_from_model
 from vllm_omni.entrypoints.log_utils import (
-    OrchestratorMetrics,
+    remove_old_logs,
     configure_orchestrator_logger,
     init_stats_paths,
-    remove_old_logs,
+    OrchestratorMetrics,
+)
+from vllm_omni.entrypoints.stage_utils import (
+    maybe_load_from_ipc as _load,
+    encode_for_ipc as _encode,
+    serialize_obj as _ser,
 )
-from vllm_omni.entrypoints.omni_stage import OmniStage
-from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
-from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
-from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
-from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -49,25 +52,23 @@ class OmniLLM:
     - max_inflight=1 per stage (serial within a stage), but pipeline across stages
     """
 
-    def __init__(
-        self,
-        model: str,
-        stage_configs_path: Optional[str] = None,
-        log_stats: bool = False,
-        log_file: Optional[str] = None,
-        init_sleep_seconds: int = 20,
-        shm_threshold_bytes: int = 65536,
-        batch_timeout: int = 10,
-        init_timeout: int = 300,
-        **kwargs,
-    ):
+    def __init__(self, model: str,
+                 stage_configs=None,
+                 log_stats: bool = False,
+                 log_file: Optional[str] = None,
+                 init_sleep_seconds: int = 20,
+                 shm_threshold_bytes: int = 65536,
+                 batch_timeout: int = 10,
+                 init_timeout: int = 300,
+                 **kwargs):
         self.batch_timeout = batch_timeout
         self._enable_stats: bool = bool(log_stats)
         # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent.
-        if stage_configs_path is None:
+        if stage_configs is None:
             self.stage_configs = load_stage_configs_from_model(model)
         else:
-            self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
+            self.stage_configs = stage_configs
+
 
         # Optional file handler for orchestrator
         self._log_file = log_file
@@ -78,15 +79,8 @@ class OmniLLM:
         self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
         self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout)
 
-    def _initialize_stages(
-        self,
-        model: str,
-        init_sleep_seconds: int,
-        shm_threshold_bytes: int,
-        init_timeout: int,
-    ) -> None:
+    def _initialize_stages(self, model: str, init_sleep_seconds: int, shm_threshold_bytes: int, init_timeout: int) -> None:
         self.stage_list: list[OmniStage] = []
-
         # Build OmniStage instances in parallel, preserve original order
         def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]:
             idx, cfg = idx_cfg
@@ -110,6 +104,7 @@ class OmniLLM:
         # Wait for all stages to report readiness before seeding
         self._stages_ready: set[int] = set()
         self._wait_for_stages_ready(timeout=init_timeout)
+        
 
     def _start_stage_processes(self, model: str) -> None:
         for stage_id, stage in enumerate(self.stage_list):
@@ -136,10 +131,7 @@ class OmniLLM:
             try:
                 q.put_nowait(None)
             except Exception as e:
-                logger.warning(
-                    "[Orchestrator] Failed to send shutdown signal to stage input queue: %s",
-                    e,
-                )
+                logger.warning("[Orchestrator] Failed to send shutdown signal to stage input queue: %s", e)
         for stage in self.stage_list:
             try:
                 stage.stop_stage_worker()
@@ -155,26 +147,17 @@ class OmniLLM:
     def generate(
         self,
         prompts: Union[PromptType, Sequence[PromptType]],
-        sampling_params_list: Optional[Union[SamplingParams, Sequence[SamplingParams]]] = None,
-    ) -> list[OmniRequestOutput]:
-        try:
-            return self._run_generation(prompts, sampling_params_list)
-        except Exception as e:
-            logger.exception("[Orchestrator] Failed to run generation: %s", e)
-            raise e
-        finally:
-            self.close()
-
-    def _run_generation(
-        self,
-        prompts: Union[PromptType, Sequence[PromptType]],
-        sampling_params_list: Optional[Union[SamplingParams, Sequence[SamplingParams]]] = None,
+        sampling_params_list: Optional[
+            Union[SamplingParams, Sequence[SamplingParams]]
+        ] = None,
     ) -> list[OmniRequestOutput]:
         logger.debug("[Orchestrator] generate() called")
         if sampling_params_list is None:
             raise ValueError("sampling_params_list is required for pipelined generation")
         if len(sampling_params_list) != len(self.stage_list):
-            raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
+            raise ValueError(
+                f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}"
+            )
 
         # Normalize prompts to a list for per-request iteration
         if not isinstance(prompts, (list, tuple)):
@@ -193,6 +176,7 @@ class OmniLLM:
         # Track per-request start time for end-to-end timing
         _req_start_ts: dict[int, float] = {}
         _wall_start_ts: float = time.time()
+        _last_finish_ts: float = _wall_start_ts
 
         # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
         final_stage_id_for_e2e = -1
@@ -203,20 +187,10 @@ class OmniLLM:
             if final_stage_id_for_e2e < 0:
                 final_stage_id_for_e2e = len(self.stage_list) - 1
         except Exception as e:
-            logger.debug(
-                "[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s",
-                e,
-                exc_info=True,
-            )
+            logger.debug("[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s", e, exc_info=True)
             final_stage_id_for_e2e = len(self.stage_list) - 1
         # Metrics/aggregation helper
-        metrics = OrchestratorMetrics(
-            num_stages,
-            self._enable_stats,
-            self._stats_file,
-            self._overall_stats_file,
-            _wall_start_ts,
-        )
+        metrics = OrchestratorMetrics(num_stages, self._enable_stats, self._stats_file, self._overall_stats_file, _wall_start_ts)
 
         # Seed stage-0 queue with all requests
         logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts))
@@ -240,11 +214,7 @@ class OmniLLM:
         completed_requests = 0
         total_requests = len(request_prompts)
 
-        logger.debug(
-            "[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d",
-            total_requests,
-            num_stages,
-        )
+        logger.debug("[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d", total_requests, num_stages)
         while completed_requests < total_requests:
             made_progress = False
             for stage_id, stage in enumerate(self.stage_list):
@@ -255,17 +225,11 @@ class OmniLLM:
                 made_progress = True
                 req_id = result.get("request_id")
                 if "error" in result:
-                    logger.error(
-                        "Stage %s error on request %s: %s",
-                        stage_id,
-                        req_id,
-                        result["error"],
-                    )
+                    logger.error("Stage %s error on request %s: %s", stage_id, req_id, result["error"])
                     continue
-
+                
                 if result.get("type") == "stage_ready":
-                    # Only happens when stage is initialized slower than expected,
-                    # so we wait for a short time and try again
+                    #Only happens when stage is initialized slower than expected, so we wait for a short time and try again
                     time.sleep(0.05)
                     continue
 
@@ -277,17 +241,8 @@ class OmniLLM:
                     if _m is not None:
                         metrics.on_stage_metrics(stage_id, req_id, _m)
                 except Exception as e:
-                    logger.exception(
-                        "[Orchestrator] Failed to process metrics for stage %s, req %s: %s",
-                        stage_id,
-                        req_id,
-                        e,
-                    )
-                logger.debug(
-                    "[Orchestrator] Stage-%s completed request %s; forwarding or finalizing",
-                    stage_id,
-                    req_id,
-                )
+                    logger.exception("[Orchestrator] Failed to process metrics for stage %s, req %s: %s", stage_id, req_id, e)
+                logger.debug("[Orchestrator] Stage-%s completed request %s; forwarding or finalizing", stage_id, req_id)
                 stage.set_engine_outputs(engine_outputs)
 
                 if getattr(stage, "final_output", False):
@@ -298,30 +253,15 @@ class OmniLLM:
                             request_output=engine_outputs,
                         )
                     )
-                    logger.debug(
-                        "[Orchestrator] Request %s finalized at stage-%s",
-                        req_id,
-                        stage_id,
-                    )
+                    logger.debug("[Orchestrator] Request %s finalized at stage-%s", req_id, stage_id)
 
-                    # End-to-end timing and time-per-token for final output
-                    # (only once per request at the designated final stage)
+                    # End-to-end timing and time-per-token for final output (only once per request at the designated final stage)
                     try:
                         rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
                         if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
-                            metrics.on_finalize_request(
-                                stage_id,
-                                req_id,
-                                engine_outputs,
-                                _req_start_ts.get(req_id, _wall_start_ts),
-                            )
+                            metrics.on_finalize_request(stage_id, req_id, engine_outputs, _req_start_ts.get(req_id, _wall_start_ts))
                     except Exception as e:
-                        logger.exception(
-                            "[Orchestrator] Finalize request handling error for req %s at stage %s: %s",
-                            req_id,
-                            stage_id,
-                            e,
-                        )
+                        logger.exception("[Orchestrator] Finalize request handling error for req %s at stage %s: %s", req_id, stage_id, e)
 
                 next_stage_id = stage_id + 1
                 if next_stage_id < num_stages:
@@ -332,7 +272,7 @@ class OmniLLM:
                         # Measure transfer size and time (encode + enqueue)
                         size_bytes = 0
                         try:
-                            size_bytes = len(_set(next_inputs))
+                            size_bytes = len(_ser(next_inputs))
                         except Exception:
                             size_bytes = 0
                         t0 = time.time()
@@ -342,51 +282,27 @@ class OmniLLM:
                             obj_key="engine_inputs",
                             shm_key="engine_inputs_shm",
                         )
-                        ipc_payload.update(
-                            {
-                                "request_id": req_id,
-                                "sampling_params": sp_next,
-                                "sent_ts": time.time(),
-                            }
-                        )
+                        ipc_payload.update({
+                            "request_id": req_id,
+                            "sampling_params": sp_next,
+                            "sent_ts": time.time(),
+                        })
                         self.stage_list[next_stage_id].submit(ipc_payload)
                         t1 = time.time()
                         tx_ms = (t1 - t0) * 1000.0
-                        metrics.on_forward(
-                            stage_id,
-                            next_stage_id,
-                            req_id,
-                            int(size_bytes),
-                            float(tx_ms),
-                            bool("engine_inputs_shm" in ipc_payload),
-                        )
+                        metrics.on_forward(stage_id, next_stage_id, req_id, int(size_bytes), float(tx_ms), bool("engine_inputs_shm" in ipc_payload))
                     except Exception as e:
-                        logger.warning(
-                            "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
-                            req_id,
-                            e,
-                        )
-                        self.stage_list[next_stage_id].submit(
-                            {
-                                "request_id": req_id,
-                                "engine_inputs": next_inputs,
-                                "sampling_params": sp_next,
-                            }
-                        )
-                    logger.debug(
-                        "[Orchestrator] Forwarded request %s to stage-%s",
-                        req_id,
-                        next_stage_id,
-                    )
+                        logger.warning("[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload", req_id, e)
+                        self.stage_list[next_stage_id].submit({
+                            "request_id": req_id,
+                            "engine_inputs": next_inputs,
+                            "sampling_params": sp_next,
+                        })
+                    logger.debug("[Orchestrator] Forwarded request %s to stage-%s", req_id, next_stage_id)
                     remaining_by_stage[next_stage_id] += 1
                 else:
                     completed_requests += 1
-                    logger.debug(
-                        "[Orchestrator] Request %s fully completed (%d/%d)",
-                        req_id,
-                        completed_requests,
-                        total_requests,
-                    )
+                    logger.debug("[Orchestrator] Request %s fully completed (%d/%d)", req_id, completed_requests, total_requests)
 
             if not made_progress:
                 time.sleep(0.005)
@@ -425,9 +341,7 @@ class OmniLLM:
             not_ready = sorted(set(range(num_stages)) - set(self._stages_ready))
             logger.warning(
                 "[Orchestrator] Initialization timeout: only %s/%s stages are ready; not ready: %s",
-                len(self._stages_ready),
-                num_stages,
-                not_ready,
+                len(self._stages_ready), num_stages, not_ready,
             )
             # Provide actionable suggestions before shutdown
             try:
@@ -438,7 +352,9 @@ class OmniLLM:
                     "Increase initialization wait time (init_sleep_seconds or call-site timeout).",
                 ]
                 if getattr(self, "_log_file", None):
-                    suggestions.append(f"Inspect per-stage log files for details: {self._log_file}.stage<id>.log")
+                    suggestions.append(
+                        f"Inspect per-stage log files for details: {self._log_file}.stage<id>.log"
+                    )
                 logger.error(
                     "[Orchestrator] Stage initialization failed, shutting down. Suggestions:\n- %s",
                     "\n- ".join(suggestions),
@@ -468,9 +384,13 @@ class OmniStageLLM(LLM):
     def __init__(
         self,
         model: str,
-        compilation_config: Optional[Union[int, dict[str, Any], CompilationConfig]] = None,
+        compilation_config: Optional[
+            Union[int, dict[str, Any], CompilationConfig]
+        ] = None,
         hf_overrides: Optional[HfOverrides] = None,
-        structured_outputs_config: Optional[Union[dict[str, Any], StructuredOutputsConfig]] = None,
+        structured_outputs_config: Optional[
+            Union[dict[str, Any], StructuredOutputsConfig]
+        ] = None,
         **kwargs,
     ):
         """LLM constructor."""
@@ -484,7 +404,9 @@ class OmniStageLLM(LLM):
             if isinstance(worker_cls, type):
                 kwargs["worker_cls"] = cloudpickle.dumps(worker_cls)
 
-        if "kv_transfer_config" in kwargs and isinstance(kwargs["kv_transfer_config"], dict):
+        if "kv_transfer_config" in kwargs and isinstance(
+            kwargs["kv_transfer_config"], dict
+        ):
             from vllm.config.kv_transfer import KVTransferConfig
 
             raw_config_dict = kwargs["kv_transfer_config"]
@@ -492,7 +414,8 @@ class OmniStageLLM(LLM):
                 kwargs["kv_transfer_config"] = KVTransferConfig(**raw_config_dict)
             except ValidationError as e:
                 logger.error(
-                    "Failed to convert 'kv_transfer_config' dict to KVTransferConfig object. Dict: %s. Error: %s",
+                    "Failed to convert 'kv_transfer_config' dict to "
+                    "KVTransferConfig object. Dict: %s. Error: %s",
                     raw_config_dict,
                     e,
                 )
@@ -505,10 +428,16 @@ class OmniStageLLM(LLM):
 
         if compilation_config is not None:
             if isinstance(compilation_config, int):
-                compilation_config_instance = CompilationConfig(level=compilation_config)
+                compilation_config_instance = CompilationConfig(
+                    level=compilation_config
+                )
             elif isinstance(compilation_config, dict):
                 compilation_config_instance = CompilationConfig(
-                    **{k: v for k, v in compilation_config.items() if is_init_field(CompilationConfig, k)}
+                    **{
+                        k: v
+                        for k, v in compilation_config.items()
+                        if is_init_field(CompilationConfig, k)
+                    }
                 )
             else:
                 compilation_config_instance = compilation_config
@@ -518,7 +447,11 @@ class OmniStageLLM(LLM):
         if structured_outputs_config is not None:
             if isinstance(structured_outputs_config, dict):
                 structured_outputs_instance = StructuredOutputsConfig(
-                    **{k: v for k, v in structured_outputs_config.items() if is_init_field(StructuredOutputsConfig, k)}
+                    **{
+                        k: v
+                        for k, v in structured_outputs_config.items()
+                        if is_init_field(StructuredOutputsConfig, k)
+                    }
                 )
             else:
                 structured_outputs_instance = structured_outputs_config
@@ -534,7 +467,9 @@ class OmniStageLLM(LLM):
         )
 
         # Create the Engine (autoselects V0 vs V1)
-        self.llm_engine = LLMEngine.from_engine_args(engine_args=engine_args, usage_context=UsageContext.LLM_CLASS)
+        self.llm_engine = LLMEngine.from_engine_args(
+            engine_args=engine_args, usage_context=UsageContext.LLM_CLASS
+        )
         self.llm_engine.output_processor = MultimodalOutputProcessor(
             tokenizer=self.llm_engine.tokenizer,
             log_stats=self.llm_engine.log_stats,
@@ -556,4 +491,6 @@ class OmniStageLLM(LLM):
 
         # Load the Input/Output processor plugin if any
         io_processor_plugin = self.llm_engine.model_config.io_processor_plugin
-        self.io_processor = get_io_processor(self.llm_engine.vllm_config, io_processor_plugin)
+        self.io_processor = get_io_processor(
+            self.llm_engine.vllm_config, io_processor_plugin
+        )

vllm_omni/entrypoints/omni_stage.py

@@ -1,72 +1,64 @@
 """
 Stage manager for orchestrating multiple engines in vLLM-omni.
 
-Enhanced to encapsulate per-stage process lifecycle and worker logic
-(device setup, LLM init, batching, shared-memory IPC), while preserving
-the original input processing utilities for cross-stage data wiring.
-
 Enhanced to encapsulate per-stage process lifecycle and worker logic
 (device setup, LLM init, batching, shared-memory IPC), while preserving
 the original input processing utilities for cross-stage data wiring.
 """
 
-import asyncio
 import importlib
 import logging
 import multiprocessing as mp
-from typing import Any, Optional, Union
+from typing import Any, Dict, List, Optional, Union
+
+from vllm_omni.entrypoints.stage_utils import (
+    _to_dict,
+    set_stage_gpu_devices,
+    maybe_load_from_ipc_with_metrics,
+    maybe_dump_to_shm,
+)
 
 from vllm.inputs import TextPrompt
-from vllm.inputs.preprocess import InputPreprocessor
-from vllm.logger import init_logger
-from vllm.sampling_params import SamplingParams
-from vllm.transformers_utils.tokenizer import AnyTokenizer
-from vllm.usage.usage_lib import UsageContext
 from vllm.v1.engine import EngineCoreOutput
 from vllm.v1.engine.async_llm import AsyncLLM
 from vllm.v1.engine.llm_engine import LLMEngine
-
-from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
-from vllm_omni.entrypoints.stage_utils import (
-    _to_dict,
-    maybe_dump_to_shm,
-    maybe_load_from_ipc_with_metrics,
-    set_stage_gpu_devices,
-)
 from vllm_omni.inputs.data import OmniTokensPrompt
 
-logger = init_logger(__name__)
-
 
 class OmniStage:
     def __init__(self, stage_config):
         self.stage_config = stage_config
         self.engine = None
         self.async_engine = None
-        self.vllm_config = None
-        self.tokenizer = None
-        self.input_preprocessor = None
-        self.is_tracing_enabled = False
         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", [])
+        if hasattr(stage_config, "engine_input_source"):
+            self.engine_input_source = stage_config.engine_input_source
+        else:
+            self.engine_input_source = []
         self.engine_output_type = stage_config.engine_args.engine_output_type
         self.engine_outputs = None
-        self.is_comprehension = getattr(stage_config, "is_comprehension", False)
         if hasattr(stage_config, "custom_process_input_func"):
             # Import the module specified in the config (already a full module path)
-            module_path, func_name = stage_config.custom_process_input_func.rsplit(".", 1)
+            module_path, func_name = stage_config.custom_process_input_func.rsplit(
+                ".", 1
+            )
             module = importlib.import_module(module_path)
             self.custom_process_input_func = getattr(module, func_name)
         else:
             self.custom_process_input_func = None
 
-        self.final_output = getattr(stage_config, "final_output", False)
-        self.final_output_type = getattr(stage_config, "final_output_type", None)
-        default_sampling_params = getattr(stage_config, "default_sampling_params", {})
-        self.default_sampling_params = SamplingParams(**_to_dict(default_sampling_params))
+        if hasattr(stage_config, "final_output"):
+            self.final_output = stage_config.final_output
+        else:
+            self.final_output = False
+
+        if hasattr(stage_config, "final_output_type"):
+            self.final_output_type = stage_config.final_output_type
+        else:
+            self.final_output_type = None
+
         # Runtime orchestration state (added)
         self._in_q: Optional[mp.Queue] = None
         self._out_q: Optional[mp.Queue] = None
@@ -83,22 +75,6 @@ class OmniStage:
         """Initialize the async engine for the stage."""
         self.async_engine = async_engine
 
-    def set_vllm_config(self, vllm_config) -> None:
-        """Set the vllm_config for the stage (received from worker process)."""
-        self.vllm_config = vllm_config
-
-    def set_tokenizer(self, tokenizer: AnyTokenizer) -> None:
-        """Set the tokenizer for the stage (received from worker process)."""
-        self.tokenizer = tokenizer
-
-    def set_input_preprocessor(self, input_preprocessor: InputPreprocessor) -> None:
-        """Set the input preprocessor for the stage (received from worker process)."""
-        self.input_preprocessor = input_preprocessor
-
-    def set_is_tracing_enabled(self, is_tracing_enabled: bool) -> None:
-        """Set the is_tracing_enabled for the stage (received from worker process)."""
-        self.is_tracing_enabled = is_tracing_enabled
-
     def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None:
         """Set the engine output for the stage."""
         self.engine_outputs = engine_outputs
@@ -112,7 +88,6 @@ class OmniStage:
         self,
         model: str,
         *,
-        is_async: bool = False,
         log_file: Optional[str] = None,
         shm_threshold_bytes: int = 65536,
         ctx: Optional[mp.context.BaseContext] = None,
@@ -125,37 +100,16 @@ class OmniStage:
         # Prepare lightweight dict config for worker
         engine_args = _to_dict(self.engine_args)
         runtime_cfg = _to_dict(getattr(self.stage_config, "runtime", {}))
-        stage_payload: dict[str, Any] = {
+        stage_payload: Dict[str, Any] = {
             "stage_id": self.stage_id,
             "engine_args": engine_args,
             "runtime": runtime_cfg,
             "shm_threshold_bytes": self._shm_threshold_bytes,
         }
-        if is_async:
-            self._proc = ctx.Process(
-                target=_stage_worker_async_entry,
-                args=(
-                    self,
-                    model,
-                    stage_payload,
-                    self._in_q,
-                    self._out_q,
-                    self._log_file,
-                    batch_timeout,
-                ),
-            )
-        else:
-            self._proc = ctx.Process(
-                target=_stage_worker,
-                args=(
-                    model,
-                    stage_payload,
-                    self._in_q,
-                    self._out_q,
-                    self._log_file,
-                    batch_timeout,
-                ),
-            )
+        self._proc = ctx.Process(
+            target=_stage_worker,
+            args=(model, stage_payload, self._in_q, self._out_q, self._log_file, batch_timeout),
+        )
         self._proc.start()
 
     def stop_stage_worker(self) -> None:
@@ -175,11 +129,11 @@ class OmniStage:
                 except Exception as e:
                     self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e)
 
-    def submit(self, payload: dict[str, Any]) -> None:
+    def submit(self, payload: Dict[str, Any]) -> None:
         assert self._in_q is not None
         self._in_q.put(payload)
 
-    def try_collect(self) -> Optional[dict[str, Any]]:
+    def try_collect(self) -> Optional[Dict[str, Any]]:
         assert self._out_q is not None
         try:
             return self._out_q.get_nowait()
@@ -188,7 +142,7 @@ class OmniStage:
 
     def process_engine_inputs(
         self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None
-    ) -> list[Union[OmniTokensPrompt, TextPrompt]]:
+    ) -> List[Union[OmniTokensPrompt, TextPrompt]]:
         """Process the engine input for the stage."""
         if self.custom_process_input_func is None:
             engine_inputs = []
@@ -196,11 +150,9 @@ class OmniStage:
                 raise ValueError("engine_input_source is empty")
             source_stage_id = self.engine_input_source[0]
             source_outputs = stage_list[source_stage_id].engine_outputs
-            if not isinstance(prompt, list):
-                prompt = [prompt]
             multi_modal_data = {
-                source_output.request_id: p.get("multi_modal_data", None)
-                for source_output, p in zip(source_outputs, prompt)
+                source_output.request_id: prompt.get("multi_modal_data", None)
+                for source_output, prompt in zip(source_outputs, prompt)
             }
 
             for source_output in source_outputs:
@@ -208,7 +160,7 @@ class OmniStage:
                     prompt_token_ids=source_output.outputs[0].token_ids,
                     multi_modal_data=(
                         multi_modal_data[source_output.request_id]
-                        if self.requires_multimodal_data and multi_modal_data
+                        if multi_modal_data
                         else None
                     ),
                 )
@@ -218,13 +170,13 @@ class OmniStage:
         else:
             engine_input_source = self.engine_input_source
             return self.custom_process_input_func(
-                stage_list, engine_input_source, prompt, self.requires_multimodal_data
+                stage_list, engine_input_source, prompt
             )
 
 
 def _stage_worker(
     model: str,
-    stage_payload: dict[str, Any],
+    stage_payload: Dict[str, Any],
     in_q: mp.Queue,
     out_q: mp.Queue,
     log_file: Optional[str] = None,
@@ -232,16 +184,15 @@ def _stage_worker(
 ) -> None:
     """Stage worker entry: device setup, LLM init, batching, SHM IPC."""
     import logging as _logging
-    import time as _time
-
-    from vllm_omni.entrypoints.log_utils import (
+    from vllm_omni.entrypoints.omni_llm import OmniStageLLM  # noqa: WPS433
+    from vllm_omni.entrypoints.log_utils import (  # noqa: WPS433
+        log_stage_running_avg,
+        log_stage_batch_stats,
         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
-
+    import queue as _queue
+    import time as _time
     # no inline JSONL/serialization imports; logging handled by utilities
 
     stage_id = stage_payload["stage_id"]
@@ -256,15 +207,11 @@ def _stage_worker(
             stage_log.setLevel(_logging.DEBUG)
             fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log")
             fh.setLevel(_logging.DEBUG)
-            fh.setFormatter(
-                _logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s")
-            )
-
+            fh.setFormatter(_logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s"))
             class _StageFilter(_logging.Filter):
                 def filter(self, record: _logging.LogRecord) -> bool:
                     setattr(record, "stage", stage_id)
                     return True
-
             fh.addFilter(_StageFilter())
             stage_log.addHandler(fh)
     except Exception:
@@ -286,11 +233,7 @@ def _stage_worker(
         _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
 
     # Init LLM
-    _logging.getLogger(__name__).debug(
-        "[Stage-%s] Initializing engine with args keys=%s",
-        stage_id,
-        list(engine_args.keys()),
-    )
+    _logging.getLogger(__name__).debug("[Stage-%s] Initializing engine with args keys=%s", stage_id, list(engine_args.keys()))
     stage_engine = OmniStageLLM(model=model, **engine_args)
     _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
     # Signal readiness to orchestrator
@@ -308,7 +251,7 @@ def _stage_worker(
             break
 
         max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
-        batch_tasks: list[dict[str, Any]] = [task]
+        batch_tasks: List[Dict[str, Any]] = [task]
         if max_batch_size > 1:
             while len(batch_tasks) < max_batch_size:
                 if not in_q.empty():
@@ -320,11 +263,11 @@ def _stage_worker(
                 else:
                     break
 
-        batch_request_ids: list[Any] = []
-        batch_engine_inputs: list[Any] = []
-        _rx_bytes_by_rid: dict[Any, int] = {}
-        _rx_decode_ms_by_rid: dict[Any, float] = {}
-        _in_flight_ms_by_rid: dict[Any, float] = {}
+        batch_request_ids: List[Any] = []
+        batch_engine_inputs: List[Any] = []
+        _rx_bytes_by_rid: Dict[Any, int] = {}
+        _rx_decode_ms_by_rid: Dict[Any, float] = {}
+        _in_flight_ms_by_rid: Dict[Any, float] = {}
         for t in batch_tasks:
             rid = t["request_id"]
             try:
@@ -335,7 +278,9 @@ def _stage_worker(
                     _in_flight_ms_by_rid[rid] = 0.0
             except Exception:
                 _in_flight_ms_by_rid[rid] = 0.0
-            ein, _rx_metrics = maybe_load_from_ipc_with_metrics(t, obj_key="engine_inputs", shm_key="engine_inputs_shm")
+            ein, _rx_metrics = maybe_load_from_ipc_with_metrics(
+                t, obj_key="engine_inputs", shm_key="engine_inputs_shm"
+            )
             _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
             _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
             batch_request_ids.append(rid)
@@ -346,21 +291,13 @@ def _stage_worker(
             else:
                 _logging.getLogger(__name__).exception("[Stage-%s] Invalid engine input type: %s", stage_id, type(ein))
         sampling_params = batch_tasks[0]["sampling_params"]
-        _logging.getLogger(__name__).debug(
-            "[Stage-%s] Received batch size=%d, request_ids=%s",
-            stage_id,
-            len(batch_tasks),
-            batch_request_ids,
-        )
+        _logging.getLogger(__name__).debug("[Stage-%s] Received batch size=%d, request_ids=%s", stage_id, len(batch_tasks), batch_request_ids)
         print("--------------------------------", flush=True)
-        print(
-            f"[Stage-{stage_id}] Received batch size={len(batch_tasks)}, request_ids={batch_request_ids}",
-            flush=True,
-        )
+        print(f"[Stage-{stage_id}] Received batch size={len(batch_tasks)}, request_ids={batch_request_ids}", flush=True)
         print("--------------------------------", flush=True)
         try:
             _batch_seq += 1
-            gen_outputs: list[Any] = []
+            gen_outputs: List[Any] = []
             _gen_t0 = _time.time()
             for ro in stage_engine.generate(batch_engine_inputs, sampling_params, use_tqdm=False):
                 gen_outputs.append(ro)
@@ -368,8 +305,8 @@ def _stage_worker(
             _gen_ms = (_gen_t1 - _gen_t0) * 1000.0
 
             # Group outputs per request id with fallback
-            req_to_outputs: dict[Any, list[Any]] = {rid: [] for rid in batch_request_ids}
-            unmapped: list[Any] = []
+            req_to_outputs: Dict[Any, List[Any]] = {rid: [] for rid in batch_request_ids}
+            unmapped: List[Any] = []
             for ro in gen_outputs:
                 rid = getattr(ro, "request_id", None)
                 if rid in req_to_outputs:
@@ -389,25 +326,12 @@ def _stage_worker(
                 _num_tokens = count_tokens_from_outputs(_r_outputs)
                 _agg_total_tokens += _num_tokens
                 _agg_total_gen_time_ms += _gen_ms
+                _tokens_per_s = (_num_tokens * 1000.0 / _gen_ms) if _gen_ms > 0 else 0.0
 
             if _stats_file:
-                _avg_tokens_per_s = (
-                    (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0
-                )
-                log_stage_running_avg(
-                    _stats_file,
-                    stage_id,
-                    int(_agg_total_tokens),
-                    float(_agg_total_gen_time_ms),
-                    float(_avg_tokens_per_s),
-                )
-                log_stage_batch_stats(
-                    _stats_file,
-                    stage_id,
-                    len(batch_tasks),
-                    float(_gen_ms),
-                    list(batch_request_ids),
-                )
+                _avg_tokens_per_s = (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0
+                log_stage_running_avg(_stats_file, stage_id, int(_agg_total_tokens), float(_agg_total_gen_time_ms), float(_avg_tokens_per_s))
+                log_stage_batch_stats(_stats_file, stage_id, len(batch_tasks), float(_gen_ms), list(batch_request_ids))
 
             # Emit per-request results
             for rid in batch_request_ids:
@@ -430,281 +354,25 @@ def _stage_worker(
                             len(batch_tasks),
                             r_outputs,
                             float(_gen_ms),
-                            int(_metrics["rx_transfer_bytes"]),  # type: ignore[index]
-                            float(_metrics["rx_decode_time_ms"]),  # type: ignore[index]
+                            int(_metrics["rx_transfer_bytes"]),   # type: ignore[index]
+                            float(_metrics["rx_decode_time_ms"]), # type: ignore[index]
                         )
                     if use_shm:
-                        out_q.put(
-                            {
-                                "request_id": rid,
-                                "stage_id": stage_id,
-                                "engine_outputs_shm": payload,
-                                "metrics": _metrics,
-                            }
-                        )
-                    else:
-                        out_q.put(
-                            {
-                                "request_id": rid,
-                                "stage_id": stage_id,
-                                "engine_outputs": payload,
-                                "metrics": _metrics,
-                            }
-                        )
-                except Exception:
-                    out_q.put(
-                        {
-                            "request_id": rid,
-                            "stage_id": stage_id,
-                            "engine_outputs": r_outputs,
-                            "metrics": {
-                                "num_tokens_out": int(count_tokens_from_outputs(r_outputs)),
-                                "stage_gen_time_ms": _gen_ms,
-                                "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
-                                "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
-                                "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
-                            },
-                        }
-                    )
-                _logging.getLogger(__name__).debug(
-                    "[Stage-%s] Enqueued result for request %s to downstream",
-                    stage_id,
-                    rid,
-                )
-        except Exception as e:
-            _logging.getLogger(__name__).exception("[Stage-%s] Failed on batch %s: %s", stage_id, batch_request_ids, e)
-            for rid in batch_request_ids:
-                out_q.put(
-                    {
-                        "request_id": rid,
-                        "stage_id": stage_id,
-                        "error": str(e),
-                    }
-                )
-
-
-def _stage_worker_async_entry(
-    omni_stage: OmniStage,
-    model: str,
-    stage_payload: dict[str, Any],
-    in_q: mp.Queue,
-    out_q: mp.Queue,
-    log_file: Optional[str] = None,
-    batch_timeout: int = 10,
-) -> None:
-    asyncio.run(_stage_worker_async(omni_stage, model, stage_payload, in_q, out_q, log_file, batch_timeout))
-
-
-async def _stage_worker_async(
-    omni_stage: OmniStage,
-    model: str,
-    stage_payload: dict[str, Any],
-    in_q: mp.Queue,
-    out_q: mp.Queue,
-    log_file: Optional[str] = None,
-    batch_timeout: int = 10,
-) -> None:
-    """Stage worker entry: device setup, LLM init, batching, SHM IPC."""
-    import logging as _logging
-    import time as _time
-
-    from vllm_omni.entrypoints.async_omni_llm import AsyncOmniStageLLM
-    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,
-    )
-
-    # no inline JSONL/serialization imports; logging handled by utilities
-
-    stage_id = stage_payload["stage_id"]
-    engine_args = stage_payload.get("engine_args", {})
-    runtime_cfg = stage_payload.get("runtime", {})
-    shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
-
-    # Per-stage file logger (optional)
-    try:
-        if log_file:
-            stage_log = _logging.getLogger(__name__)
-            stage_log.setLevel(_logging.DEBUG)
-            fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log")
-            fh.setLevel(_logging.DEBUG)
-            fh.setFormatter(
-                _logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s")
-            )  # noqa: E501
-
-            class _StageFilter(_logging.Filter):
-                def filter(self, record: _logging.LogRecord) -> bool:
-                    setattr(record, "stage", stage_id)
-                    return True
-
-            fh.addFilter(_StageFilter())
-            stage_log.addHandler(fh)
-    except Exception:
-        pass
-
-    # Stage stats JSONL file
-    _stats_file = f"{log_file}.stage{stage_id}.stats.jsonl" if log_file else None
-
-    # Aggregates for running average
-    _agg_total_tokens = 0
-    _agg_total_gen_time_ms = 0.0
-    # Monotonic batch id per stage process for orchestrator dedup on time
-    # aggregation
-    _batch_seq = 0
-
-    # Device mapping
-    try:
-        set_stage_gpu_devices(stage_id, runtime_cfg.get("devices"))
-    except Exception as e:
-        _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
-
-    # Init LLM
-    _logging.getLogger(__name__).debug(
-        "[Stage-%s] Initializing engine with args keys=%s",
-        stage_id,
-        list(engine_args.keys()),
-    )
-    omni_engine_args = AsyncOmniEngineArgs(model=model, **engine_args)
-    usage_context = UsageContext.OPENAI_API_SERVER
-    vllm_config = omni_engine_args.create_engine_config(usage_context=usage_context)
-    stage_engine = AsyncOmniStageLLM.from_vllm_config(
-        vllm_config=vllm_config,
-        usage_context=usage_context,
-        engine_args=omni_engine_args,
-    )
-    omni_stage.set_async_engine(stage_engine)
-    # Don't keep the dummy data in memory
-    await stage_engine.reset_mm_cache()
-    _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
-    # Signal readiness to orchestrator and send vllm_config back to main process
-    try:
-        # Send vllm_config back to main process so it can be accessed via
-        # get_vllm_config(). This is needed because async_engine is only available
-        # in the worker process
-
-        # input_preprocessor = await stage_engine.get_input_preprocessor()
-        out_q.put(
-            {
-                "type": "stage_ready",
-                "stage_id": stage_id,
-                "vllm_config": vllm_config,
-                "tokenizer": getattr(stage_engine, "tokenizer", None),
-                "is_tracing_enabled": await stage_engine.is_tracing_enabled(),
-                # "input_preprocessor": input_preprocessor,
-            }
-        )
-    except Exception as e:
-        _logging.getLogger(__name__).warning("[Stage-%s] Failed to send stage ready signal: %s", stage_id, e)
-
-    # Batch processing loop
-    while True:
-        task = in_q.get()
-        _recv_dequeue_ts = _time.time()
-        if task is None:
-            _logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id)
-            break
-
-        _rx_bytes_by_rid: dict[Any, int] = {}
-        _rx_decode_ms_by_rid: dict[Any, float] = {}
-        _in_flight_ms_by_rid: dict[Any, float] = {}
-
-        rid = task["request_id"]
-        try:
-            sent_ts = float(task.get("sent_ts", None)) if isinstance(task, dict) else None
-            if sent_ts is not None:
-                _in_flight_ms_by_rid[rid] = (_recv_dequeue_ts - sent_ts) * 1000.0
-            else:
-                _in_flight_ms_by_rid[rid] = 0.0
-        except Exception:
-            _in_flight_ms_by_rid[rid] = 0.0
-        ein, _rx_metrics = maybe_load_from_ipc_with_metrics(task, obj_key="engine_inputs", shm_key="engine_inputs_shm")
-        _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
-        _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
-
-        sampling_params = task["sampling_params"]
-        _logging.getLogger(__name__).debug("[Stage-%s] Received batch size=1, request_ids=%d", stage_id, rid)
-        print("--------------------------------", flush=True)
-        print(f"[Stage-{stage_id}] Received batch size=1, request_ids={rid}", flush=True)
-        print("--------------------------------", flush=True)
-        try:
-            _batch_seq += 1
-            _gen_t0 = _time.time()
-            if isinstance(ein, list):
-                ein = ein[0]
-
-            async for res in stage_engine.generate(ein, sampling_params, rid):
-                gen_output = res
-            _gen_t1 = _time.time()
-            _gen_ms = (_gen_t1 - _gen_t0) * 1000.0
-
-            r_outputs = [gen_output]
-            _num_tokens = count_tokens_from_outputs(r_outputs)
-            _agg_total_tokens += _num_tokens
-            _agg_total_gen_time_ms += _gen_ms
-
-            if _stats_file:
-                _avg_tokens_per_s = (
-                    (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0
-                )
-                log_stage_running_avg(
-                    _stats_file,
-                    stage_id,
-                    int(_agg_total_tokens),
-                    float(_agg_total_gen_time_ms),
-                    float(_avg_tokens_per_s),
-                )
-                log_stage_batch_stats(_stats_file, stage_id, 1, float(_gen_ms), [rid])
-
-            try:
-                use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes)
-                _metrics = {
-                    "num_tokens_out": int(count_tokens_from_outputs(r_outputs)),
-                    "stage_gen_time_ms": _gen_ms,
-                    "batch_id": int(_batch_seq),
-                    "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
-                    "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
-                    "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
-                }
-                if _stats_file:
-                    compute_and_log_stage_request_stats(
-                        _stats_file,
-                        stage_id,
-                        rid,
-                        1,
-                        r_outputs,
-                        float(_gen_ms),
-                        int(_metrics["rx_transfer_bytes"]),  # type: ignore[index]
-                        float(_metrics["rx_decode_time_ms"]),  # type: ignore[index]
-                    )
-                if use_shm:
-                    out_q.put(
-                        {
+                        out_q.put({
                             "request_id": rid,
                             "stage_id": stage_id,
                             "engine_outputs_shm": payload,
                             "metrics": _metrics,
-                        }
-                    )
-                else:
-                    out_q.put(
-                        {
+                        })
+                    else:
+                        out_q.put({
                             "request_id": rid,
                             "stage_id": stage_id,
                             "engine_outputs": payload,
                             "metrics": _metrics,
-                        }
-                    )
-            except Exception as e:
-                _logging.getLogger(__name__).exception(
-                    "[Stage-%s] Failed to enqueue result for request %s: %s",
-                    stage_id,
-                    rid,
-                    e,
-                )
-                out_q.put(
-                    {
+                        })
+                except Exception:
+                    out_q.put({
                         "request_id": rid,
                         "stage_id": stage_id,
                         "engine_outputs": r_outputs,
@@ -715,16 +383,13 @@ async def _stage_worker_async(
                             "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
                             "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
                         },
-                    }
-                )
-            _logging.getLogger(__name__).debug("[Stage-%s] Enqueued result for request %s to downstream", stage_id, rid)
-
+                    })
+                _logging.getLogger(__name__).debug("[Stage-%s] Enqueued result for request %s to downstream", stage_id, rid)
         except Exception as e:
-            _logging.getLogger(__name__).exception("[Stage-%s] Failed on request %s: %s", stage_id, rid, e)
-            out_q.put(
-                {
+            _logging.getLogger(__name__).exception("[Stage-%s] Failed on batch %s: %s", stage_id, batch_request_ids, e)
+            for rid in batch_request_ids:
+                out_q.put({
                     "request_id": rid,
                     "stage_id": stage_id,
                     "error": str(e),
-                }
-            )
+                })

vllm_omni/entrypoints/openai/api_server.py

@@ -1,313 +0,0 @@
-import multiprocessing
-import multiprocessing.forkserver as forkserver
-import os
-from argparse import Namespace
-from collections.abc import AsyncIterator
-from contextlib import asynccontextmanager
-from http import HTTPStatus
-from typing import Any, Optional
-
-import vllm.envs as envs
-from fastapi import Depends, HTTPException, Request
-from fastapi.responses import JSONResponse, StreamingResponse
-from starlette.datastructures import State
-from vllm.config import VllmConfig
-from vllm.engine.protocol import EngineClient
-from vllm.entrypoints.chat_utils import load_chat_template, resolve_hf_chat_template, resolve_mistral_chat_template
-from vllm.entrypoints.launcher import serve_http
-from vllm.entrypoints.logger import RequestLogger
-from vllm.entrypoints.openai.api_server import (
-    base,
-    build_app,
-    load_log_config,
-    maybe_register_tokenizer_info_endpoint,
-    router,
-    setup_server,
-    validate_json_request,
-)
-from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ChatCompletionResponse, ErrorResponse
-from vllm.entrypoints.openai.serving_models import BaseModelPath, LoRAModulePath, OpenAIServingModels
-from vllm.entrypoints.openai.tool_parsers import ToolParserManager
-
-# yapf conflicts with isort for this block
-# yapf: disable
-# yapf: enable
-from vllm.entrypoints.tool_server import DemoToolServer, MCPToolServer, ToolServer
-from vllm.entrypoints.utils import load_aware_call, with_cancellation
-from vllm.logger import init_logger
-from vllm.transformers_utils.tokenizer import MistralTokenizer
-from vllm.utils import decorate_logs
-
-from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
-from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
-
-logger = init_logger(__name__)
-
-
-async def omni_run_server(args, **uvicorn_kwargs) -> None:
-    """Run a single-worker API server."""
-
-    # Add process-specific prefix to stdout and stderr.
-    decorate_logs("APIServer")
-
-    listen_address, sock = setup_server(args)
-    await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
-
-
-async def omni_run_server_worker(listen_address, sock, args, client_config=None, **uvicorn_kwargs) -> None:
-    """Run a single API server worker."""
-
-    if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:
-        ToolParserManager.import_tool_parser(args.tool_parser_plugin)
-
-    # Load logging config for uvicorn if specified
-    log_config = load_log_config(args.log_config_file)
-    if log_config is not None:
-        uvicorn_kwargs["log_config"] = log_config
-
-    async with build_async_omni_llm(
-        args,
-        client_config=client_config,
-    ) as engine_client:
-        maybe_register_tokenizer_info_endpoint(args)
-        app = build_app(args)
-
-        vllm_config = await engine_client.get_vllm_config()
-        await omni_init_app_state(engine_client, vllm_config, app.state, args)
-
-        logger.info(
-            "Starting vLLM API server %d on %s",
-            vllm_config.parallel_config._api_process_rank,
-            listen_address,
-        )
-        shutdown_task = await serve_http(
-            app,
-            sock=sock,
-            enable_ssl_refresh=args.enable_ssl_refresh,
-            host=args.host,
-            port=args.port,
-            log_level=args.uvicorn_log_level,
-            # NOTE: When the 'disable_uvicorn_access_log' value is True,
-            # no access log will be output.
-            access_log=not args.disable_uvicorn_access_log,
-            timeout_keep_alive=envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE,
-            ssl_keyfile=args.ssl_keyfile,
-            ssl_certfile=args.ssl_certfile,
-            ssl_ca_certs=args.ssl_ca_certs,
-            ssl_cert_reqs=args.ssl_cert_reqs,
-            h11_max_incomplete_event_size=args.h11_max_incomplete_event_size,
-            h11_max_header_count=args.h11_max_header_count,
-            **uvicorn_kwargs,
-        )
-
-    # NB: Await server shutdown only after the backend context is exited
-    try:
-        await shutdown_task
-    finally:
-        sock.close()
-
-
-@asynccontextmanager
-async def build_async_omni_llm(
-    args: Namespace,
-    *,
-    disable_frontend_multiprocessing: Optional[bool] = None,
-    client_config: Optional[dict[str, Any]] = None,
-) -> AsyncIterator[EngineClient]:
-    if os.getenv("VLLM_WORKER_MULTIPROC_METHOD") == "forkserver":
-        # The executor is expected to be mp.
-        # Pre-import heavy modules in the forkserver process
-        logger.debug("Setup forkserver with pre-imports")
-        multiprocessing.set_start_method("forkserver")
-        multiprocessing.set_forkserver_preload(["vllm.v1.engine.async_llm"])
-        forkserver.ensure_running()
-        logger.debug("Forkserver setup complete!")
-
-    # Context manager to handle async_omni_llm lifecycle
-    # Ensures everything is shutdown and cleaned up on error/exit
-    async with build_async_omni_llm_from_stage_config(
-        args,
-        disable_frontend_multiprocessing=disable_frontend_multiprocessing,
-    ) as async_omni_llm:
-        yield async_omni_llm
-
-
-@asynccontextmanager
-async def build_async_omni_llm_from_stage_config(
-    args: Namespace,
-    *,
-    disable_frontend_multiprocessing: bool = False,
-    client_config: Optional[dict[str, Any]] = None,
-) -> AsyncIterator[EngineClient]:
-    """
-    Create AsyncOmniLLM, either:
-        - in-process using the AsyncOmniLLM Directly
-        - multiprocess using AsyncOmniLLM RPC
-
-    Returns the AsyncOmniLLM or None if the creation failed.
-    """
-
-    # V1 AsyncLLM.
-    assert envs.VLLM_USE_V1
-
-    if disable_frontend_multiprocessing:
-        logger.warning(
-            "V1 is enabled, but got --disable-frontend-multiprocessing. "
-            "To disable frontend multiprocessing, set VLLM_USE_V1=0."
-        )
-
-    async_omni_llm: Optional[EngineClient] = None
-
-    try:
-        if getattr(args, "stage_configs_path", None):
-            async_omni_llm = AsyncOmniLLM(model=args.model, stage_configs_path=args.stage_configs_path)
-        else:
-            async_omni_llm = AsyncOmniLLM(model=args.model)
-
-        # # Don't keep the dummy data in memory
-        # await async_llm.reset_mm_cache()
-
-        yield async_omni_llm
-    finally:
-        if async_omni_llm:
-            async_omni_llm.shutdown()
-
-
-async def omni_init_app_state(
-    engine_client: EngineClient,
-    vllm_config: VllmConfig,
-    state: State,
-    args: Namespace,
-) -> None:
-    if args.served_model_name is not None:
-        served_model_names = args.served_model_name
-    else:
-        served_model_names = [args.model]
-
-    if args.enable_log_requests:
-        request_logger = RequestLogger(max_log_len=args.max_log_len)
-    else:
-        request_logger = None
-
-    base_model_paths = [BaseModelPath(name=name, model_path=args.model) for name in served_model_names]
-    state.engine_client = engine_client
-    state.log_stats = not args.disable_log_stats
-    state.vllm_config = vllm_config
-    model_config = vllm_config.model_config
-    state.log_stats = not args.disable_log_stats
-
-    # For omni models
-    state.stage_configs = engine_client.stage_configs
-
-    resolved_chat_template = load_chat_template(args.chat_template)
-    if resolved_chat_template is not None:
-        # Get the tokenizer to check official template
-        tokenizer = await engine_client.get_tokenizer()
-
-        if isinstance(tokenizer, MistralTokenizer):
-            # The warning is logged in resolve_mistral_chat_template.
-            resolved_chat_template = resolve_mistral_chat_template(chat_template=resolved_chat_template)
-        else:
-            hf_chat_template = resolve_hf_chat_template(
-                tokenizer=tokenizer,
-                chat_template=None,
-                tools=None,
-                model_config=vllm_config.model_config,
-            )
-
-            if hf_chat_template != resolved_chat_template:
-                logger.warning(
-                    "Using supplied chat template: %s\nIt is different from official chat template '%s'. This discrepancy may lead to performance degradation.",  # noqa: E501
-                    resolved_chat_template,
-                    args.model,
-                )
-
-    if args.tool_server == "demo":
-        tool_server: Optional[ToolServer] = DemoToolServer()
-        assert isinstance(tool_server, DemoToolServer)
-        await tool_server.init_and_validate()
-    elif args.tool_server:
-        tool_server = MCPToolServer()
-        await tool_server.add_tool_server(args.tool_server)
-    else:
-        tool_server = None
-
-    # Merge default_mm_loras into the static lora_modules
-    default_mm_loras = vllm_config.lora_config.default_mm_loras if vllm_config.lora_config is not None else {}
-
-    lora_modules = args.lora_modules
-    if default_mm_loras:
-        default_mm_lora_paths = [
-            LoRAModulePath(
-                name=modality,
-                path=lora_path,
-            )
-            for modality, lora_path in default_mm_loras.items()
-        ]
-        if args.lora_modules is None:
-            lora_modules = default_mm_lora_paths
-        else:
-            lora_modules += default_mm_lora_paths
-
-    state.openai_serving_models = OpenAIServingModels(
-        engine_client=engine_client,
-        model_config=model_config,
-        base_model_paths=base_model_paths,
-        lora_modules=lora_modules,
-    )
-    await state.openai_serving_models.init_static_loras()
-    state.openai_serving_chat = OmniOpenAIServingChat(
-        engine_client,
-        model_config,
-        state.openai_serving_models,
-        args.response_role,
-        request_logger=request_logger,
-        chat_template=resolved_chat_template,
-        chat_template_content_format=args.chat_template_content_format,
-        trust_request_chat_template=args.trust_request_chat_template,
-        return_tokens_as_token_ids=args.return_tokens_as_token_ids,
-        enable_auto_tools=args.enable_auto_tool_choice,
-        exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
-        tool_parser=args.tool_call_parser,
-        reasoning_parser=args.structured_outputs_config.reasoning_parser,
-        enable_prompt_tokens_details=args.enable_prompt_tokens_details,
-        enable_force_include_usage=args.enable_force_include_usage,
-        enable_log_outputs=args.enable_log_outputs,
-        log_error_stack=args.log_error_stack,
-    )
-
-    state.enable_server_load_tracking = args.enable_server_load_tracking
-    state.server_load_metrics = 0
-
-
-def Omnichat(request: Request) -> Optional[OmniOpenAIServingChat]:
-    return request.app.state.openai_serving_chat
-
-
-@router.post(
-    "/v1/chat/completions",
-    dependencies=[Depends(validate_json_request)],
-    responses={
-        HTTPStatus.OK.value: {"content": {"text/event-stream": {}}},
-        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
-        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
-        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
-    },
-)
-@with_cancellation
-@load_aware_call
-async def create_chat_completion(request: ChatCompletionRequest, raw_request: Request):
-    handler = Omnichat(raw_request)
-    if handler is None:
-        return base(raw_request).create_error_response(message="The model does not support Chat Completions API")
-    try:
-        generator = await handler.create_chat_completion(request, raw_request)
-    except Exception as e:
-        raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e
-    if isinstance(generator, ErrorResponse):
-        return JSONResponse(content=generator.model_dump(), status_code=generator.error.code)
-
-    elif isinstance(generator, ChatCompletionResponse):
-        return JSONResponse(content=generator.model_dump())
-
-    return StreamingResponse(content=generator, media_type="text/event-stream")

vllm_omni/entrypoints/openai/serving_chat.py

@@ -1,779 +0,0 @@
-import asyncio
-import base64
-import json
-import time
-import uuid
-from collections.abc import AsyncGenerator, AsyncIterator, Sequence
-from datetime import datetime, timedelta, timezone
-from io import BytesIO
-from typing import Any, Callable, Optional, Union
-
-import jinja2
-from fastapi import Request
-from pydantic import TypeAdapter
-
-try:
-    import soundfile
-except ImportError:
-    soundfile = None
-
-from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio
-from vllm.entrypoints.chat_utils import (
-    ChatCompletionMessageParam,
-    ChatTemplateContentFormatOption,
-    ConversationMessage,
-    apply_hf_chat_template,
-    apply_mistral_chat_template,
-    get_history_tool_calls_cnt,
-    make_tool_call_id,
-    resolve_chat_template_content_format,
-)
-from vllm.entrypoints.harmony_utils import parse_chat_output
-from vllm.entrypoints.openai.protocol import (
-    ChatCompletionNamedToolChoiceParam,
-    ChatCompletionRequest,
-    ChatCompletionResponse,
-    ChatCompletionResponseChoice,
-    ChatMessage,
-    ErrorResponse,
-    FunctionCall,
-    FunctionDefinition,
-    PromptTokenUsageInfo,
-    RequestResponseMetadata,
-    ToolCall,
-    UsageInfo,
-)
-from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
-from vllm.entrypoints.openai.serving_engine import (
-    ChatLikeRequest,
-    EngineTokensPrompt,
-    RequestPrompt,
-    ResponsesRequest,
-    TextTokensPrompt,
-    clamp_prompt_logprobs,
-    is_list_of,
-)
-from vllm.entrypoints.openai.tool_parsers import ToolParser
-from vllm.entrypoints.openai.tool_parsers.mistral_tool_parser import MistralToolCall
-from vllm.inputs.data import PromptType
-from vllm.logger import init_logger
-from vllm.lora.request import LoRARequest
-from vllm.outputs import RequestOutput
-from vllm.sampling_params import SamplingParams
-from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer
-from vllm.transformers_utils.tokenizers import (
-    maybe_serialize_tool_calls,
-    truncate_tool_call_ids,
-    validate_request_params,
-)
-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__)
-
-
-class OmniOpenAIServingChat(OpenAIServingChat):
-    async def create_chat_completion(
-        self,
-        request: ChatCompletionRequest,
-        raw_request: Optional[Request] = None,
-    ) -> Union[AsyncGenerator[str, None], ChatCompletionResponse, ErrorResponse]:
-        """
-        Chat Completion API similar to OpenAI's API.
-
-        See https://platform.openai.com/docs/api-reference/chat/create
-        for the API specification. This API mimics the OpenAI
-        Chat Completion API.
-        """
-        error_check_ret = await self._check_model(request)
-        if error_check_ret is not None:
-            logger.error("Error with model %s", error_check_ret)
-            return error_check_ret
-
-        # If the engine is dead, raise the engine's DEAD_ERROR.
-        # This is required for the streaming case, where we return a
-        # success status before we actually start generating text :).
-        if self.engine_client.errored:
-            raise self.engine_client.dead_error
-
-        try:
-            lora_request = self._maybe_get_adapters(request, supports_default_mm_loras=True)
-
-            model_name = self.models.model_name(lora_request)
-
-            tokenizer = await self.engine_client.get_tokenizer()
-
-            tool_parser = self.tool_parser
-
-            if isinstance(tokenizer, MistralTokenizer):
-                # because of issues with pydantic we need to potentially
-                # re-serialize the tool_calls field of the request
-                # for more info: see comment in `maybe_serialize_tool_calls`
-                maybe_serialize_tool_calls(request)
-                truncate_tool_call_ids(request)
-                validate_request_params(request)
-
-            if (
-                request.tool_choice == "auto"
-                and not (self.enable_auto_tools and tool_parser is not None)
-                and not isinstance(tokenizer, MistralTokenizer)
-                and not self.use_harmony
-            ):
-                # for hf tokenizers, "auto" tools requires
-                # --enable-auto-tool-choice and --tool-call-parser
-                return self.create_error_response(
-                    '"auto" tool choice requires --enable-auto-tool-choice and --tool-call-parser to be set'
-                )
-
-            if request.tools is None or (request.tool_choice == "none" and self.exclude_tools_when_tool_choice_none):
-                tool_dicts = None
-            else:
-                tool_dicts = [tool.model_dump() for tool in request.tools]
-
-            # Common case.
-            request_chat_template = request.chat_template
-            chat_template_kwargs = request.chat_template_kwargs
-            if not self.trust_request_chat_template and (
-                request_chat_template is not None
-                or (chat_template_kwargs and chat_template_kwargs.get("chat_template") is not None)
-            ):
-                return self.create_error_response(
-                    "Chat template is passed with request, but --trust-request-chat-template is not set. "
-                    "Refused request with untrusted chat template."
-                )
-            (
-                conversation,
-                request_prompts,
-                engine_prompts,
-            ) = await self._preprocess_chat(
-                request,
-                tokenizer,
-                request.messages,
-                chat_template=request_chat_template or self.chat_template,
-                chat_template_content_format=self.chat_template_content_format,
-                add_generation_prompt=request.add_generation_prompt,
-                continue_final_message=request.continue_final_message,
-                tool_dicts=tool_dicts,
-                documents=request.documents,
-                chat_template_kwargs=request.chat_template_kwargs,
-                tool_parser=tool_parser,
-                add_special_tokens=request.add_special_tokens,
-            )
-
-        except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e:
-            logger.exception("Error in preprocessing prompt inputs")
-            return self.create_error_response(f"{e} {e.__cause__}")
-
-        request_id = f"chatcmpl-{self._base_request_id(raw_request, request.request_id)}"
-
-        request_metadata = RequestResponseMetadata(request_id=request_id)
-        if raw_request:
-            raw_request.state.request_metadata = request_metadata
-
-        # Schedule the request and get the result generator.
-        generators: list[AsyncGenerator[RequestOutput, None]] = []
-        try:
-            for i, engine_prompt in enumerate(engine_prompts):
-                if hasattr(request, "sampling_params_list"):
-                    sampling_params_list = self._to_sampling_params_list(request.sampling_params_list)
-                else:
-                    sampling_params_list = None
-
-                self._log_inputs(
-                    request_id,
-                    request_prompts[i],
-                    params_list=sampling_params_list,
-                    lora_request=lora_request,
-                )
-
-                trace_headers = None if raw_request is None else await self._get_trace_headers(raw_request.headers)
-
-                generator = self.engine_client.generate(
-                    prompt=engine_prompt,
-                    request_id=request_id,
-                    sampling_params_list=sampling_params_list,
-                    lora_request=lora_request,
-                    trace_headers=trace_headers,
-                    priority=request.priority,
-                )
-
-                generators.append(generator)
-        except ValueError as e:
-            # TODO: Use a vllm-specific Validation Error
-            return self.create_error_response(str(e))
-
-        assert len(generators) == 1
-        (result_generator,) = generators
-
-        # Streaming response
-        if request.stream:
-            return self.chat_completion_stream_generator(
-                request,
-                result_generator,
-                request_id,
-                model_name,
-                conversation,
-                tokenizer,
-                request_metadata,
-                enable_force_include_usage=self.enable_force_include_usage,
-            )
-
-        try:
-            return await self.chat_completion_full_generator(
-                request,
-                result_generator,
-                request_id,
-                model_name,
-                conversation,
-                tokenizer,
-                request_metadata,
-            )
-        except ValueError as e:
-            # 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:
-            if isinstance(sampling_params, dict):
-                final_sampling_params_list.append(SamplingParams(**sampling_params))
-            elif isinstance(sampling_params, SamplingParams):
-                final_sampling_params_list.append(sampling_params)
-            else:
-                raise ValueError(f"Invalid sampling params: {sampling_params}")
-        return final_sampling_params_list
-
-    def _log_inputs(
-        self,
-        request_id: str,
-        inputs: Union[RequestPrompt, PromptType],
-        params_list: Optional[list[SamplingParams]],
-        lora_request: Optional[LoRARequest],
-    ) -> None:
-        if self.request_logger is None:
-            return
-        prompt, prompt_token_ids, prompt_embeds = None, None, None
-        if isinstance(inputs, str):
-            prompt = inputs
-        elif isinstance(inputs, list):
-            prompt_token_ids = inputs
-        else:
-            prompt = getattr(inputs, "prompt", None)
-            prompt_token_ids = getattr(inputs, "prompt_token_ids", None)
-
-        logger.info(
-            "Received request %s: prompt: %r, params_list: %s, prompt_token_ids: %s, prompt_embeds shape: %s, lora_request: %s.",  # noqa: E501
-            request_id,
-            prompt,
-            params_list,
-            prompt_token_ids,
-            prompt_embeds.shape if prompt_embeds is not None else None,
-            lora_request,
-        )
-
-    async def chat_completion_full_generator(
-        self,
-        request: ChatCompletionRequest,
-        result_generator: AsyncIterator[RequestOutput],
-        request_id: str,
-        model_name: str,
-        conversation: list[ConversationMessage],
-        tokenizer: AnyTokenizer,
-        request_metadata: RequestResponseMetadata,
-    ) -> Union[ErrorResponse, ChatCompletionResponse]:
-        created_time = int(time.time())
-        final_res: Optional[RequestOutput] = None
-
-        final_outputs: list[OmniRequestOutput] = []
-        try:
-            async for res in result_generator:
-                final_outputs.append(res)
-        except asyncio.CancelledError:
-            return self.create_error_response("Client disconnected")
-        except ValueError as e:
-            # TODO: Use a vllm-specific Validation Error
-            return self.create_error_response(str(e))
-
-        assert final_outputs is not None
-
-        choices: list[ChatCompletionResponseChoice] = []
-
-        usage = UsageInfo(prompt_tokens=0, completion_tokens=0, total_tokens=0)
-        role = self.get_chat_request_role(request)
-        prompt_logprobs = None
-        prompt_token_ids = None
-        kv_transfer_params = None
-
-        for omni_outputs in final_outputs:
-            choices_data = []
-            if omni_outputs.final_output_type == "text":
-                (
-                    choices_data,
-                    usage,
-                    prompt_logprobs,
-                    prompt_token_ids,
-                    kv_transfer_params,
-                ) = self._create_text_choice(request, omni_outputs, tokenizer, conversation, role)
-            elif omni_outputs.final_output_type == "audio":
-                choices_data = self._create_audio_choice(omni_outputs, role)
-            else:
-                logger.warning(f"Unsupported final output type: {omni_outputs.final_output_type}")
-                continue
-            choices.extend(choices_data)
-
-        response = ChatCompletionResponse(
-            id=request_id,
-            created=created_time,
-            model=model_name,
-            choices=choices,
-            usage=usage,
-            prompt_logprobs=prompt_logprobs,
-            prompt_token_ids=prompt_token_ids,
-            kv_transfer_params=kv_transfer_params,
-        )
-
-        # Log complete response if output logging is enabled
-        if self.enable_log_outputs and self.request_logger:
-            for choice in choices:
-                output_text = ""
-                if choice.message.content:
-                    output_text = choice.message.content
-                elif choice.message.tool_calls:
-                    # For tool calls, log the function name and arguments
-                    tool_call_descriptions = []
-                    for tc in choice.message.tool_calls:
-                        if hasattr(tc.function, "name") and hasattr(tc.function, "arguments"):
-                            tool_call_descriptions.append(f"{tc.function.name}({tc.function.arguments})")
-                    tool_calls_str = ", ".join(tool_call_descriptions)
-                    output_text = f"[tool_calls: {tool_calls_str}]"
-
-                if output_text:
-                    # Get the corresponding output token IDs
-                    output_token_ids = None
-                    if choice.index < len(final_res.outputs):
-                        output_token_ids = final_res.outputs[choice.index].token_ids
-
-                    self.request_logger.log_outputs(
-                        request_id=request_id,
-                        outputs=output_text,
-                        output_token_ids=output_token_ids,
-                        finish_reason=choice.finish_reason,
-                        is_streaming=False,
-                        delta=False,
-                    )
-
-        return response
-
-    def _create_text_choice(
-        self,
-        request: ChatCompletionRequest,
-        omni_outputs: OmniRequestOutput,
-        tokenizer: AnyTokenizer,
-        conversation: list[ConversationMessage],
-        role: str,
-    ):
-        final_res = omni_outputs.request_output
-        if self.tool_call_id_type == "kimi_k2":
-            history_tool_call_cnt = get_history_tool_calls_cnt(conversation)
-        else:
-            history_tool_call_cnt = 0
-
-        choices: list[ChatCompletionResponseChoice] = []
-
-        for output in final_res.outputs:
-            token_ids = output.token_ids
-            out_logprobs = output.logprobs
-            tool_call_info = None
-
-            if request.logprobs and request.top_logprobs is not None:
-                assert out_logprobs is not None, "Did not output logprobs"
-                logprobs = self._create_chat_logprobs(
-                    token_ids=token_ids,
-                    top_logprobs=out_logprobs,
-                    num_output_top_logprobs=request.top_logprobs,
-                    tokenizer=tokenizer,
-                    return_as_token_id=request.return_tokens_as_token_ids,
-                )
-            else:
-                logprobs = None
-
-            if self.use_harmony:
-                reasoning_content, content, _ = parse_chat_output(token_ids)
-                if not request.include_reasoning:
-                    reasoning_content = None
-
-                if self.tool_parser is not None:
-                    tool_parser = self.tool_parser(tokenizer)
-                    # NOTE: We use token_ids for openai tool parser
-                    tool_call_info = tool_parser.extract_tool_calls(
-                        "",
-                        request=request,
-                        token_ids=token_ids,  # type: ignore
-                    )
-                    content = tool_call_info.content
-                    message = ChatMessage(
-                        role=role,
-                        reasoning_content=reasoning_content,
-                        content=content,
-                        tool_calls=tool_call_info.tool_calls,
-                    )
-                else:
-                    message = ChatMessage(
-                        role=role,
-                        reasoning_content=reasoning_content,
-                        content=content,
-                    )
-
-                choice_data = ChatCompletionResponseChoice(
-                    index=output.index,
-                    message=message,
-                    logprobs=logprobs,
-                    finish_reason=(
-                        "tool_calls"
-                        if (tool_call_info is not None and tool_call_info.tools_called)
-                        else (output.finish_reason if output.finish_reason else "stop")
-                    ),
-                    stop_reason=output.stop_reason,
-                )
-                choices.append(choice_data)
-                continue
-
-            if self.reasoning_parser:
-                try:
-                    reasoning_parser = self.reasoning_parser(tokenizer)
-                except RuntimeError as e:
-                    logger.exception("Error in reasoning parser creation.")
-                    return self.create_error_response(str(e))
-                # If the reasoning parser is enabled,
-                # tool calls are extracted exclusively from the content.
-                reasoning_content, content = reasoning_parser.extract_reasoning_content(output.text, request=request)
-                if not request.include_reasoning:
-                    reasoning_content = None
-            else:
-                reasoning_content = None
-                content = output.text
-
-            auto_tools_called = False
-            # if auto tools are not enabled, and a named tool choice using
-            #   outlines is not being used
-            if (not self.enable_auto_tools or not self.tool_parser) and (
-                not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
-                and request.tool_choice != "required"
-            ):
-                message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content)
-
-            # if the request uses tools and specified a tool choice
-            elif request.tool_choice and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam:
-                tool_call_class = MistralToolCall if isinstance(tokenizer, MistralTokenizer) else ToolCall
-                message = ChatMessage(
-                    role=role,
-                    reasoning_content=reasoning_content,
-                    content="",
-                    tool_calls=[
-                        tool_call_class(
-                            function=FunctionCall(
-                                name=request.tool_choice.function.name,
-                                arguments=content,
-                            )
-                        )
-                    ],
-                )
-
-            elif request.tool_choice and request.tool_choice == "required":
-                tool_call_class = MistralToolCall if isinstance(tokenizer, MistralTokenizer) else ToolCall
-
-                # the fields of FunctionDefinition are a superset of the
-                # tool call outputs and can be used for parsing
-                assert content is not None
-                tool_calls = TypeAdapter(list[FunctionDefinition]).validate_json(content)
-                tool_call_ids = []
-                for tool_call in tool_calls:
-                    tool_call_ids.append(
-                        make_tool_call_id(
-                            id_type=self.tool_call_id_type,
-                            func_name=tool_call.name,
-                            idx=history_tool_call_cnt,
-                        )
-                    )
-                    history_tool_call_cnt += 1
-                message = ChatMessage(
-                    role=role,
-                    content="",
-                    tool_calls=[
-                        tool_call_class(
-                            id=tool_call_ids[i],
-                            function=FunctionCall(
-                                name=tool_call.name,
-                                arguments=json.dumps(tool_call.parameters, ensure_ascii=False),
-                            ),
-                        )
-                        for i, tool_call in enumerate(tool_calls)
-                    ],
-                    reasoning_content=reasoning_content,
-                )
-
-            # if the request doesn't use tool choice
-            # OR specifies to not use a tool
-            elif not request.tool_choice or request.tool_choice == "none":
-                message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content)
-
-            # handle when there are tools and tool choice is auto
-            elif (
-                request.tools
-                and (request.tool_choice == "auto" or request.tool_choice is None)
-                and self.enable_auto_tools
-                and self.tool_parser
-            ):
-                try:
-                    tool_parser = self.tool_parser(tokenizer)
-                except RuntimeError as e:
-                    logger.exception("Error in tool parser creation.")
-                    return self.create_error_response(str(e))
-
-                tool_call_info = tool_parser.extract_tool_calls(content if content is not None else "", request=request)
-                # In the OpenAI API the finish_reason is "tools_called"
-                # if the tool choice is auto and the model produced a tool
-                # call. The same is not true for named function calls
-                auto_tools_called = tool_call_info.tools_called
-                if tool_call_info.tools_called:
-                    message = ChatMessage(
-                        role=role,
-                        reasoning_content=reasoning_content,
-                        content=tool_call_info.content,
-                        tool_calls=tool_call_info.tool_calls,
-                    )
-
-                else:
-                    # FOR NOW make it a chat message; we will have to detect
-                    # the type to make it later.
-                    ret_content = content
-
-                    # try to use content return from tool parser first,
-                    # tool parser may do some modify for the content.
-                    if tool_call_info.content and len(tool_call_info.content) > 0:
-                        ret_content = tool_call_info.content
-                    message = ChatMessage(
-                        role=role,
-                        reasoning_content=reasoning_content,
-                        content=ret_content,
-                    )
-
-            # undetermined case that is still important to handle
-            else:
-                logger.error(
-                    "Error in chat_completion_full_generator - cannot determine if tools should be extracted. "
-                    "Returning a standard chat completion."
-                )
-                message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content)
-
-            choice_data = ChatCompletionResponseChoice(
-                index=output.index,
-                message=message,
-                logprobs=logprobs,
-                finish_reason=(
-                    "tool_calls" if auto_tools_called else output.finish_reason if output.finish_reason else "stop"
-                ),
-                stop_reason=output.stop_reason,
-                token_ids=(as_list(output.token_ids) if request.return_token_ids else None),
-            )
-            choices.append(choice_data)
-
-        if request.echo:
-            last_msg_content: Union[str, list[dict[str, str]]] = ""
-            if conversation and "content" in conversation[-1] and conversation[-1].get("role") == role:
-                last_msg_content = conversation[-1]["content"] or ""
-            if isinstance(last_msg_content, list):
-                last_msg_content = "\n".join(msg["text"] for msg in last_msg_content)
-
-            for choice in choices:
-                full_message = last_msg_content + (choice.message.content or "")
-                choice.message.content = full_message
-
-        assert final_res.prompt_token_ids is not None
-        num_prompt_tokens = len(final_res.prompt_token_ids)
-        if final_res.encoder_prompt_token_ids is not None:
-            num_prompt_tokens += len(final_res.encoder_prompt_token_ids)
-        num_generated_tokens = sum(len(output.token_ids) for output in final_res.outputs)
-        usage = UsageInfo(
-            prompt_tokens=num_prompt_tokens,
-            completion_tokens=num_generated_tokens,
-            total_tokens=num_prompt_tokens + num_generated_tokens,
-        )
-        if self.enable_prompt_tokens_details and final_res.num_cached_tokens:
-            usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=final_res.num_cached_tokens)
-
-        prompt_logprobs = clamp_prompt_logprobs(final_res.prompt_logprobs)
-        prompt_token_ids = final_res.prompt_token_ids if request.return_token_ids else None
-        kv_transfer_params = final_res.kv_transfer_params
-
-        return choices, usage, prompt_logprobs, prompt_token_ids, kv_transfer_params
-
-    def _create_audio_choice(self, omni_outputs: OmniRequestOutput, role: str):
-        choices: list[ChatCompletionResponseChoice] = []
-        final_res = omni_outputs.request_output
-        audio_tensor = final_res.multimodal_output["audio"].detach().cpu().numpy()
-
-        # Convert numpy array to WAV bytes and encode as base64
-        if soundfile is None:
-            raise ImportError(
-                "soundfile is required for audio generation. Please install it with: pip install soundfile"
-            )
-
-        # Default sample rate for TTS models (typically 24000 Hz)
-        # You may need to adjust this based on your model's configuration
-        sample_rate = 24000
-
-        # Ensure audio is 1D (flatten if needed)
-        if audio_tensor.ndim > 1:
-            audio_tensor = audio_tensor.flatten()
-
-        # Convert to WAV format and encode as base64
-        with BytesIO() as buffer:
-            soundfile.write(buffer, audio_tensor, sample_rate, format="WAV")
-            wav_bytes = buffer.getvalue()
-
-        audio_base64 = base64.b64encode(wav_bytes).decode("utf-8")
-
-        # Generate unique ID for the audio
-        audio_id = f"audio-{uuid.uuid4().hex[:16]}"
-
-        # Set expiration time (e.g., 24 hours from now) as Unix timestamp
-        expires_at = int((datetime.now(timezone.utc) + timedelta(hours=24)).timestamp())
-
-        # Create OpenAIChatCompletionAudio object with all required fields
-        audio_obj = OpenAIChatCompletionAudio(
-            id=audio_id,
-            data=audio_base64,
-            expires_at=expires_at,
-            transcript="",  # Empty transcript if not available
-        )
-
-        for output in final_res.outputs:
-            choice_data = ChatCompletionResponseChoice(
-                index=output.index,
-                message=ChatMessage(role=role, audio=audio_obj),
-                logprobs=None,
-                finish_reason="stop",
-                stop_reason=None,
-            )
-            choices.append(choice_data)
-        return choices

vllm_omni/entrypoints/stage_utils.py

@@ -1,19 +1,19 @@
 from __future__ import annotations
 
-import json
+from typing import Any, Dict, Optional, Tuple, Union
+
 import logging
+import json
 import os
 import pickle
 from multiprocessing import shared_memory as _shm
-from typing import Any
 
 import cloudpickle
-from omegaconf import OmegaConf
 
 logger = logging.getLogger(__name__)
 
 
-def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None:
+def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) -> None:
     """Configure per-stage CUDA visibility and current device.
 
     Behavior
@@ -26,8 +26,8 @@ def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None:
     - Otherwise: set CUDA_VISIBLE_DEVICES to the provided single device string.
     """
     try:
-        selected_physical: int | None = None
-        logical_idx: int | None = None
+        selected_physical: Optional[int] = None
+        logical_idx: Optional[int] = None
 
         if isinstance(devices, str) and "," in devices:
             os.environ["CUDA_VISIBLE_DEVICES"] = devices
@@ -37,9 +37,7 @@ def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None:
                     selected_physical = int(toks[0])
                     logger.debug(
                         "[Stage-%s] Set CUDA_VISIBLE_DEVICES to %s; logical 0 -> physical %s",
-                        stage_id,
-                        devices,
-                        selected_physical,
+                        stage_id, devices, selected_physical,
                     )
                 except Exception as e:
                     logger.debug("[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e)
@@ -53,70 +51,43 @@ def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None:
                     if 0 <= logical_idx < len(mapping):
                         selected_physical = mapping[logical_idx]
                 except Exception as e:
-                    logger.debug(
-                        "[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s",
-                        stage_id,
-                        e,
-                    )
+                    logger.debug("[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", stage_id, e)
                     selected_physical = None
             if selected_physical is None:
                 selected_physical = int(logical_idx)
             os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
             logger.debug(
                 "[Stage-%s] Logical index %d -> physical %s; set CUDA_VISIBLE_DEVICES to single device",
-                stage_id,
-                logical_idx + 1,
-                selected_physical,
+                stage_id, logical_idx + 1, selected_physical,
             )
         elif devices in (None, "cpu"):
-            logger.debug(
-                "[Stage-%s] Using default device visibility (devices=%s)",
-                stage_id,
-                devices,
-            )
+            logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices)
         else:
             selected_physical = int(str(devices))
             os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
-            logger.debug(
-                "[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)",
-                stage_id,
-                selected_physical,
-            )
+            logger.debug("[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", stage_id, selected_physical)
 
         try:
             import torch  # noqa: WPS433
-
             if torch.cuda.is_available():
                 try:
                     torch.cuda.set_device(0)
                 except Exception as e:
-                    logger.debug(
-                        "[Stage-%s] torch.cuda.set_device(0) failed: %s",
-                        stage_id,
-                        e,
-                        exc_info=True,
-                    )
+                    logger.debug("[Stage-%s] torch.cuda.set_device(0) failed: %s", stage_id, e, exc_info=True)
                 num = torch.cuda.device_count()
                 info = []
                 for i in range(num):
                     total = torch.cuda.get_device_properties(i).total_memory
                     free, _ = torch.cuda.mem_get_info(i)
-                    info.append(
-                        {
-                            "idx": i,
-                            "name": torch.cuda.get_device_name(i),
-                            "total": int(total),
-                            "free": int(free),
-                        }
-                    )
+                    info.append({
+                        "idx": i,
+                        "name": torch.cuda.get_device_name(i),
+                        "total": int(total),
+                        "free": int(free),
+                    })
                 logger.debug("[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info)
         except Exception as e:
-            logger.debug(
-                "[Stage-%s] Failed to query CUDA devices: %s",
-                stage_id,
-                e,
-                exc_info=True,
-            )
+            logger.debug("[Stage-%s] Failed to query CUDA devices: %s", stage_id, e, exc_info=True)
     except Exception as e:
         logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e)
 
@@ -126,7 +97,7 @@ def serialize_obj(obj: Any) -> bytes:
     return cloudpickle.dumps(obj)
 
 
-def shm_write_bytes(payload: bytes) -> dict[str, Any]:
+def shm_write_bytes(payload: bytes) -> Dict[str, Any]:
     """Write bytes into SharedMemory and return meta dict {name,size}.
 
     Caller should close the segment; the receiver should unlink.
@@ -139,11 +110,11 @@ def shm_write_bytes(payload: bytes) -> dict[str, Any]:
     try:
         shm.close()
     except Exception as e:
-        logger.debug("Failed to close shared memory: %s", e)
+        logger.debug("Ensure parent dir failed for %s: %s", path, e)
     return meta
 
 
-def shm_read_bytes(meta: dict[str, Any]) -> bytes:
+def shm_read_bytes(meta: Dict[str, Any]) -> bytes:
     """Read bytes from SharedMemory by meta {name,size} and cleanup."""
     shm = _shm.SharedMemory(name=meta["name"])  # type: ignore[index]
     mv = memoryview(shm.buf)
@@ -170,7 +141,7 @@ def _ensure_parent_dir(path: str) -> None:
         pass
 
 
-def append_jsonl(path: str, record: dict[str, Any]) -> None:
+def append_jsonl(path: str, record: Dict[str, Any]) -> None:
     """Append a JSON record as one line to a JSONL file (best-effort).
 
     This is safe to call from multiple processes when each process writes
@@ -187,7 +158,7 @@ def append_jsonl(path: str, record: dict[str, Any]) -> None:
         logger.exception("Failed to append JSONL to %s", path)
 
 
-def maybe_dump_to_shm(obj: Any, threshold: int) -> tuple[bool, Any]:
+def maybe_dump_to_shm(obj: Any, threshold: int) -> Tuple[bool, Any]:
     """Dump object to SHM if serialized size exceeds threshold.
 
     Returns (True, meta) when dumped; otherwise (False, original_obj).
@@ -198,7 +169,7 @@ def maybe_dump_to_shm(obj: Any, threshold: int) -> tuple[bool, Any]:
     return False, obj
 
 
-def maybe_load_from_ipc(container: dict[str, Any], obj_key: str, shm_key: str) -> Any:
+def maybe_load_from_ipc(container: Dict[str, Any], obj_key: str, shm_key: str) -> Any:
     """Load object from container that may carry SHM or inline object.
 
     Deprecated: prefer `maybe_load_from_ipc_with_metrics` to also obtain
@@ -209,9 +180,7 @@ def maybe_load_from_ipc(container: dict[str, Any], obj_key: str, shm_key: str) -
     return container[obj_key]
 
 
-def maybe_load_from_ipc_with_metrics(
-    container: dict[str, Any], obj_key: str, shm_key: str
-) -> tuple[Any, dict[str, float]]:
+def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, shm_key: str) -> tuple[Any, Dict[str, float]]:
     """Load object and return (object, metrics) with RX bytes and decode time.
 
     Metrics keys:
@@ -219,7 +188,6 @@ def maybe_load_from_ipc_with_metrics(
       - rx_decode_time_ms: float
     """
     import time as _time  # local import to avoid overhead at module import
-
     t0 = _time.time()
     if shm_key in container:
         meta = container[shm_key]  # type: ignore[index]
@@ -243,13 +211,13 @@ def maybe_load_from_ipc_with_metrics(
     }
 
 
-def encode_for_ipc(obj: Any, threshold: int, obj_key: str, shm_key: str) -> dict[str, Any]:
+def encode_for_ipc(obj: Any, threshold: int, obj_key: str, shm_key: str) -> Dict[str, Any]:
     """Return a dict payload for IPC: inline (obj_key) or SHM (shm_key).
 
     When serialized size exceeds threshold, returns {shm_key: {name,size}};
     otherwise returns {obj_key: obj}.
     """
-    payload: dict[str, Any] = {}
+    payload: Dict[str, Any] = {}
     use_shm, data = maybe_dump_to_shm(obj, threshold)
     if use_shm:
         payload[shm_key] = data
@@ -259,7 +227,7 @@ def encode_for_ipc(obj: Any, threshold: int, obj_key: str, shm_key: str) -> dict
 
 
 # Convert OmegaConf/objects to plain dicts
-def _to_dict(x: Any) -> dict[str, Any]:
+def _to_dict(x: Any) -> Dict[str, Any]:
     try:
         if isinstance(x, dict):
             return dict(x)

vllm_omni/entrypoints/utils.py

@@ -1,13 +1,9 @@
-from __future__ import annotations
-
-import logging
 import os
 from pathlib import Path
 
 from omegaconf import OmegaConf
-from vllm.transformers_utils.config import get_config
 
-logger = logging.getLogger(__name__)
+from vllm.transformers_utils.config import get_config
 
 # Get the project root directory (2 levels up from this file)
 PROJECT_ROOT = Path(__file__).parent.parent.parent

vllm_omni/inputs/data.py

@@ -1,12 +1,7 @@
-from typing import Any, Optional
-
-try:
-    from typing import NotRequired
-except ImportError:
-    # Python < 3.11: use typing_extensions
-    from typing_extensions import NotRequired
+from typing import Any, NotRequired, Optional
 
 import torch
+
 from vllm.inputs.data import EmbedsPrompt, TokenInputs, TokensPrompt
 
 
@@ -44,7 +39,8 @@ def token_inputs_omni(
     prompt_embeds: Optional[torch.Tensor] = None,
     additional_information: Optional[dict[str, Any]] = None,
 ) -> OmniTokenInputs:
-    """Construct token inputs with optional embeddings and metadata."""
+    """Construct [`TokenInputs`][vllm.inputs.data.TokenInputs] from optional
+    values."""
     inputs = OmniTokenInputs(type="token", prompt_token_ids=prompt_token_ids)
 
     if prompt is not None:

vllm_omni/inputs/parse.py

@@ -17,9 +17,15 @@ def parse_singleton_prompt_omni(prompt: SingletonPrompt) -> ParsedSingletonPromp
         # Priority tokens: When both tokens and embeds exist, keep both and
         # follow the tokens path
         if "prompt_token_ids" in prompt:
-            return ParsedTokensPrompt(type="tokens", content=prompt)  # type: ignore[typeddict-item]
+            return ParsedTokensPrompt(
+                type="tokens", content=prompt
+            )  # type: ignore[typeddict-item]
         elif "prompt_embeds" in prompt:
-            return ParsedEmbedsPrompt(type="embeds", content=prompt)  # type: ignore[typeddict-item]
+            return ParsedEmbedsPrompt(
+                type="embeds", content=prompt
+            )  # type: ignore[typeddict-item]
         elif "prompt" in prompt:
             return ParsedTextPrompt(type="text", content=prompt)
-    raise TypeError("inputs must be a string, TextPrompt, TokensPrompt, or EmbedsPrompt")
+    raise TypeError(
+        "inputs must be a string, TextPrompt, TokensPrompt, or EmbedsPrompt"
+    )

vllm_omni/inputs/preprocess.py

@@ -1,16 +1,13 @@
 from typing import Any, Optional, Union
 
 from typing_extensions import assert_never
+
 from vllm.inputs.data import SingletonInputs, SingletonPrompt, TextPrompt, TokensPrompt
 from vllm.inputs.preprocess import InputPreprocessor
-from vllm.logger import init_logger
 from vllm.multimodal.inputs import MultiModalInputs, MultiModalUUIDDict
-
 from vllm_omni.inputs.data import OmniTokenInputs, token_inputs_omni
 from vllm_omni.inputs.parse import parse_singleton_prompt_omni
 
-logger = init_logger(__name__)
-
 
 class OmniInputPreprocessor(InputPreprocessor):
     def _process_tokens(
@@ -20,7 +17,9 @@ class OmniInputPreprocessor(InputPreprocessor):
         *,
         mm_uuids: Optional[MultiModalUUIDDict] = None,
     ) -> Union[OmniTokenInputs, MultiModalInputs]:
-        prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs)
+        prompt_token_ids = self._truncate_inputs(
+            parsed_content["prompt_token_ids"], tokenization_kwargs
+        )
         prompt_embeds = parsed_content.get("prompt_embeds")
         additional_information = parsed_content.get("additional_information")
 
@@ -63,7 +62,7 @@ class OmniInputPreprocessor(InputPreprocessor):
 
         Returns:
 
-        * Input container compatible with vLLM's singleton prompt handling.
+        * [`SingletonInputs`][vllm.inputs.data.SingletonInputs] instance
         """
         parsed = parse_singleton_prompt_omni(prompt)
 

vllm_omni/model_executor/layers/mrope.py

@@ -4,6 +4,7 @@ from typing import Optional, Union
 import numpy as np
 import torch
 from transformers import PretrainedConfig
+
 from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding
 
 
@@ -248,7 +249,9 @@ class MRotaryEmbedding(RotaryEmbedding):
                     input_token_type.append("text")
 
             input_type_group: list[tuple[str, int, int]] = []
-            for key, group_iter in itertools.groupby(enumerate(input_token_type), lambda x: x[1]):
+            for key, group_iter in itertools.groupby(
+                enumerate(input_token_type), lambda x: x[1]
+            ):
                 group_list = list(group_iter)
                 start_index = group_list[0][0]
                 end_index = group_list[-1][0] + 1
@@ -257,7 +260,9 @@ class MRotaryEmbedding(RotaryEmbedding):
             video_frame_num = 1
             mm_data_idx = 0
             for modality_type, start_idx, end_idx in input_type_group:
-                st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                st_idx = (
+                    llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                )
                 if modality_type == "image":
                     t, h, w = (
                         image_grid_thw[mm_data_idx][0],
@@ -270,10 +275,27 @@ class MRotaryEmbedding(RotaryEmbedding):
                         w // spatial_merge_size,
                     )
 
-                    t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
-                    h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
-                    w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
-                    llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx)
+                    t_index = (
+                        torch.arange(llm_grid_t)
+                        .view(-1, 1)
+                        .expand(-1, llm_grid_h * llm_grid_w)
+                        .flatten()
+                    )
+                    h_index = (
+                        torch.arange(llm_grid_h)
+                        .view(1, -1, 1)
+                        .expand(llm_grid_t, -1, llm_grid_w)
+                        .flatten()
+                    )
+                    w_index = (
+                        torch.arange(llm_grid_w)
+                        .view(1, 1, -1)
+                        .expand(llm_grid_t, llm_grid_h, -1)
+                        .flatten()
+                    )
+                    llm_pos_ids_list.append(
+                        torch.stack([t_index, h_index, w_index]) + st_idx
+                    )
                     mm_data_idx += 1
 
                 elif modality_type == "video":
@@ -289,17 +311,36 @@ class MRotaryEmbedding(RotaryEmbedding):
                     )
 
                     for t_idx in range(llm_grid_t):
-                        t_index = torch.tensor(t_idx).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
-                        h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(1, -1, llm_grid_w).flatten()
-                        w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(1, llm_grid_h, -1).flatten()
-                        llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx)
+                        t_index = (
+                            torch.tensor(t_idx)
+                            .view(-1, 1)
+                            .expand(-1, llm_grid_h * llm_grid_w)
+                            .flatten()
+                        )
+                        h_index = (
+                            torch.arange(llm_grid_h)
+                            .view(1, -1, 1)
+                            .expand(1, -1, llm_grid_w)
+                            .flatten()
+                        )
+                        w_index = (
+                            torch.arange(llm_grid_w)
+                            .view(1, 1, -1)
+                            .expand(1, llm_grid_h, -1)
+                            .flatten()
+                        )
+                        llm_pos_ids_list.append(
+                            torch.stack([t_index, h_index, w_index]) + st_idx
+                        )
 
                     mm_data_idx += 1
                     video_frame_num += 1
 
                 else:
                     text_len = end_idx - start_idx
-                    llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+                    llm_pos_ids_list.append(
+                        torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
+                    )
                     video_frame_num = 1
 
         else:
@@ -331,7 +372,9 @@ class MRotaryEmbedding(RotaryEmbedding):
         tokens_per_second = getattr(hf_config.vision_config, "tokens_per_second", 1.0)
 
         input_tokens_tensor = torch.tensor(input_tokens)
-        vision_start_indices = torch.argwhere(input_tokens_tensor == vision_start_token_id).squeeze(1)
+        vision_start_indices = torch.argwhere(
+            input_tokens_tensor == vision_start_token_id
+        ).squeeze(1)
         vision_tokens = input_tokens_tensor[vision_start_indices + 1]
         image_nums = (vision_tokens == image_token_id).sum()
         video_nums = (vision_tokens == video_token_id).sum()
@@ -381,11 +424,15 @@ class MRotaryEmbedding(RotaryEmbedding):
             text_len = ed - st
 
             st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
-            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+            llm_pos_ids_list.append(
+                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
+            )
 
             t_index = (
                 (
-                    torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w)
+                    torch.arange(llm_grid_t)
+                    .view(-1, 1)
+                    .expand(-1, llm_grid_h * llm_grid_w)
                     * video_second_per_grid_t
                     * tokens_per_second
                 )
@@ -393,15 +440,29 @@ class MRotaryEmbedding(RotaryEmbedding):
                 .flatten()
             )
 
-            h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
-            w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
-            llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+            h_index = (
+                torch.arange(llm_grid_h)
+                .view(1, -1, 1)
+                .expand(llm_grid_t, -1, llm_grid_w)
+                .flatten()
+            )
+            w_index = (
+                torch.arange(llm_grid_w)
+                .view(1, 1, -1)
+                .expand(llm_grid_t, llm_grid_h, -1)
+                .flatten()
+            )
+            llm_pos_ids_list.append(
+                torch.stack([t_index, h_index, w_index]) + text_len + st_idx
+            )
             st = ed + llm_grid_t * llm_grid_h * llm_grid_w
 
         if st < len(input_tokens):
             st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
             text_len = len(input_tokens) - st
-            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+            llm_pos_ids_list.append(
+                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
+            )
 
         llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
         mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()
@@ -451,7 +512,9 @@ class MRotaryEmbedding(RotaryEmbedding):
         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)
+        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)
@@ -471,13 +534,21 @@ class MRotaryEmbedding(RotaryEmbedding):
         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
+            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:
+                    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:
+                    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])
@@ -500,40 +571,58 @@ class MRotaryEmbedding(RotaryEmbedding):
                     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)
+                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()
+                t_index = (
+                    torch.arange(grid_t)
+                    * second_per_grid_ts[video_idx]
+                    * tokens_per_second
+                ).long()
                 llm_pos_ids = cls._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)
+                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)
+                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 = cls._split_list_into_ranges(t_index, t_ntoken_per_chunk)
+                t_index = (
+                    torch.arange(grid_t)
+                    * second_per_grid_ts[video_idx]
+                    * tokens_per_second
+                ).long()
+                t_index_split_chunk = cls._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)
+                    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 = cls._get_llm_pos_ids_for_vision(
                         start_idx,
@@ -544,23 +633,38 @@ class MRotaryEmbedding(RotaryEmbedding):
                         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])
+                    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
+                        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)
+                            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)
+                    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])
+                    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
+                        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
@@ -569,7 +673,9 @@ class MRotaryEmbedding(RotaryEmbedding):
             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)
+        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
@@ -586,10 +692,25 @@ class MRotaryEmbedding(RotaryEmbedding):
         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()
+        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()
+            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)
@@ -610,7 +731,14 @@ class MRotaryEmbedding(RotaryEmbedding):
         context_len: int,
         seq_len: int,
     ) -> list[list[int]]:
-        return [list(range(context_len + mrope_position_delta, seq_len + mrope_position_delta)) for _ in range(3)]
+        return [
+            list(
+                range(
+                    context_len + mrope_position_delta, seq_len + mrope_position_delta
+                )
+            )
+            for _ in range(3)
+        ]
 
     @staticmethod
     def get_next_input_positions_tensor(
@@ -620,6 +748,7 @@ class MRotaryEmbedding(RotaryEmbedding):
         context_len: int,
         num_new_tokens: int,
     ):
+
         values = np.arange(
             mrope_position_delta + context_len,
             mrope_position_delta + context_len + num_new_tokens,
@@ -650,19 +779,25 @@ class MRotaryEmbedding(RotaryEmbedding):
         audio_end_token_id = thinker_config.audio_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)
+        tokens_per_second = getattr(
+            thinker_config.vision_config, "tokens_per_second", 25
+        )
 
         grid_t = video_grid_thw[0]
         grid_h = video_grid_thw[1]
         grid_w = video_grid_thw[2]
         t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
-        t_index = (torch.arange(grid_t) * video_second_per_grid_t * tokens_per_second).long()
+        t_index = (
+            torch.arange(grid_t) * video_second_per_grid_t * tokens_per_second
+        ).long()
         t_index_split_chunk = cls._split_list_into_ranges(t_index, t_ntoken_per_chunk)
 
         updates = [audio_start_token_id]
         added_audio_len = 0
         for t_chunk in t_index_split_chunk:
-            vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
+            vision_ntoken_per_chunk = (
+                len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
+            )
             updates.extend([video_token_id] * vision_ntoken_per_chunk)
 
             audio_chunk_size = min(t_ntoken_per_chunk, audio_len - added_audio_len)

vllm_omni/model_executor/model_loader/weight_utils.py

@@ -4,6 +4,7 @@ from typing import Optional, Union
 
 import huggingface_hub
 from huggingface_hub import snapshot_download
+
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader.weight_utils import DisabledTqdm, get_lock
 

vllm_omni/model_executor/models/qwen2_5_omni.py

@@ -1,21 +1,26 @@
 import glob
 import os
-from collections.abc import Iterable
 from functools import cached_property
-from typing import NamedTuple, Optional, Union
+from typing import Dict, Iterable, NamedTuple, Optional, Set, Tuple, 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,
     Qwen2_5OmniThinkerConfig,
 )
+
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.qwen2_5_omni_thinker import (
+    Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniThinkerDummyInputsBuilder,
+    Qwen2_5OmniThinkerMultiModalProcessor,
+    Qwen2_5OmniThinkerProcessingInfo,
+)
 from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
 
 # from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav
@@ -24,16 +29,10 @@ from vllm.sequence import IntermediateTensors
 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.model_loader.weight_utils import download_weights_from_hf_specific
-from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
-    Qwen2_5OmniConditionalGenerationMixin,
-    Qwen2_5OmniThinkerDummyInputsBuilder,
-    Qwen2_5OmniThinkerMultiModalProcessor,
-    Qwen2_5OmniThinkerProcessingInfo,
+from vllm_omni.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf_specific,
 )
-from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
-from vllm_omni.model_executor.models.vision import get_llm_pos_ids_for_vision
+from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
 
 TALKER_CODEC_EOS_TOKEN_ID = 8294
 TALKER_CODEC_BOS_TOKEN_ID = 8293
@@ -56,8 +55,9 @@ logger = init_logger(__name__)
     dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
 class Qwen2_5OmniForConditionalGeneration(
-    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin, SupportsMRoPE
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
 ):
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.have_multimodal_outputs = True
@@ -117,15 +117,17 @@ class Qwen2_5OmniForConditionalGeneration(
                     architectures=["Qwen2_5OmniToken2WavModel"],
                 )
             # voice resources (loaded on demand)
-            self._token2wav_conds: dict[str, torch.Tensor] = {}
-            self._token2wav_ref_mels: dict[str, torch.Tensor] = {}
+            self._token2wav_conds: Dict[str, torch.Tensor] = {}
+            self._token2wav_ref_mels: Dict[str, torch.Tensor] = {}
             self.model = self.token2wav
         else:
             raise ValueError("Invalid model stage")
 
         # Set up intermediate tensors
         self.make_empty_intermediate_tensors = (
-            (self.thinker.make_empty_intermediate_tensors) if self.model_stage == "thinker" else lambda: None
+            (self.thinker.make_empty_intermediate_tensors)
+            if self.model_stage == "thinker"
+            else lambda: None
         )
 
     # -------------------- Device utilities --------------------
@@ -174,7 +176,11 @@ class Qwen2_5OmniForConditionalGeneration(
         multimodal_embeddings=None,
     ) -> torch.Tensor:
         if self.model_stage == "code2wav":
-            return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.vllm_config.model_config.get_hidden_size())
+            return (
+                torch.zeros_like(input_ids)
+                .reshape(-1, 1)
+                .repeat(1, self.vllm_config.model_config.get_hidden_size())
+            )
         return self.model.get_input_embeddings(input_ids, multimodal_embeddings)
 
     def get_multimodal_embeddings(self, **kwargs):
@@ -223,7 +229,9 @@ class Qwen2_5OmniForConditionalGeneration(
             # if input_ids is None, set it to a zero tensor, in the length of the
             # same as the embedding seq length
             if input_ids is None:
-                input_ids = torch.zeros(inputs_embeds.shape[1], dtype=torch.long, device=thinker_dev).unsqueeze(
+                input_ids = torch.zeros(
+                    inputs_embeds.shape[1], dtype=torch.long, device=thinker_dev
+                ).unsqueeze(
                     0
                 )  # (1, 0)
                 added_batch_dim = True
@@ -251,7 +259,11 @@ class Qwen2_5OmniForConditionalGeneration(
 
             # Text-only path
             return OmniOutput(
-                text_hidden_states=(text_hidden_states.squeeze(0) if added_batch_dim else text_hidden_states),
+                text_hidden_states=(
+                    text_hidden_states.squeeze(0)
+                    if added_batch_dim
+                    else text_hidden_states
+                ),
                 multimodal_outputs=None,
             )
 
@@ -261,17 +273,13 @@ class Qwen2_5OmniForConditionalGeneration(
             # Rules:
             # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS]
             # - Decode segments consist of a single non-special token.
-            # - If additional_information is provided
-            # (can be a list split by request or a concatenated tensor plus a list of shapes),
+            # - If additional_information is provided (can be a list split by request or a concatenated tensor plus a list of shapes),
             #   then for each request, reconstruct the thinker→talker input embeddings for the Prefill segments;
-            # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them;
-            # otherwise, keep the original embedding.
+            # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them; otherwise, keep the original embedding.
 
             if input_ids is None and additional_information is None:
                 input_ids = torch.zeros(
-                    inputs_embeds.shape[0],
-                    dtype=torch.long,
-                    device=inputs_embeds.device,
+                    inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device
                 )
                 additional_information = {}
                 self.thinker_reply_part = torch.zeros_like(inputs_embeds)
@@ -285,17 +293,13 @@ class Qwen2_5OmniForConditionalGeneration(
 
             # ------- Request-scoped additional information (no cross-request concat) -------
             request_ids: Optional[list[str]] = kwargs.get("request_ids")  # ordered
-            request_token_spans: Optional[list[tuple[int, int]]] = kwargs.get("request_token_spans")
+            request_token_spans: Optional[list[tuple[int,int]]] = kwargs.get("request_token_spans")
             addi_by_req: Optional[dict] = kwargs.get("additional_information_by_req_id")
             runtime_addi = kwargs.get("runtime_additional_information")
 
             # Normalize runtime_addi into a mapping by request_id for convenience
             runtime_addi_by_req: dict[str, dict] = {}
-            if (
-                isinstance(request_ids, list)
-                and isinstance(runtime_addi, list)
-                and len(runtime_addi) == len(request_ids)
-            ):
+            if isinstance(request_ids, list) and isinstance(runtime_addi, list) and len(runtime_addi) == len(request_ids):
                 for i, rid in enumerate(request_ids):
                     if isinstance(rid, str) and isinstance(runtime_addi[i], dict):
                         runtime_addi_by_req[rid] = runtime_addi[i]
@@ -328,19 +332,9 @@ class Qwen2_5OmniForConditionalGeneration(
                     otoks = info.get("thinker_output_token_ids")  # list[int]
 
                     if not isinstance(pe, torch.Tensor):
-                        pe = torch.zeros(
-                            0,
-                            self.talker.config.hidden_size,
-                            dtype=inputs_embeds.dtype,
-                            device=self._module_device(self.model),
-                        )
+                        pe = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
                     if not isinstance(tr, torch.Tensor):
-                        tr = torch.zeros(
-                            0,
-                            self.talker.config.hidden_size,
-                            dtype=inputs_embeds.dtype,
-                            device=self._module_device(self.model),
-                        )
+                        tr = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
                     if not isinstance(ptoks, (list, torch.Tensor)):
                         ptoks = []
                     if not isinstance(otoks, (list, torch.Tensor)):
@@ -360,9 +354,7 @@ class Qwen2_5OmniForConditionalGeneration(
 
                     # Prepare per-request reply queue for subsequent decode: drop first row
                     if tr.ndim == 2 and tr.shape[0] > 0:
-                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = (
-                            tr[1:].detach().to("cpu").contiguous()
-                        )
+                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = tr[1:].detach().to("cpu").contiguous()
 
             # ------- Decode: span_len == 1 -------
             if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list):
@@ -386,11 +378,7 @@ class Qwen2_5OmniForConditionalGeneration(
                         dv = info.get("decode_output_prompt_embeds") if isinstance(info, dict) else None
                         if isinstance(dv, torch.Tensor) and dv.numel() > 0:
                             step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1)
-                        elif (
-                            hasattr(self, "thinker_reply_part")
-                            and isinstance(self.thinker_reply_part, torch.Tensor)
-                            and self.thinker_reply_part.numel() > 0
-                        ):
+                        elif hasattr(self, "thinker_reply_part") and isinstance(self.thinker_reply_part, torch.Tensor) and self.thinker_reply_part.numel() > 0:
                             # C) fallback shared pool
                             step_vec = self.thinker_reply_part[0:1]
                             self.thinker_reply_part = self.thinker_reply_part[1:]
@@ -415,8 +403,7 @@ class Qwen2_5OmniForConditionalGeneration(
                 multimodal_outputs = {"additional_information_update_by_req_id": update_by_req_id}
 
             if sampling_metadata is not None:
-                # the padding token id is set to text model's pad token id,
-                # which do not match with the talker model's word embedding size
+                # the padding token id is set to text model's pad token id, which do not match with the talker model's word embedding size
                 sampling_metadata.prompt_token_ids[sampling_metadata.prompt_token_ids == 152064] = 8448
 
             return OmniOutput(
@@ -434,12 +421,16 @@ class Qwen2_5OmniForConditionalGeneration(
                     device=inputs_embeds.device,
                 )
             )
-
+            
             code = code[:-1] if code[-1] == TALKER_CODEC_EOS_TOKEN_ID else code
             code = code[1:] if code[0] == TALKER_CODEC_BOS_TOKEN_ID else code
-
-            audio_tensor = self.generate_audio(code, voice_type)
-            return OmniOutput(text_hidden_states=None, multimodal_outputs={"audio": audio_tensor})
+            
+            audio_tensor = self.generate_audio(
+                code, voice_type
+            )
+            return OmniOutput(
+                text_hidden_states=None, multimodal_outputs={"audio": audio_tensor}
+            )
 
         return OmniOutput(
             text_hidden_states=torch.cat(
@@ -450,7 +441,9 @@ class Qwen2_5OmniForConditionalGeneration(
                     ).to(self._module_device(self.model)),
                     self.talker.thinker_to_talker_proj(
                         self.talker.get_input_embeddings(
-                            torch.tensor([TALKER_CODEC_BOS_TOKEN_ID, TALKER_CODEC_EOS_TOKEN_ID])
+                            torch.tensor(
+                                [TALKER_CODEC_BOS_TOKEN_ID, TALKER_CODEC_EOS_TOKEN_ID]
+                            )
                             .to(torch.bfloat16)
                             .to(self._module_device(self.model))
                         )
@@ -461,170 +454,6 @@ class Qwen2_5OmniForConditionalGeneration(
             multimodal_outputs=None,
         )
 
-    def get_mrope_input_positions(
-        self,
-        input_tokens: list[int],
-        hf_config: PretrainedConfig,
-        image_grid_thw: list[list[int]] | torch.Tensor,
-        video_grid_thw: list[list[int]] | torch.Tensor,
-        second_per_grid_ts: list[float] | None = None,
-        context_len: int = 0,
-        seq_len: int | None = None,
-        audio_feature_lengths: torch.Tensor | None = None,
-        use_audio_in_video: bool = False,
-    ) -> tuple[torch.Tensor, int]:
-        """Get mrope input positions and delta value (Qwen2.5-Omni version).
-
-        Differences from MRotaryEmbedding:
-            1. Add audio support (and related `audio_feature_lengths`).
-            2. Add `use_audio_in_video` option to read audio from video inputs.
-                In this case, audio and vision position ids will be split into
-                chunks and interleaved.
-
-        Example:
-
-            (V_i are vision position ids, A_i are audio position ids)
-
-            |V_1 ...    V_n|A_1 ...   A_n|V_n+1 ... V_2n|A_n+1 ... A_2n|...
-            |vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |...
-        """
-
-        # TODO(fyabc): refactor and share more code with
-        #  _vl_get_input_positions_tensor.
-
-        thinker_config = hf_config.thinker_config
-        audio_token_id = thinker_config.audio_token_index
-        image_token_id = thinker_config.image_token_index
-        video_token_id = thinker_config.video_token_index
-        audio_start_token_id = thinker_config.audio_start_token_id
-        audio_end_token_id = thinker_config.audio_end_token_id
-        vision_start_token_id = thinker_config.vision_start_token_id
-        vision_end_token_id = thinker_config.vision_end_token_id
-        seconds_per_chunk = thinker_config.seconds_per_chunk
-        spatial_merge_size = thinker_config.vision_config.spatial_merge_size
-        tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25)
-
-        if isinstance(image_grid_thw, list):
-            image_grid_thw = torch.tensor(image_grid_thw)
-        if isinstance(video_grid_thw, list):
-            video_grid_thw = torch.tensor(video_grid_thw)
-
-        src_item = input_tokens
-        audio_seqlens = audio_feature_lengths
-        if not second_per_grid_ts:
-            second_per_grid_ts = [1] * video_grid_thw.shape[0]
-        audio_idx = 0
-        video_idx = 0
-        image_idx = 0
-        new_src_item: list[int] = []
-        llm_pos_ids_list: list[torch.Tensor] = []
-
-        idx = 0
-        while idx < len(src_item):
-            new_src_item_len = len(new_src_item)
-            start_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
-            if src_item[idx] not in [audio_token_id, video_token_id, image_token_id]:
-                if use_audio_in_video and idx > 0:
-                    if src_item[idx] == vision_end_token_id and src_item[idx - 1] == audio_end_token_id:
-                        # processing the <|audio_eos|> before <|vision_eos|>
-                        start_idx -= 1
-                    elif src_item[idx] == audio_start_token_id and src_item[idx - 1] == vision_start_token_id:
-                        # processing the <|audio_bos|> after <|vision_eos|>
-                        start_idx -= 1
-                new_src_item.append(src_item[idx])
-                llm_pos_ids = torch.tensor([start_idx], dtype=torch.long).expand(3, -1)
-                llm_pos_ids_list.append(llm_pos_ids)
-            elif src_item[idx] == audio_token_id:
-                assert audio_seqlens is not None
-                audio_seqlen = audio_seqlens[audio_idx]
-                place_num = ((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1
-                new_src_item.extend([audio_token_id] * place_num)
-                llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx
-                llm_pos_ids_list.append(llm_pos_ids)
-                audio_idx += 1
-            elif src_item[idx] == image_token_id:
-                grid_t = image_grid_thw[image_idx][0]
-                grid_hs = image_grid_thw[:, 1]
-                grid_ws = image_grid_thw[:, 2]
-                t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long()
-                llm_pos_ids = get_llm_pos_ids_for_vision(
-                    start_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
-                )
-                llm_pos_ids_list.append(llm_pos_ids)
-                vision_seqlen = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
-                new_src_item.extend([image_token_id] * vision_seqlen)
-                image_idx += 1
-            elif src_item[idx] == video_token_id and not use_audio_in_video:
-                grid_t = video_grid_thw[video_idx][0]
-                grid_hs = video_grid_thw[:, 1]
-                grid_ws = video_grid_thw[:, 2]
-                t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long()
-                llm_pos_ids = get_llm_pos_ids_for_vision(
-                    start_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
-                )
-                llm_pos_ids_list.append(llm_pos_ids)
-                vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
-                new_src_item.extend([video_token_id] * vision_seqlen)
-                video_idx += 1
-            else:
-                # read audio from video
-                assert audio_seqlens is not None
-                audio_seqlen = audio_seqlens[audio_idx]
-                vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
-                grid_t = video_grid_thw[video_idx][0]
-                grid_h = video_grid_thw[video_idx][1]
-                grid_w = video_grid_thw[video_idx][2]
-                grid_hs = video_grid_thw[:, 1]
-                grid_ws = video_grid_thw[:, 2]
-                t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
-                t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long()
-                t_index_split_chunk = split_list_into_ranges(t_index, t_ntoken_per_chunk)
-                place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2
-                pure_audio_len = place_num - 2
-                added_audio_len = 0
-                audio_llm_pos_ids_list: list[torch.Tensor] = []
-                for t_chunk in t_index_split_chunk:
-                    vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
-                    new_src_item.extend([video_token_id] * vision_ntoken_per_chunk)
-                    vision_llm_pos_ids_list = get_llm_pos_ids_for_vision(
-                        start_idx,
-                        video_idx,
-                        spatial_merge_size,
-                        t_chunk,
-                        grid_hs,
-                        grid_ws,
-                    ).split(1, dim=1)
-                    llm_pos_ids_list.extend(vision_llm_pos_ids_list)
-                    new_src_item.extend(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) * [audio_token_id])
-                    audio_start_idx = (
-                        start_idx if len(audio_llm_pos_ids_list) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1
-                    )
-                    if min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) > 0:
-                        audio_llm_pos_ids_list = (
-                            torch.arange(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)).expand(3, -1)
-                            + audio_start_idx
-                        ).split(1, dim=1)
-                    else:
-                        audio_llm_pos_ids_list = []
-                    added_audio_len += min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)
-                    llm_pos_ids_list.extend(audio_llm_pos_ids_list)
-                if added_audio_len < pure_audio_len:
-                    new_src_item.extend((pure_audio_len - added_audio_len) * [audio_token_id])
-                    audio_llm_pos_ids_list = (
-                        torch.arange(pure_audio_len - added_audio_len).expand(3, -1) + llm_pos_ids_list[-1].max() + 1
-                    ).split(1, dim=1)
-                    llm_pos_ids_list.extend(audio_llm_pos_ids_list)
-                audio_idx += 1
-                video_idx += 1
-            # move to the next token
-            idx += len(new_src_item) - new_src_item_len
-
-        llm_positions = torch.cat(llm_pos_ids_list, dim=1)
-        mrope_position_delta = torch.cat(llm_pos_ids_list, dim=1).max() + 1 - len(src_item)
-        llm_positions = llm_positions[:, context_len:seq_len]
-
-        return llm_positions, mrope_position_delta
-
     def generate_audio(self, code, voice_type):
         token2wav_dev = self._module_device(self.token2wav)
         if isinstance(code, torch.Tensor):
@@ -741,6 +570,7 @@ class Qwen2_5OmniForConditionalGeneration(
         thinker_prompt_embeds,
         prompt_token_ids,
     ):
+
         talker_hf_config = self.talker_config
         if hasattr(talker_hf_config, "talker_config"):
             talker_hf_config = talker_hf_config.talker_config
@@ -763,7 +593,9 @@ class Qwen2_5OmniForConditionalGeneration(
         input_tokens_len = len(prompt_token_ids_processed)
         # the code below is from model runner in Qwen, may need to further discuss later
         if input_tokens_len > 2:
-            prompt_token_ids_processed = [self.talker_config.tts_codec_mask_token_id] * (input_tokens_len - 2) + [
+            prompt_token_ids_processed = [
+                self.talker_config.tts_codec_mask_token_id
+            ] * (input_tokens_len - 2) + [
                 self.talker_config.tts_codec_pad_token_id,
                 self.talker_config.tts_codec_start_token_id,
             ]
@@ -774,7 +606,9 @@ class Qwen2_5OmniForConditionalGeneration(
             ][-input_tokens_len:]
         if isinstance(prompt_token_ids_processed, list):
             prompt_token_ids_processed = (
-                torch.Tensor(prompt_token_ids_processed).to(torch.int64).to(self._module_device(self.talker))
+                torch.Tensor(prompt_token_ids_processed)
+                .to(torch.int64)
+                .to(self._module_device(self.talker))
             )
         return prompt_token_ids_processed, prompt_embeds
 
@@ -783,12 +617,14 @@ class Qwen2_5OmniForConditionalGeneration(
         output_prompt_embeds,
         output_token_ids,
     ):
-        processed_output_token_embeds = output_prompt_embeds + self.talker.get_input_embeddings(
-            output_token_ids
+        processed_output_token_embeds = (
+            output_prompt_embeds + self.talker.get_input_embeddings(output_token_ids)
         )  # for decode
         return output_token_ids, processed_output_token_embeds
 
-    def compute_logits(self, hidden_states: Union[torch.Tensor, OmniOutput]) -> Optional[torch.Tensor]:
+    def compute_logits(
+        self, hidden_states: Union[torch.Tensor, OmniOutput]
+    ) -> Optional[torch.Tensor]:
         # Handle OmniOutput type
         if isinstance(hidden_states, OmniOutput):
             hidden_states = hidden_states.text_hidden_states
@@ -804,7 +640,7 @@ class Qwen2_5OmniForConditionalGeneration(
         # Use thinker model for sampling
         return self.model.sample(logits, sampling_metadata)
 
-    def generate_speech(self, text_tokens: torch.Tensor, voice_type: str = "default") -> torch.Tensor:
+    def generate_speech(self, text_tokens: torch.Tensor, voice_type: str = "default"):
         """
         Generate speech from text tokens using the talker and token2wav models.
         This method is kept for backward compatibility and direct speech generation.
@@ -817,7 +653,9 @@ class Qwen2_5OmniForConditionalGeneration(
             Audio tensor
         """
         # Generate codec tokens using talker model
-        talker_output = self.talker(input_ids=None, positions=None, inputs_embeds=text_tokens)
+        talker_output = self.talker(
+            input_ids=None, positions=None, inputs_embeds=text_tokens
+        )
 
         # Convert talker output to codec tokens
         codec_tokens = self._convert_to_codec_tokens(talker_output)
@@ -844,7 +682,9 @@ class Qwen2_5OmniForConditionalGeneration(
             # Suppress only codec_bos, consistent with HF generate's
             # suppress_tokens behavior
             bos_id = None
-            if hasattr(self, "talker_config") and hasattr(self.talker_config, "tts_codec_start_token_id"):
+            if hasattr(self, "talker_config") and hasattr(
+                self.talker_config, "tts_codec_start_token_id"
+            ):
                 bos_id = int(getattr(self.talker_config, "tts_codec_start_token_id"))
             if bos_id is not None:
                 logits[..., bos_id] = -1e9
@@ -863,8 +703,12 @@ class Qwen2_5OmniForConditionalGeneration(
         conds = getattr(self.token2wav_config, "conds", None)
         ref_mels = getattr(self.token2wav_config, "ref_mels", None)
         if isinstance(conds, dict) and isinstance(ref_mels, dict):
-            self._token2wav_conds = {k: torch.as_tensor(v, device=device) for k, v in conds.items()}
-            self._token2wav_ref_mels = {k: torch.as_tensor(v, device=device) for k, v in ref_mels.items()}
+            self._token2wav_conds = {
+                k: torch.as_tensor(v, device=device) for k, v in conds.items()
+            }
+            self._token2wav_ref_mels = {
+                k: torch.as_tensor(v, device=device) for k, v in ref_mels.items()
+            }
         # legacy: load from directory if provided
         model_path = hf_model_folder
         if isinstance(model_path, str) and os.path.isdir(model_path):
@@ -876,14 +720,24 @@ class Qwen2_5OmniForConditionalGeneration(
                     self._token2wav_ref_mels[key] = value["ref_mel"].to(device)
             else:
                 # legacy npy inputs
-                for f in sorted(glob.glob(os.path.join(model_path, "inputs", "*spk_emb.npy"))):
+                for f in sorted(
+                    glob.glob(os.path.join(model_path, "inputs", "*spk_emb.npy"))
+                ):
                     key = os.path.basename(f).split("_")[0].lower()
-                    self._token2wav_conds[key] = torch.as_tensor(np.load(f), device=device)
-                for f in sorted(glob.glob(os.path.join(model_path, "inputs", "*ref_mel.npy"))):
+                    self._token2wav_conds[key] = torch.as_tensor(
+                        np.load(f), device=device
+                    )
+                for f in sorted(
+                    glob.glob(os.path.join(model_path, "inputs", "*ref_mel.npy"))
+                ):
                     key = os.path.basename(f).split("_")[0].lower()
-                    self._token2wav_ref_mels[key] = torch.as_tensor(np.load(f), device=device)
+                    self._token2wav_ref_mels[key] = torch.as_tensor(
+                        np.load(f), device=device
+                    )
 
-    def _codec_to_audio(self, codec_tokens: torch.Tensor, voice_type: str = "default") -> Optional[torch.Tensor]:
+    def _codec_to_audio(
+        self, codec_tokens: torch.Tensor, voice_type: str = "default"
+    ) -> Optional[torch.Tensor]:
         if self.token2wav is None:
             self._init_token2wav_model()
         if self.token2wav is None:
@@ -917,7 +771,9 @@ class Qwen2_5OmniForConditionalGeneration(
             if codec.ndim == 1:
                 codec = codec.unsqueeze(0)
         else:
-            codec = torch.as_tensor(codec_tokens, dtype=torch.long, device=token2wav_dev).unsqueeze(0)
+            codec = torch.as_tensor(
+                codec_tokens, dtype=torch.long, device=token2wav_dev
+            ).unsqueeze(0)
 
         # Streaming with chunked process and boundary alignment
         # (rely on token2wav.process_chunk)
@@ -932,7 +788,9 @@ class Qwen2_5OmniForConditionalGeneration(
         steps = 10
 
         # Prepare initial noise for the whole sequence
-        y_all = torch.randn((1, total_mel, mel_dim), dtype=ref_mel.dtype, device=token2wav_dev)
+        y_all = torch.randn(
+            (1, total_mel, mel_dim), dtype=ref_mel.dtype, device=token2wav_dev
+        )
 
         logger.info(
             "Currently, we do not use the chunked process, we only use the "
@@ -944,7 +802,9 @@ class Qwen2_5OmniForConditionalGeneration(
         for i in range(codec.shape[1]):
             chunk_code_length = i * 2 - 24
             finished = i == (codec.shape[1] - 1)
-            if (chunk_code_length > 0 and chunk_code_length % chunk_size == 0) or finished:
+            if (
+                chunk_code_length > 0 and chunk_code_length % chunk_size == 0
+            ) or finished:
                 chunk_ends.append(i)
 
         # Number of chunks in mel domain
@@ -973,7 +833,7 @@ class Qwen2_5OmniForConditionalGeneration(
         waveform = np.concatenate(wav_chunks)
         return torch.as_tensor(waveform, device=token2wav_dev)
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
         """Load weights for all components of the omni model."""
         loaded_weights = set()
         thinker_weights = []
@@ -1002,14 +862,20 @@ class Qwen2_5OmniForConditionalGeneration(
         if talker_weights and self.talker is not None:
             # Map talker weights to appropriate components
             if self.thinker is None:
-                thinker_embedding_weights = [w for n, w in thinker_weights if n == "thinker.model.embed_tokens.weight"]
+                thinker_embedding_weights = [
+                    w
+                    for n, w in thinker_weights
+                    if n == "thinker.model.embed_tokens.weight"
+                ]
                 if thinker_embedding_weights:
                     self.thinker_embedding = nn.Embedding(
                         thinker_embedding_weights[0].shape[0],
                         thinker_embedding_weights[0].shape[1],
                     )
                     self.thinker_embedding.weight = nn.Parameter(
-                        thinker_embedding_weights[0].to(self._module_device(self.talker))
+                        thinker_embedding_weights[0].to(
+                            self._module_device(self.talker)
+                        )
                     )
             talker_loaded = self.talker.load_weights(talker_weights)
             talker_loaded = add_prefix_to_loaded_weights(talker_loaded, "talker")
@@ -1030,7 +896,9 @@ class Qwen2_5OmniForConditionalGeneration(
                     allow_patterns=["*.pt"],
                 )
             self._init_token2wav_model(hf_model_folder)
-            t2w_loaded = self.token2wav.load_weights(token2wav_weights, os.path.join(hf_model_folder, "spk_dict.pt"))
+            t2w_loaded = self.token2wav.load_weights(
+                token2wav_weights, os.path.join(hf_model_folder, "spk_dict.pt")
+            )
             t2w_loaded = add_prefix_to_loaded_weights(t2w_loaded, "token2wav")
             loaded_weights.update(t2w_loaded)
 

vllm_omni/model_executor/models/qwen2_5_omni_talker.py

@@ -1,17 +1,24 @@
-from collections.abc import Iterable
 from functools import cached_property
-from typing import Optional, Union
+from typing import Iterable, Optional, Set, Tuple, Union
 
 import torch
 import torch.nn as nn
-from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniTalkerConfig
-from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
+from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
+    Qwen2_5OmniTalkerConfig,
+)
+from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
+    Qwen2_5OmniAudioEncoder,
+)
 
 # from vllm.attention import AttentionMetadata  # unused import
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import ColumnParallelLinear
-from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    SupportsPP,
+)
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
     Qwen2_5OmniThinkerDummyInputsBuilder,
@@ -65,7 +72,9 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
 
         if hasattr(config, "talker_config"):
             self.config = config.talker_config
-            vllm_config.model_config.hf_text_config = vllm_config.model_config.hf_config.talker_config
+            vllm_config.model_config.hf_text_config = (
+                vllm_config.model_config.hf_config.talker_config
+            )
         else:
             self.config = config
 
@@ -83,7 +92,9 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
             hf_config=getattr(self.config, "text_config", self.config),
             architectures=["Qwen2ForCausalLM_old"],
         )
-        self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
+        self.make_empty_intermediate_tensors = (
+            self.language_model.make_empty_intermediate_tensors
+        )
 
     def init_multi_modal(self, thinker_config):
         self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
@@ -111,6 +122,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
     ) -> torch.Tensor:
         inputs_embeds = self.language_model.get_input_embeddings(input_ids)
         if multimodal_embeddings is not None and len(multimodal_embeddings) != 0:
+
             # TODO (ywang96): support overlapping modalitiy embeddings so that
             # `use_audio_in_video` will work on V1.
             inputs_embeds = merge_multimodal_embeddings(
@@ -133,7 +145,9 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
         inputs_embeds: Optional[torch.Tensor] = None,
         **kwargs: object,
     ) -> Union[torch.Tensor, IntermediateTensors]:
-        assert input_ids is not None or inputs_embeds is not None, "input_ids or inputs_embeds must be provided"
+        assert (
+            input_ids is not None or inputs_embeds is not None
+        ), "input_ids or inputs_embeds must be provided"
         # forward_context: ForwardContext = get_forward_context()  # unused variable
 
         if intermediate_tensors is not None:
@@ -152,16 +166,8 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
         )
         return hidden_states
 
-    def bad_word_processor(self, logits: torch.Tensor) -> torch.Tensor:
-        if hasattr(self.config, "tts_codec_start_token_id"):
-            bos_id = int(getattr(self.config, "tts_codec_start_token_id"))
-            logits[..., bos_id] = -1e9
-        return logits
-
     def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]:
-        logits = self.language_model.compute_logits(hidden_states)
-        logits = self.bad_word_processor(logits)
-        return logits
+        return self.language_model.compute_logits(hidden_states)
 
     def sample(
         self,
@@ -170,7 +176,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
     ) -> Optional[SamplerOutput]:
         return self.language_model.sample(logits, sampling_metadata)
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
         loader = AutoWeightsLoader(
             self,
             skip_prefixes=["thinker.", "token2wav."],
@@ -206,15 +212,31 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
         # Preserve the order of modalities if there are multiple of them
         # from the order of kwargs.
         for input_key in kwargs:
-            if input_key in ("pixel_values", "image_embeds") and "image" not in mm_input_by_modality:
-                mm_input_by_modality["image"] = self._parse_and_validate_image_input(**kwargs)
-            if input_key in ("pixel_values_videos", "video_embeds") and "video" not in mm_input_by_modality:
-                mm_input_by_modality["video"] = self._parse_and_validate_video_input(**kwargs)
-            if input_key in ("input_audio_features") and "audio" not in mm_input_by_modality:
-                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(**kwargs)
+            if (
+                input_key in ("pixel_values", "image_embeds")
+                and "image" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["image"] = self._parse_and_validate_image_input(
+                    **kwargs
+                )
+            if (
+                input_key in ("pixel_values_videos", "video_embeds")
+                and "video" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["video"] = self._parse_and_validate_video_input(
+                    **kwargs
+                )
+            if (
+                input_key in ("input_audio_features")
+                and "audio" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(
+                    **kwargs
+                )
         return mm_input_by_modality
 
     def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings:
+
         mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
         if not mm_input_by_modality:
             return []

vllm_omni/model_executor/models/qwen2_5_omni_thinker.py

@@ -1,18 +1,30 @@
-from collections import defaultdict
 from collections.abc import Iterable, Mapping, Sequence
+from copy import copy
 from functools import partial
-from typing import Annotated, Any, Callable, Literal, Optional, Union
+from typing import Any, Optional, Union
 
 import torch
 import torch.nn as nn
 from transformers.feature_extraction_utils import BatchFeature
-from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniConfig, Qwen2_5OmniThinkerConfig
-from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
-from transformers.models.qwen2_5_omni.processing_qwen2_5_omni import Qwen2_5OmniProcessor
+from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
+    Qwen2_5OmniConfig,
+    Qwen2_5OmniThinkerConfig,
+)
+from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
+    Qwen2_5OmniAudioEncoder,
+)
+from transformers.models.qwen2_5_omni.processing_qwen2_5_omni import (
+    Qwen2_5OmniProcessor,
+)
 from transformers.models.whisper import WhisperFeatureExtractor
+
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMRoPE, SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    SupportsPP,
+)
 from vllm.model_executor.models.qwen2_5_vl import (
     Qwen2_5_VisionTransformer,
     Qwen2_5_VLImageEmbeddingInputs,
@@ -24,6 +36,7 @@ from vllm.model_executor.models.qwen2_5_vl import (
     Qwen2_5_VLVideoPixelInputs,
 )
 from vllm.model_executor.models.qwen2_audio import (
+    Qwen2AudioInputs,
     Qwen2AudioProcessingInfo,
     _get_feat_extract_output_lengths,
 )
@@ -41,7 +54,7 @@ from vllm.multimodal.inputs import (
     ModalityData,
     MultiModalDataDict,
     MultiModalFieldConfig,
-    MultiModalKwargsItems,
+    MultiModalKwargs,
     NestedTensors,
 )
 from vllm.multimodal.parse import (
@@ -53,18 +66,13 @@ from vllm.multimodal.parse import (
 )
 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
 
 try:
@@ -75,27 +83,6 @@ except (ImportError, ModuleNotFoundError):
 logger = init_logger(__name__)
 
 
-class Qwen2_5OmniAudioFeatureInputs(TensorSchema):
-    """
-    Dimensions:
-        - na: Number of audios
-        - nmb: Number of mel bins
-        - msl: Maximum sequence length
-        - tsl: Total sequence length
-    """
-
-    type: Literal["audio_features"]
-    input_features: Annotated[
-        Union[torch.Tensor, list[torch.Tensor]],
-        TensorShape("nmb", "tsl"),
-    ]
-
-    feature_attention_mask: Annotated[
-        torch.Tensor,
-        TensorShape("na", "msl"),
-    ]
-
-
 def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
     audio_feature_lengths = hf_inputs.get("audio_feature_lengths", torch.empty((0,)))
 
@@ -106,53 +93,25 @@ def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
     video_grid_sizes = video_grid_thw.prod(-1)
 
     return dict(
-        input_audio_features=MultiModalFieldConfig.flat_from_sizes("audio", audio_feature_lengths, dim=1),
+        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_grid_sizes),
         image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
         image_grid_thw=MultiModalFieldConfig.batched("image"),
-        pixel_values_videos=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes),
+        pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
+            "video", video_grid_sizes
+        ),
         video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes),
         video_grid_thw=MultiModalFieldConfig.batched("video"),
         second_per_grid_ts=MultiModalFieldConfig.batched("video"),
     )
 
 
-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,
         data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],
@@ -168,7 +127,10 @@ class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
         return super()._parse_audio_data(data)
 
 
-class Qwen2_5OmniThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo):
+class Qwen2_5OmniThinkerProcessingInfo(
+    Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo
+):
+
     def get_hf_config(self):
         return self.ctx.get_hf_config(Qwen2_5OmniConfig).thinker_config
 
@@ -189,7 +151,10 @@ class Qwen2_5OmniThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProce
         return {"audio": None, "image": None, "video": None}
 
 
-class Qwen2_5OmniThinkerDummyInputsBuilder(BaseDummyInputsBuilder[Qwen2_5OmniThinkerProcessingInfo]):
+class Qwen2_5OmniThinkerDummyInputsBuilder(
+    BaseDummyInputsBuilder[Qwen2_5OmniThinkerProcessingInfo]
+):
+
     def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
         num_audios = mm_counts.get("audio", 0)
         num_images = mm_counts.get("image", 0)
@@ -201,7 +166,11 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(BaseDummyInputsBuilder[Qwen2_5OmniThi
         image_token: str = hf_processor.image_token
         video_token: str = hf_processor.video_token
 
-        return audio_token * num_audios + image_token * num_images + video_token * num_videos
+        return (
+            audio_token * num_audios
+            + image_token * num_images
+            + video_token * num_videos
+        )
 
     def get_dummy_mm_data(
         self,
@@ -222,25 +191,36 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(BaseDummyInputsBuilder[Qwen2_5OmniThi
             * feature_extractor.sampling_rate
         )
         target_width, target_height = self.info.get_image_size_with_most_features()
-        target_num_frames = self.info.get_num_frames_with_most_features(seq_len, mm_counts)
+        target_num_frames = self.info.get_num_frames_with_most_features(
+            seq_len, mm_counts
+        )
 
         mm_data = {
-            "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),
+            "audio": self._get_dummy_audios(
+                length=target_audio_length, num_audios=num_audios
+            ),
+            "image": self._get_dummy_images(
+                width=target_width, height=target_height, num_images=num_images
+            ),
             "video": self._get_dummy_videos(
-                width=target_width, height=target_height, num_frames=target_num_frames, num_videos=num_videos
+                width=target_width,
+                height=target_height,
+                num_frames=target_num_frames,
+                num_videos=num_videos,
             ),
         }
 
         return mm_data
 
 
-class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]):
+class Qwen2_5OmniThinkerMultiModalProcessor(
+    BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]
+):
+
     def _get_data_parser(self) -> MultiModalDataParser:
         feature_extractor = self.info.get_feature_extractor()
         return Qwen2_5OmniThinkerMultiModalDataParser(
-            spatial_merge_size=self.info.get_hf_config().vision_config.spatial_merge_size,
-            target_sr=feature_extractor.sampling_rate,
+            target_sr=feature_extractor.sampling_rate
         )
 
     def _call_hf_processor(
@@ -272,18 +252,15 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
         feature_attention_mask = hf_inputs.get("feature_attention_mask", None)
         if "input_audio_features" not in hf_inputs and input_features is not None:
             if feature_attention_mask is not None:
-                input_features = input_features.permute(0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
+                input_features = input_features.permute(0, 2, 1)[
+                    feature_attention_mask.bool()
+                ].permute(1, 0)
             hf_inputs["input_audio_features"] = input_features
-        if "audio_feature_lengths" not in hf_inputs and feature_attention_mask is not None:
+        if (
+            "audio_feature_lengths" not in hf_inputs
+            and feature_attention_mask is not None
+        ):
             hf_inputs["audio_feature_lengths"] = feature_attention_mask.sum(-1)
-
-        video_second_per_grid = hf_inputs.get("video_second_per_grid", None)
-        if video_second_per_grid is not None:
-            hf_inputs["second_per_grid_ts"] = video_second_per_grid
-
-        use_audio_in_video = mm_kwargs.get("use_audio_in_video", False)
-        hf_inputs["use_audio_in_video"] = torch.tensor(use_audio_in_video)
-
         return hf_inputs
 
     def _get_mm_fields_config(
@@ -291,177 +268,70 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
         hf_inputs: BatchFeature,
         hf_processor_mm_kwargs: Mapping[str, object],
     ) -> Mapping[str, MultiModalFieldConfig]:
-        return create_qwen2_5_omni_thinker_field_factory(self.info.get_hf_config().vision_config.spatial_merge_size)(
-            hf_inputs
-        )
-
-    def _apply_prompt_updates(
-        self,
-        token_ids: list[int],
-        mm_prompt_updates: MultiModalPromptUpdates,
-    ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
-        tokenizer = self.info.get_tokenizer()
-
-        new_token_ids, match_result = self._apply_token_matches(
-            token_ids,
-            mm_prompt_updates,
-        )
-
-        # If the search text does not represent a special token,
-        # it may have different token IDs in the prompt, because
-        # the tokens may go across the boundaries of the search text.
-        # ----
-        # e.g. when searching for "foo" in "food", if "food" itself makes
-        # up a token, then the token ID of "foo" will not appear at all
-        # ----
-        # Since it is inefficient to search for all possible tokenizations
-        # of the search text in the prompt, we instead perform string-based
-        # updates on the decoded token IDs, then encode them back.
-        if not all(all(update_idx is not None for update_idx in update_idxs) for update_idxs in match_result.values()):
-            new_text, match_result = self._apply_text_matches(
-                decode_tokens(tokenizer, token_ids),
-                mm_prompt_updates,
-            )
-
-            new_token_ids = encode_tokens(
-                tokenizer,
-                new_text,
-                add_special_tokens=False,
-            )
-
-        matched_updates = defaultdict[str, list[Sequence[ResolvedPromptUpdate]]](list)
-        unmatched_audio_items = []
-
-        for modality, update_idxs in match_result.items():
-            for item_idx, update_idx in enumerate(update_idxs):
-                if update_idx is None:
-                    # Check if this is audio that might be embedded in video
-                    if modality == "audio":
-                        unmatched_audio_items.append(item_idx)
-                        continue
-                    else:
-                        assert False, f"Failed to apply prompt replacement for mm_items[{modality!r}][{item_idx}]"
-
-                matched_updates[modality].append([mm_prompt_updates[modality][item_idx][update_idx]])
-
-        # If there are unmatched audio items, check if we can derive them from video
-        if unmatched_audio_items:
-            # Check if video exists in the original updates (not just matched ones)
-            if "video" in mm_prompt_updates:
-                num_videos = len(mm_prompt_updates.get("video", []))
-                num_audios = len(mm_prompt_updates.get("audio", []))
-                # If counts match, audio might be embedded in video - skip assertion
-                # The placeholders will be derived from video in _maybe_apply_prompt_updates
-                if num_audios == num_videos and "video" in matched_updates:
-                    pass  # Will be handled by deriving from video placeholders
-                else:
-                    # Audio items exist but can't be matched
-                    # Check if video was also matched
-                    if "video" not in matched_updates:
-                        # Neither audio nor video matched - this is a real error
-                        assert False, (
-                            f"Failed to apply prompt replacement for {len(unmatched_audio_items)} "
-                            f"audio item(s). Audio items cannot be matched in prompt."
-                        )
-                    elif num_audios != num_videos:
-                        # Audio and video counts don't match
-                        assert False, (
-                            f"Failed to apply prompt replacement for {len(unmatched_audio_items)} "
-                            f"audio item(s). Audio items cannot be matched in prompt and "
-                            f"audio count ({num_audios}) does not match video count ({num_videos})."
-                        )
-                    else:
-                        # Video matched but audio didn't, and counts match - allow it
-                        pass
-            else:
-                # Audio items can't be matched and there are no video items
-                assert False, (
-                    f"Failed to apply prompt replacement for {len(unmatched_audio_items)} "
-                    f"audio item(s). Audio items cannot be matched in prompt."
-                )
-
-        placeholders = self._find_mm_placeholders(
-            new_token_ids,
-            dict(matched_updates),
-        )
-
-        return new_token_ids, placeholders
+        return _qwen2_5_omni_thinker_field_config(hf_inputs)
 
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
         prompt_ids: list[int],
-        mm_kwargs: MultiModalKwargsItems,
-        mm_prompt_updates: MultiModalPromptUpdates,
+        mm_kwargs: MultiModalKwargs,
         is_update_applied: bool,
     ) -> tuple[list[int], str, Mapping[str, list[PlaceholderFeaturesInfo]]]:
         """
         Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.
         """
+        unbound_prompt_updates = self._get_prompt_updates(
+            mm_items,
+            hf_processor_mm_kwargs,
+            mm_kwargs,
+        )
+        mm_prompt_updates = self._bind_and_group_updates(unbound_prompt_updates)
+
         mm_item_counts = mm_items.get_all_counts()
         self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
 
-        use_audio_in_video = False
-        if "video" in mm_kwargs:
-            video_items = [item for item in mm_kwargs["video"] if item is not None]
-            # only check video items (if there are any)
-            if video_items:
-                use_audio_in_video = all(item["use_audio_in_video"].data for item in video_items)
+        use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)
 
         if is_update_applied:
             mm_placeholders = self._find_mm_placeholders(
+                mm_prompt_updates,
+                prompt_ids,
+                mm_item_counts,
+            )
+            self._validate_mm_placeholders(
+                mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video
+            )
+
+            tokenizer = self.info.get_tokenizer()
+            prompt = decode_tokens(tokenizer, prompt_ids)
+        else:
+            (
+                prompt_ids,
+                prompt,
+                mm_placeholders,
+            ) = self._apply_prompt_updates(
                 prompt_ids,
                 mm_prompt_updates,
-            )
-            self._validate_mm_placeholders(
-                mm_placeholders,
                 mm_item_counts,
             )
-
-        else:
-            # When use_audio_in_video=True, audio tokens are not in the prompt,
-            # so we need to filter out audio updates before applying replacements
-            if use_audio_in_video and "audio" in mm_prompt_updates:
-                # Remove audio from prompt updates (it won't match anything)
-                filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"}
-                prompt_ids, mm_placeholders = self._apply_prompt_updates(
-                    prompt_ids,
-                    filtered_updates,
-                )
-                # Derive audio placeholders from video placeholders
-                mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates)
-            else:
-                # Apply prompt updates normally
-                # _apply_prompt_updates will handle unmatched audio items
-                # when video exists with matching counts
-                prompt_ids, mm_placeholders = self._apply_prompt_updates(
-                    prompt_ids,
-                    mm_prompt_updates,
-                )
-
-                # Check if audio placeholders are missing but should exist
-                # This can happen when audio items are embedded in video tokens
-                if "audio" in mm_prompt_updates and "audio" not in mm_placeholders and "video" in mm_placeholders:
-                    num_audios = len(mm_prompt_updates.get("audio", []))
-                    num_videos = len(mm_placeholders.get("video", []))
-                    # If counts match, derive audio placeholders from video
-                    if num_audios == num_videos:
-                        mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates)
             self._validate_mm_placeholders(
-                mm_placeholders,
-                mm_item_counts,
+                mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video
             )
 
         tokenizer = self.info.get_tokenizer()
         prompt = decode_tokens(tokenizer, prompt_ids)
 
+        if use_audio_in_video:
+            mm_kwargs["use_audio_in_video"] = True
+
         return prompt_ids, prompt, mm_placeholders
 
     def _get_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
         hf_processor_mm_kwargs: Mapping[str, Any],
-        out_mm_kwargs: MultiModalKwargsItems,
+        out_mm_kwargs: MultiModalKwargs,
     ) -> Sequence[PromptUpdate]:
         processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
         tokenizer = self.info.get_tokenizer()
@@ -475,17 +345,20 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
         image_token_id = vocab[image_token]
         video_token_id = vocab[video_token]
 
-        out_mm_data = out_mm_kwargs.get_data()
-        audio_feature_lengths = out_mm_data.get("audio_feature_lengths")
-        feature_attention_mask = out_mm_data.get("feature_attention_mask")
+        audio_feature_lengths = out_mm_kwargs.get("audio_feature_lengths")
+        feature_attention_mask = out_mm_kwargs.get("feature_attention_mask")
         if audio_feature_lengths is None and feature_attention_mask is None:
             audio_output_lengths = []
         elif audio_feature_lengths is not None:
-            _, audio_output_lens = _get_feat_extract_output_lengths(audio_feature_lengths)
+            _, audio_output_lens = _get_feat_extract_output_lengths(
+                audio_feature_lengths
+            )
             audio_output_lengths = audio_output_lens.tolist()
         elif feature_attention_mask is not None:
             assert isinstance(feature_attention_mask, torch.Tensor)
-            _, audio_output_lens = _get_feat_extract_output_lengths(feature_attention_mask.sum(-1))
+            _, audio_output_lens = _get_feat_extract_output_lengths(
+                feature_attention_mask.sum(-1)
+            )
             audio_output_lengths = audio_output_lens.tolist()
 
         # number of audios read from video.
@@ -499,13 +372,14 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
                 audios = mm_items.get_items("audio", AudioProcessorItems)
                 audio = audios.get(item_idx)
                 raise ValueError(
-                    f"The audio {audio} (len={len(audio)}) is too short to be represented inside the model"
+                    f"The audio {audio} (len={len(audio)}) is too short "
+                    "to be represented inside the model"
                 )
 
             return [audio_token_id] * num_features
 
         def get_replacement_qwen2_vision(item_idx: int, modality: str):
-            grid_thw = out_mm_data[f"{modality}_grid_thw"][item_idx]
+            grid_thw = out_mm_kwargs[f"{modality}_grid_thw"][item_idx]
             assert isinstance(grid_thw, torch.Tensor)
             merge_length = image_processor.merge_size**2
 
@@ -518,8 +392,10 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
         def get_replacement_qwen2_use_audio_in_video(item_idx: int):
             nonlocal audio_in_video_item_idx
 
-            audio_num_features = audio_output_lengths[audio_in_video_item_idx + item_idx]
-            video_grid_thw = out_mm_data["video_grid_thw"][item_idx]
+            audio_num_features = audio_output_lengths[
+                audio_in_video_item_idx + item_idx
+            ]
+            video_grid_thw = out_mm_kwargs["video_grid_thw"][item_idx]
 
             audio_in_video_item_idx += 1
 
@@ -529,15 +405,13 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
             else:
                 video_second_per_grid_t = 1.0
 
-            placeholder = MRotaryEmbedding.omni_get_updates_use_audio_in_video(
+            return MRotaryEmbedding.omni_get_updates_use_audio_in_video(
                 thinker_config=thinker_config,
                 audio_len=audio_num_features,
                 video_grid_thw=video_grid_thw,
                 video_second_per_grid_t=video_second_per_grid_t,
             )
 
-            return PromptUpdateDetails.select_token_id(placeholder, embed_token_id=video_token_id)
-
         video_replacement_fn = (
             get_replacement_qwen2_use_audio_in_video
             if use_audio_in_video
@@ -562,50 +436,6 @@ class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniT
             ),
         ]
 
-    def _derive_audio_from_video_placeholders(
-        self,
-        placeholders: Mapping[str, list[PlaceholderFeaturesInfo]],
-        mm_prompt_updates: MultiModalPromptUpdates,
-    ) -> Mapping[str, list[PlaceholderFeaturesInfo]]:
-        """
-        Helper to derive audio placeholders from video placeholders when
-        use_audio_in_video=True.
-        """
-        if "video" not in placeholders:
-            return placeholders
-
-        # Validate audio and video counts match
-        num_videos = len(placeholders["video"])
-        num_audios = len(mm_prompt_updates.get("audio", []))
-        if num_audios != num_videos:
-            raise ValueError(
-                f"use_audio_in_video requires equal number of audio and video items, got {num_audios=}, {num_videos=}"
-            )
-
-        tokenizer = self.info.get_tokenizer()
-        processor = self.info.get_hf_processor()
-        audio_token_id = tokenizer.get_vocab()[processor.audio_token]
-
-        result_placeholders = dict(placeholders)
-        audio_placeholders = []
-
-        # Each video is paired with one audio
-        for video_idx, video_placeholder in enumerate(placeholders["video"]):
-            # Create is_embed mask selecting only audio tokens
-            audio_is_embed = torch.tensor(video_placeholder.tokens) == audio_token_id
-
-            audio_placeholder = PlaceholderFeaturesInfo(
-                modality="audio",
-                item_idx=video_idx,
-                start_idx=video_placeholder.start_idx,
-                tokens=video_placeholder.tokens,
-                is_embed=audio_is_embed,
-            )
-            audio_placeholders.append(audio_placeholder)
-
-        result_placeholders["audio"] = audio_placeholders
-        return result_placeholders
-
     def _apply_hf_processor_main(
         self,
         prompt: Union[str, list[int]],
@@ -664,33 +494,55 @@ 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:
+
+    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)}")
+            raise ValueError(
+                f"Incorrect type of {name}. " f"Got type: {type(mm_input)}"
+            )
         if isinstance(mm_input, torch.Tensor):
-            if dim == 0:
-                return mm_input.reshape(-1, *mm_input.shape[2:])
             return torch.concat(list(mm_input), dim=dim)
         else:
             return torch.concat(mm_input, dim=dim)
 
-    def _parse_and_validate_audio_input(self, **kwargs: object) -> Optional[Qwen2_5OmniAudioFeatureInputs]:
+    def _parse_and_validate_audio_input(
+        self, **kwargs: object
+    ) -> Optional[Qwen2AudioInputs]:
         input_audio_features = kwargs.pop("input_audio_features", None)
         audio_feature_lengths = kwargs.pop("audio_feature_lengths", None)
         feature_attention_mask = kwargs.pop("feature_attention_mask", None)
         if input_audio_features is None:
             return None
-        input_audio_features = self._validate_and_reshape_mm_tensor(input_audio_features, "input_audio_features", dim=1)
+        input_audio_features = self._validate_and_reshape_mm_tensor(
+            input_audio_features, "input_audio_features", dim=1
+        )
         if feature_attention_mask is not None:
             feature_attention_mask = self._validate_and_reshape_mm_tensor(
                 feature_attention_mask, "feature_attention_mask"
             )
         if not isinstance(input_audio_features, (torch.Tensor, list)):
-            raise ValueError(f"Incorrect type of audio input features. Got type: {type(input_audio_features)}")
-        return Qwen2_5OmniAudioFeatureInputs(
-            type="audio_features",
+            raise ValueError(
+                "Incorrect type of audio input features. "
+                f"Got type: {type(input_audio_features)}"
+            )
+        return Qwen2AudioInputs(
             input_features=input_audio_features,
             audio_feature_lengths=audio_feature_lengths,
             feature_attention_mask=feature_attention_mask,
@@ -708,24 +560,42 @@ class Qwen2_5OmniConditionalGenerationMixin:
             return None
 
         if pixel_values is not None:
-            pixel_values = self._validate_and_reshape_mm_tensor(pixel_values, "image pixel values")
-            image_grid_thw = self._validate_and_reshape_mm_tensor(image_grid_thw, "image grid_thw")
+            pixel_values = self._validate_and_reshape_mm_tensor(
+                pixel_values, "image pixel values"
+            )
+            image_grid_thw = self._validate_and_reshape_mm_tensor(
+                image_grid_thw, "image grid_thw"
+            )
 
             if not isinstance(pixel_values, (torch.Tensor, list)):
-                raise ValueError(f"Incorrect type of image pixel values. Got type: {type(pixel_values)}")
+                raise ValueError(
+                    "Incorrect type of image pixel values. "
+                    f"Got type: {type(pixel_values)}"
+                )
 
             return Qwen2_5_VLImagePixelInputs(
-                type="pixel_values", pixel_values=pixel_values, image_grid_thw=image_grid_thw
+                type="pixel_values",
+                pixel_values=pixel_values,
+                image_grid_thw=image_grid_thw,
             )
 
         if image_embeds is not None:
-            image_embeds = self._validate_and_reshape_mm_tensor(image_embeds, "image embeds")
-            image_grid_thw = self._validate_and_reshape_mm_tensor(image_grid_thw, "image grid_thw")
+            image_embeds = self._validate_and_reshape_mm_tensor(
+                image_embeds, "image embeds"
+            )
+            image_grid_thw = self._validate_and_reshape_mm_tensor(
+                image_grid_thw, "image grid_thw"
+            )
 
             if not isinstance(image_embeds, torch.Tensor):
-                raise ValueError(f"Incorrect type of image embeddings. Got type: {type(image_embeds)}")
+                raise ValueError(
+                    "Incorrect type of image embeddings. "
+                    f"Got type: {type(image_embeds)}"
+                )
             return Qwen2_5_VLImageEmbeddingInputs(
-                type="image_embeds", image_embeds=image_embeds, image_grid_thw=image_grid_thw
+                type="image_embeds",
+                image_embeds=image_embeds,
+                image_grid_thw=image_grid_thw,
             )
 
     def _parse_and_validate_video_input(
@@ -740,8 +610,12 @@ class Qwen2_5OmniConditionalGenerationMixin:
             return None
 
         if pixel_values_videos is not None:
-            pixel_values_videos = self._validate_and_reshape_mm_tensor(pixel_values_videos, "video pixel values")
-            video_grid_thw = self._validate_and_reshape_mm_tensor(video_grid_thw, "video grid_thw")
+            pixel_values_videos = self._validate_and_reshape_mm_tensor(
+                pixel_values_videos, "video pixel values"
+            )
+            video_grid_thw = self._validate_and_reshape_mm_tensor(
+                video_grid_thw, "video grid_thw"
+            )
 
             return Qwen2_5_VLVideoPixelInputs(
                 type="pixel_values_videos",
@@ -750,35 +624,48 @@ class Qwen2_5OmniConditionalGenerationMixin:
             )
 
         if video_embeds is not None:
-            video_embeds = self._validate_and_reshape_mm_tensor(video_embeds, "video embeds")
-            video_grid_thw = self._validate_and_reshape_mm_tensor(video_grid_thw, "video grid_thw")
+            video_embeds = self._validate_and_reshape_mm_tensor(
+                video_embeds, "video embeds"
+            )
+            video_grid_thw = self._validate_and_reshape_mm_tensor(
+                video_grid_thw, "video grid_thw"
+            )
 
             if not isinstance(video_embeds, torch.Tensor):
-                raise ValueError(f"Incorrect type of video embeddings. Got type: {type(video_embeds)}")
+                raise ValueError(
+                    "Incorrect type of video embeddings. "
+                    f"Got type: {type(video_embeds)}"
+                )
             return Qwen2_5_VLVideoEmbeddingInputs(
-                type="video_embeds", video_embeds=video_embeds, video_grid_thw=video_grid_thw
+                type="video_embeds",
+                video_embeds=video_embeds,
+                video_grid_thw=video_grid_thw,
             )
 
     def _process_audio_input(
         self,
-        audio_input: Qwen2_5OmniAudioFeatureInputs,
+        audio_input: Qwen2AudioInputs,
         audio_hashes: list[str] = None,
         cached_audio_features: torch.Tensor = None,
     ) -> torch.Tensor:
+
         input_features = audio_input["input_features"]
         audio_feature_lengths = audio_input["audio_feature_lengths"]
         if input_features.ndim == 3:
             assert input_features.shape[0] == 1
             input_features = input_features.squeeze(0)
         if audio_feature_lengths.ndim == 2:
-            assert audio_feature_lengths.shape[0] == 1 or audio_feature_lengths.shape[1] == 1
+            assert (
+                audio_feature_lengths.shape[0] == 1
+                or audio_feature_lengths.shape[1] == 1
+            )
             if audio_feature_lengths.shape[0] == 1:
                 audio_feature_lengths = audio_feature_lengths.squeeze(0)
             else:
                 audio_feature_lengths = audio_feature_lengths.squeeze(1)
 
-        audio_feat_lengths, audio_output_lengths = self.audio_tower._get_feat_extract_output_lengths(
-            audio_feature_lengths
+        audio_feat_lengths, audio_output_lengths = (
+            self.audio_tower._get_feat_extract_output_lengths(audio_feature_lengths)
         )
 
         audio_outputs = self.audio_tower(
@@ -786,9 +673,12 @@ class Qwen2_5OmniConditionalGenerationMixin:
             feature_lens=audio_feature_lengths,
             aftercnn_lens=audio_feat_lengths,
         )
-        return audio_outputs.last_hidden_state.split(audio_output_lengths.tolist())
+        audio_features = audio_outputs.last_hidden_state
+        return audio_features.split(audio_output_lengths.tolist())
 
-    def _process_image_input(self, image_input: Qwen2_5_VLImageInputs) -> tuple[torch.Tensor, ...]:
+    def _process_image_input(
+        self, image_input: Qwen2_5_VLImageInputs
+    ) -> tuple[torch.Tensor, ...]:
         if image_input["type"] == "image_embeds":
             return image_input["image_embeds"].type(self.visual.dtype)
 
@@ -830,7 +720,7 @@ class Qwen2_5OmniConditionalGenerationMixin:
     dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
 class Qwen2_5OmniThinkerForConditionalGeneration(
-    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin, SupportsMRoPE
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
 ):
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
@@ -868,24 +758,17 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
         else:
             logger.warning(
                 "flash_attn is not available, the model may not yield the "
-                "exactly same result as the transformers implementation in the audio tower part."
+                "exactly same result as the transformers implementation "
+                "in the audio tower part."
             )
 
-        if multimodal_config.get_limit_per_prompt("audio"):
-            self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
-        else:
-            self.audio_tower = None
-
-        if multimodal_config.get_limit_per_prompt("image") or multimodal_config.get_limit_per_prompt("video"):
-            self.visual = Qwen2_5_VisionTransformer(
-                vision_config=thinker_config.vision_config,
-                norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
-                quant_config=quant_config,
-                prefix=maybe_prefix(prefix, "visual"),
-                # attn_backend_override=None,
-            )
-        else:
-            self.visual = None
+        self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
+        self.visual = Qwen2_5_VisionTransformer(
+            vision_config=thinker_config.vision_config,
+            norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "visual"),
+        )
         self.quant_config = quant_config
         self.language_model = init_vllm_registered_model(
             vllm_config=vllm_config,
@@ -894,7 +777,9 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
             architectures=["Qwen2ForCausalLM"],
         )
 
-        self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
+        self.make_empty_intermediate_tensors = (
+            self.language_model.make_empty_intermediate_tensors
+        )
 
     def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
         mm_input_by_modality = {}
@@ -902,18 +787,34 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
         # Preserve the order of modalities if there are multiple of them
         # from the order of kwargs.
         for input_key in kwargs:
-            if input_key in ("pixel_values", "image_embeds") and "image" not in mm_input_by_modality:
-                mm_input_by_modality["image"] = self._parse_and_validate_image_input(**kwargs)
-            if input_key in ("pixel_values_videos", "video_embeds") and "video" not in mm_input_by_modality:
-                mm_input_by_modality["video"] = self._parse_and_validate_video_input(**kwargs)
-            if input_key in ("input_audio_features") and "audio" not in mm_input_by_modality:
-                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(**kwargs)
+            if (
+                input_key in ("pixel_values", "image_embeds")
+                and "image" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["image"] = self._parse_and_validate_image_input(
+                    **kwargs
+                )
+            if (
+                input_key in ("pixel_values_videos", "video_embeds")
+                and "video" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["video"] = self._parse_and_validate_video_input(
+                    **kwargs
+                )
+            if (
+                input_key in ("input_audio_features")
+                and "audio" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(
+                    **kwargs
+                )
         return mm_input_by_modality
 
     def get_language_model(self) -> torch.nn.Module:
         return self.language_model
 
     def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings:
+
         mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
         if not mm_input_by_modality:
             return []
@@ -944,6 +845,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
     ) -> torch.Tensor:
         inputs_embeds = self.language_model.get_input_embeddings(input_ids)
         if multimodal_embeddings is not None and len(multimodal_embeddings) != 0:
+
             # TODO (ywang96): support overlapping modalitiy embeddings so that
             # `use_audio_in_video` will work on V1.
             inputs_embeds = merge_multimodal_embeddings(
@@ -995,7 +897,9 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
                 placeholder_token_id = self.config.image_token_index
             if modality == "video":
                 placeholder_token_id = self.config.video_token_index
-            inputs_embeds = merge_multimodal_embeddings(input_ids, inputs_embeds, embeddings, placeholder_token_id)
+            inputs_embeds = merge_multimodal_embeddings(
+                input_ids, inputs_embeds, embeddings, placeholder_token_id
+            )
         return inputs_embeds
 
     def forward(
@@ -1017,7 +921,10 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
             text_inputs_embeds = self.get_input_embeddings(
                 input_ids,
                 (
-                    [(torch.zeros_like(embeddings), modality) for embeddings, modality in multimodal_embeddings]
+                    [
+                        (torch.zeros_like(embeddings), modality)
+                        for embeddings, modality in multimodal_embeddings
+                    ]
                     if multimodal_embeddings is not None
                     else None
                 ),

vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py

@@ -3,8 +3,7 @@
 ############################
 
 import math
-from collections.abc import Iterable
-from typing import Optional, Union
+from typing import Iterable, List, Optional, Set, Tuple, Union
 
 import numpy as np
 import torch
@@ -17,17 +16,24 @@ from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
     Qwen2_5OmniDiTConfig,
     Qwen2_5OmniToken2WavConfig,
 )
-from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniPreTrainedModel
+from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
+    Qwen2_5OmniPreTrainedModel,
+)
 
 # Bring in HF base classes, configs and utilities used below
 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.models.interfaces import SupportsPP
-from vllm.model_executor.models.utils import AutoWeightsLoader as _Vllm_AutoWeightsLoader
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader as _Vllm_AutoWeightsLoader,
+)
 from vllm.model_executor.models.utils import WeightsMapper as _Vllm_WeightsMapper
-from vllm.model_executor.models.utils import init_vllm_registered_model as _vllm_init_vllm_registered_model
-from vllm.model_executor.models.utils import maybe_prefix as _vllm_maybe_prefix
+from vllm.model_executor.models.utils import (
+    init_vllm_registered_model as _Vllm_init_vllm_registered_model,
+)
+from vllm.model_executor.models.utils import maybe_prefix as _Vllm_maybe_prefix
 from vllm.sequence import IntermediateTensors
 from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
@@ -212,7 +218,9 @@ class AttentiveStatisticsPooling(nn.Module):
 
     def _compute_statistics(self, x, m, dim=2):
         mean = (m * x).sum(dim)
-        std = torch.sqrt((m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(self.eps))
+        std = torch.sqrt(
+            (m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(self.eps)
+        )
         return mean, std
 
     def forward(self, hidden_states):
@@ -274,7 +282,9 @@ class SqueezeExcitationRes2NetBlock(nn.Module):
             kernel_size=1,
             dilation=1,
         )
-        self.res2net_block = Res2NetBlock(out_channels, out_channels, res2net_scale, kernel_size, dilation)
+        self.res2net_block = Res2NetBlock(
+            out_channels, out_channels, res2net_scale, kernel_size, dilation
+        )
         self.tdnn2 = TimeDelayNetBlock(
             out_channels,
             out_channels,
@@ -302,10 +312,13 @@ class ECAPA_TimeDelayNet(torch.nn.Module):
 
     def __init__(self, config: Qwen2_5OmniDiTConfig):
         super().__init__()
-        if len(config.enc_channels) != len(config.enc_kernel_sizes) or len(config.enc_channels) != len(
-            config.enc_dilations
-        ):
-            raise ValueError("enc_channels, enc_kernel_sizes and enc_dilations should have same length")
+        if len(config.enc_channels) != len(config.enc_kernel_sizes) or len(
+            config.enc_channels
+        ) != len(config.enc_dilations):
+            raise ValueError(
+                "enc_channels, enc_kernel_sizes and enc_dilations should have "
+                "same length"
+            )
         self.channels = config.enc_channels
         self.blocks = nn.ModuleList()
 
@@ -399,14 +412,26 @@ class DiTInputEmbedding(nn.Module):
     ):
         if apply_cfg:
             hidden_states = torch.cat([hidden_states, hidden_states], dim=0)
-            speaker_embedding = torch.cat([speaker_embedding, torch.zeros_like(speaker_embedding)], dim=0)
-            condition_vector = torch.cat([condition_vector, torch.zeros_like(condition_vector)], dim=0)
+            speaker_embedding = torch.cat(
+                [speaker_embedding, torch.zeros_like(speaker_embedding)], dim=0
+            )
+            condition_vector = torch.cat(
+                [condition_vector, torch.zeros_like(condition_vector)], dim=0
+            )
             code_embed = torch.cat([code_embed, code_embed_uncond], dim=0)
         elif drop_audio_cond:  # cfg for cond audio
             condition_vector = torch.zeros_like(condition_vector)
             speaker_embedding = torch.zeros_like(speaker_embedding)
-        condition_vector = self.spk_encoder(condition_vector).unsqueeze(1).repeat(1, hidden_states.size(1), 1)
-        hidden_states = self.proj(torch.cat((hidden_states, condition_vector, code_embed, speaker_embedding), dim=-1))
+        condition_vector = (
+            self.spk_encoder(condition_vector)
+            .unsqueeze(1)
+            .repeat(1, hidden_states.size(1), 1)
+        )
+        hidden_states = self.proj(
+            torch.cat(
+                (hidden_states, condition_vector, code_embed, speaker_embedding), dim=-1
+            )
+        )
 
         return hidden_states
 
@@ -440,9 +465,13 @@ class Qwen2_5_OmniAdaLayerNormZero(nn.Module):
 
     def forward(self, hidden_states, emb=None):
         emb = self.linear(self.silu(emb))
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(
+            emb, 6, dim=1
+        )
 
-        hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        hidden_states = (
+            self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        )
         return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp
 
 
@@ -461,7 +490,9 @@ class Qwen2_5_OmniAdaLayerNormZero_Final(nn.Module):
         emb = self.linear(self.silu(emb))
         scale, shift = torch.chunk(emb, 2, dim=1)
 
-        hidden_states = self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
+        hidden_states = (
+            self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
+        )
         return hidden_states
 
 
@@ -541,7 +572,9 @@ class DiTAttention(nn.Module):
         self.to_k = nn.Linear(config.hidden_size, self.inner_dim)
         self.to_v = nn.Linear(config.hidden_size, self.inner_dim)
 
-        self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, config.hidden_size), nn.Dropout(config.dropout)])
+        self.to_out = nn.ModuleList(
+            [nn.Linear(self.inner_dim, config.hidden_size), nn.Dropout(config.dropout)]
+        )
 
     def forward(
         self,
@@ -567,7 +600,9 @@ class DiTAttention(nn.Module):
         # Due to training process, only first head is applied with RoPE,
         # will be fixed at next release
         cos, sin = position_embeddings
-        query[:, :1], key[:, :1] = apply_rotary_pos_emb(query[:, :1], key[:, :1], cos, sin)
+        query[:, :1], key[:, :1] = apply_rotary_pos_emb(
+            query[:, :1], key[:, :1], cos, sin
+        )
 
         attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
         attention_weights, _ = attention_interface(
@@ -581,7 +616,9 @@ class DiTAttention(nn.Module):
 
         # mask. e.g. inference got a batch with different target durations,
         # mask out the padding
-        attention_weights = attention_weights.reshape(batch_size, -1, self.heads * head_dim)
+        attention_weights = attention_weights.reshape(
+            batch_size, -1, self.heads * head_dim
+        )
         attention_weights = attention_weights.to(query.dtype)
 
         # linear proj
@@ -611,7 +648,9 @@ class DiTTimestepEmbedding(nn.Module):
     def __init__(self, dim, freq_embed_dim=256):
         super().__init__()
         self.time_embed = SinusPositionEmbedding(freq_embed_dim)
-        self.time_mlp = nn.ModuleList([nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)])
+        self.time_mlp = nn.ModuleList(
+            [nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)]
+        )
 
     def forward(self, timestep):  # noqa: F821
         time_hidden = self.time_embed(timestep)
@@ -622,21 +661,29 @@ class DiTTimestepEmbedding(nn.Module):
 
 
 class DiTDecoderLayer(nn.Module):
-    def __init__(self, config: Qwen2_5OmniDiTConfig, look_ahead_block=0, look_backward_block=0):
+    def __init__(
+        self, config: Qwen2_5OmniDiTConfig, look_ahead_block=0, look_backward_block=0
+    ):
         super().__init__()
         self.attn_norm = Qwen2_5_OmniAdaLayerNormZero(config.hidden_size)
 
         self.attn = DiTAttention(config)
         self.look_ahead_block = look_ahead_block
         self.look_backward_block = look_backward_block
-        self.ff_norm = nn.LayerNorm(config.hidden_size, elementwise_affine=False, eps=1e-6)
-        self.ff = DiTMLP(dim=config.hidden_size, mult=config.ff_mult, dropout=config.dropout)
+        self.ff_norm = nn.LayerNorm(
+            config.hidden_size, elementwise_affine=False, eps=1e-6
+        )
+        self.ff = DiTMLP(
+            dim=config.hidden_size, mult=config.ff_mult, dropout=config.dropout
+        )
 
     def forward(
         self, hidden_states, timestep, position_embeddings=None, block_diff=None
     ):  # x: noised input, t: time embedding
         # pre-norm & modulation for attention input
-        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(hidden_states, emb=timestep)
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(
+            hidden_states, emb=timestep
+        )
 
         # attention
         attn_output = self.attn(
@@ -649,7 +696,9 @@ class DiTDecoderLayer(nn.Module):
         # process attention output for input x
         hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_output
 
-        norm = self.ff_norm(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        norm = (
+            self.ff_norm(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        )
         ff_output = self.ff(norm)
         hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
 
@@ -692,14 +741,14 @@ class SnakeBeta(nn.Module):
         beta = self.beta.unsqueeze(0).unsqueeze(-1)
         alpha = torch.exp(alpha)
         beta = torch.exp(beta)
-        hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(
-            torch.sin(hidden_states * alpha), 2
-        )
+        hidden_states = hidden_states + (
+            1.0 / (beta + self.no_div_by_zero)
+        ) * torch.pow(torch.sin(hidden_states * alpha), 2)
 
         return hidden_states
 
 
-def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
+def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
     """Generates a 1D Kaiser-windowed sinc filter.
 
     Args:
@@ -724,7 +773,9 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->
     else:
         beta = 0.0
 
-    kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
+    kaiser_window = torch.kaiser_window(
+        kernel_size, beta=beta, periodic=False, dtype=torch.float32
+    )
 
     # Compute time indices
     if is_even:
@@ -734,7 +785,9 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->
 
     # Compute sinc filter
     if cutoff == 0:
-        return torch.zeros((1, 1, kernel_size), dtype=torch.float32)  # Ensures correct shape
+        return torch.zeros(
+            (1, 1, kernel_size), dtype=torch.float32
+        )  # Ensures correct shape
 
     sinc_filter = torch.sinc(2 * cutoff * time_indices)
     normalized_filter = 2 * cutoff * kaiser_window * sinc_filter
@@ -749,19 +802,27 @@ class UpSample1d(nn.Module):
     def __init__(self, ratio=2, kernel_size=None):
         super().__init__()
         self.ratio = ratio
-        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
         self.stride = ratio
         self.pad = self.kernel_size // ratio - 1
         self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
-        self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+        self.pad_right = (
+            self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+        )
 
-        filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
+        filter = kaiser_sinc_filter1d(
+            cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size
+        )
         self.register_buffer("filter", filter, persistent=False)
 
     def forward(self, hidden_states):
         channels = hidden_states.shape[1]
         hidden_states_dtype = hidden_states.dtype
-        hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(self.filter.dtype)
+        hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(
+            self.filter.dtype
+        )
         hidden_states = self.ratio * F.conv_transpose1d(
             hidden_states,
             self.filter.expand(channels, -1, -1),
@@ -794,7 +855,9 @@ class DownSample1d(nn.Module):
     def forward(self, hidden_states):
         channels = hidden_states.shape[1]
         hidden_states_dtype = hidden_states.dtype
-        hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate").to(self.filter.dtype)
+        hidden_states = F.pad(
+            hidden_states, (self.pad_left, self.pad_right), mode="replicate"
+        ).to(self.filter.dtype)
         out = F.conv1d(
             hidden_states,
             self.filter.expand(channels, -1, -1),
@@ -895,10 +958,15 @@ class AMPBlock(torch.nn.Module):
             ]
         )
 
-        self.num_layers = len(self.convs1) + len(self.convs2)  # total number of conv layers
+        self.num_layers = len(self.convs1) + len(
+            self.convs2
+        )  # total number of conv layers
 
         self.activations = nn.ModuleList(
-            [TorchActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)]
+            [
+                TorchActivation1d(activation=SnakeBeta(channels))
+                for _ in range(self.num_layers)
+            ]
         )
 
     def _get_padding(self, kernel_size, dilation=1):
@@ -931,7 +999,9 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
         self.num_residual_blocks = len(config.resblock_kernel_sizes)
         self.num_upsample_layers = len(config.upsample_rates)
 
-        self.conv_pre = nn.Conv1d(config.mel_dim, config.upsample_initial_channel, 7, 1, padding=3)
+        self.conv_pre = nn.Conv1d(
+            config.mel_dim, config.upsample_initial_channel, 7, 1, padding=3
+        )
 
         # Removing extra ModuleList breaks official state dict
         ups = [
@@ -946,7 +1016,9 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
                     )
                 ]
             )
-            for layer_idx, (stride, kernel_size) in enumerate(zip(config.upsample_rates, config.upsample_kernel_sizes))
+            for layer_idx, (stride, kernel_size) in enumerate(
+                zip(config.upsample_rates, config.upsample_kernel_sizes)
+            )
         ]
         self.ups = nn.ModuleList(ups)
 
@@ -958,12 +1030,16 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
                     dilation,
                 )
                 for layer_idx in range(self.num_upsample_layers)
-                for kernel_size, dilation in zip(config.resblock_kernel_sizes, config.resblock_dilation_sizes)
+                for kernel_size, dilation in zip(
+                    config.resblock_kernel_sizes, config.resblock_dilation_sizes
+                )
             ]
         )
 
         self.activation_post = TorchActivation1d(
-            activation=SnakeBeta(config.upsample_initial_channel // (2**self.num_upsample_layers))
+            activation=SnakeBeta(
+                config.upsample_initial_channel // (2**self.num_upsample_layers)
+            )
         )
         self.conv_post = nn.Conv1d(
             config.upsample_initial_channel // (2**self.num_upsample_layers),
@@ -1003,7 +1079,9 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
         for layer_index in range(self.num_upsample_layers):
             hidden_representation = self.ups[layer_index][0](hidden_representation)
             residual_output = sum(
-                self.resblocks[layer_index * self.num_residual_blocks + block_index](hidden_representation)
+                self.resblocks[layer_index * self.num_residual_blocks + block_index](
+                    hidden_representation
+                )
                 for block_index in range(self.num_residual_blocks)
             )
             residual_output = residual_output / self.num_residual_blocks
@@ -1031,7 +1109,11 @@ class RungeKutta4ODESolver:
         value_start,
         function_value_start=None,
     ):
-        k1 = function_value_start if function_value_start is not None else function(time_start, value_start)
+        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,
@@ -1057,7 +1139,9 @@ class RungeKutta4ODESolver:
             function_value_start,
         )
 
-    def _linear_interpolation(self, time_start, time_end, value_start, value_end, time_point):
+    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:
@@ -1078,10 +1162,15 @@ class RungeKutta4ODESolver:
         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)
+            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]:
+            while (
+                current_index < len(time_points)
+                and time_end >= time_points[current_index]
+            ):
                 solution[current_index] = self._linear_interpolation(
                     time_start,
                     time_end,
@@ -1112,7 +1201,9 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         self.repeats = config.repeats
         self.time_embed = DiTTimestepEmbedding(config.hidden_size)
 
-        self.text_embed = DiTCodecEmbedding(config.num_embeds, config.emb_dim, config.repeats)
+        self.text_embed = DiTCodecEmbedding(
+            config.num_embeds, config.emb_dim, config.repeats
+        )
         self.input_embed = DiTInputEmbedding(config)
 
         self.rotary_embed = Qwen2_5OmniDiTRotaryEmbedding(config.head_dim)
@@ -1132,12 +1223,16 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
                 )
             )
 
-        self.norm_out = Qwen2_5_OmniAdaLayerNormZero_Final(config.hidden_size)  # final modulation
+        self.norm_out = Qwen2_5_OmniAdaLayerNormZero_Final(
+            config.hidden_size
+        )  # final modulation
         self.proj_out = nn.Linear(config.hidden_size, config.mel_dim)
 
     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]
+        block_indices = (
+            torch.arange(seq_len, device=hidden_states.device) // self.block_size
+        )  # [seq_length]
 
         block_i = block_indices.unsqueeze(1)  # [seq_length, 1]
         block_j = block_indices.unsqueeze(0)  # [1, seq_length]
@@ -1162,8 +1257,12 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
 
         # Compute embeddings
         time_embedding = self.time_embed(time_step)
-        text_embedding = self.text_embed(quantized_code, drop_code=False if apply_cfg else drop_code)
-        text_embedding_unconditioned = self.text_embed(quantized_code, drop_code=True) if apply_cfg else None
+        text_embedding = self.text_embed(
+            quantized_code, drop_code=False if apply_cfg else drop_code
+        )
+        text_embedding_unconditioned = (
+            self.text_embed(quantized_code, drop_code=True) if apply_cfg else None
+        )
 
         hidden_states = self.input_embed(
             hidden_states,
@@ -1203,11 +1302,17 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         guidance_scale=0.5,
         sway_coefficient=-1.0,
     ):
-        noise_initialization = torch.randn([1, 30000, self.mel_dim], dtype=reference_mel_spectrogram.dtype)
+        noise_initialization = torch.randn(
+            [1, 30000, self.mel_dim], dtype=reference_mel_spectrogram.dtype
+        )
         maximum_duration = quantized_code.shape[1] * self.repeats
-        initial_state = noise_initialization[:, :maximum_duration].to(quantized_code.device)
+        initial_state = noise_initialization[:, :maximum_duration].to(
+            quantized_code.device
+        )
         batch_size = reference_mel_spectrogram.shape[0]
-        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, maximum_duration, 1)
+        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(
+            1, maximum_duration, 1
+        )
 
         if batch_size != 1:
             raise ValueError("Only batch size = 1 is currently supported")
@@ -1234,7 +1339,10 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
                 apply_cfg=True,
             )
             guided_prediction, null_prediction = torch.chunk(model_output, 2, dim=0)
-            return guided_prediction + (guided_prediction - null_prediction) * guidance_scale
+            return (
+                guided_prediction
+                + (guided_prediction - null_prediction) * guidance_scale
+            )
 
         initial_time = 0
         time_embedding = torch.linspace(
@@ -1246,9 +1354,13 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         )
 
         if sway_coefficient is not None:
-            time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
+            time_embedding += sway_coefficient * (
+                torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding
+            )
 
-        ode_solver = RungeKutta4ODESolver(function=ode_function, initial_value=initial_state)
+        ode_solver = RungeKutta4ODESolver(
+            function=ode_function, initial_value=initial_state
+        )
         solution_trajectory = ode_solver.integrate(time_embedding)
 
         generated_waveform = solution_trajectory[-1]
@@ -1258,14 +1370,14 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
     @torch.no_grad()
     def fast_block_sample(
         self,
-        conditioning_vector: torch.Tensor,
-        reference_mel_spectrogram: torch.Tensor,
-        quantized_code: torch.Tensor,
+        conditioning_vector,
+        reference_mel_spectrogram,
+        quantized_code,
         y0: torch.Tensor,
-        num_steps: int = 10,
-        guidance_scale: float = 0.5,
+        num_steps=10,
+        guidance_scale=0.5,
         sway_coefficient: Optional[float] = -1.0,
-    ) -> torch.Tensor:
+    ):
         """
         Block-wise ODE sampling starting from provided initial state y0.
 
@@ -1279,7 +1391,9 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         """
         initial_state = y0.to(quantized_code.device)
         batch_size = reference_mel_spectrogram.shape[0]
-        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, initial_state.shape[1], 1)
+        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(
+            1, initial_state.shape[1], 1
+        )
 
         if batch_size != 1:
             raise ValueError("Only batch size = 1 is currently supported")
@@ -1306,7 +1420,10 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
                 apply_cfg=True,
             )
             guided_prediction, null_prediction = torch.chunk(model_output, 2, dim=0)
-            return guided_prediction + (guided_prediction - null_prediction) * guidance_scale
+            return (
+                guided_prediction
+                + (guided_prediction - null_prediction) * guidance_scale
+            )
 
         initial_time = 0
         time_embedding = torch.linspace(
@@ -1318,9 +1435,13 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         )
 
         if sway_coefficient is not None:
-            time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
+            time_embedding += sway_coefficient * (
+                torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding
+            )
 
-        ode_solver = RungeKutta4ODESolver(function=ode_function, initial_value=initial_state)
+        ode_solver = RungeKutta4ODESolver(
+            function=ode_function, initial_value=initial_state
+        )
         solution_trajectory = ode_solver.integrate(time_embedding)
 
         generated_waveform = solution_trajectory[-1]
@@ -1356,7 +1477,8 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
             attn_impl = "sdpa"
         elif config._attn_implementation == "eager":
             logger.warning_once(
-                "Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa"
+                "Qwen2_5OmniToken2WavModel does not support eager attention "
+                "implementation, fall back to sdpa"
             )
             attn_impl = "sdpa"
         self.code2wav_dit_model = Qwen2_5OmniToken2WavDiTModel._from_config(
@@ -1379,7 +1501,9 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
 
         # codec embedding size for masking EOS out-of-range
         try:
-            self.codec_embed_size = self.code2wav_dit_model.text_embed.codec_embed.weight.size(0)
+            self.codec_embed_size = (
+                self.code2wav_dit_model.text_embed.codec_embed.weight.size(0)
+            )
         except Exception:
             self.codec_embed_size = -1
 
@@ -1461,7 +1585,7 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
         steps: int,
         prev_generated: torch.Tensor,
         finished: bool = False,
-    ) -> tuple[Optional[torch.Tensor], torch.Tensor]:
+    ) -> Tuple[Optional[torch.Tensor], torch.Tensor]:
         """Streaming per small chunk: returns (mel_or_None, audio_slice)."""
         start_index = max(i * self.chunk_size - self.past_cache_size, 0)
         end_index = min(
@@ -1469,8 +1593,16 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
             codec_all.shape[1] * self.factor,
         )
 
-        y0 = y_all[:, start_index:end_index].reshape(1, -1, self.code2wav_dit_model.mel_dim).contiguous()
-        codec = codec_all[:, start_index // self.factor : end_index // self.factor].reshape(1, -1).contiguous()
+        y0 = (
+            y_all[:, start_index:end_index]
+            .reshape(1, -1, self.code2wav_dit_model.mel_dim)
+            .contiguous()
+        )
+        codec = (
+            codec_all[:, start_index // self.factor : end_index // self.factor]
+            .reshape(1, -1)
+            .contiguous()
+        )
 
         # generate mel for current window (B, mel_dim, T)
         generated = self.process_chunk_dit_batch(
@@ -1500,9 +1632,9 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
         y_all: torch.Tensor,
         i: int,
         steps: int,
-        prev_generated: Union[torch.Tensor, list[torch.Tensor]],
+        prev_generated: Union[torch.Tensor, List[torch.Tensor]],
         finished: bool = False,
-    ) -> tuple[Union[torch.Tensor, list[torch.Tensor]], torch.Tensor]:
+    ) -> Tuple[Union[torch.Tensor, List[torch.Tensor]], torch.Tensor]:
         """High-level chunk API aligning to qwen2_code2wav_dit signature."""
         if not isinstance(prev_generated, torch.Tensor):
             prev_generated = prev_generated[0] if len(prev_generated) > 0 else None
@@ -1527,7 +1659,7 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
         finished: bool,
         prev_generated: Optional[torch.Tensor],
         generated: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         """
         Align mel and audio boundaries for 50Hz-like streaming.
 
@@ -1544,13 +1676,21 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
             mel = generated[:, :, : self.chunk_size]
         elif finished:
             mel_trim = generated[:, :, self.past_cache_size :]
-            mel = torch.cat([prev_generated[:, :, -self.future_size * 2 :], mel_trim], dim=2)
+            mel = torch.cat(
+                [prev_generated[:, :, -self.future_size * 2 :], mel_trim], dim=2
+            )
         else:
             if start_index == 0:
-                mel_trim = generated[:, :, i * self.chunk_size : -self.future_cache_size]
+                mel_trim = generated[
+                    :, :, i * self.chunk_size : -self.future_cache_size
+                ]
             else:
-                mel_trim = generated[:, :, self.past_cache_size : -self.future_cache_size]
-            mel = torch.cat([prev_generated[:, :, -self.future_size * 2 :], mel_trim], dim=2)
+                mel_trim = generated[
+                    :, :, self.past_cache_size : -self.future_cache_size
+                ]
+            mel = torch.cat(
+                [prev_generated[:, :, -self.future_size * 2 :], mel_trim], dim=2
+            )
 
         audio = self.code2wav_bigvgan_model(mel)
         if i == 0:
@@ -1558,7 +1698,11 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
         elif finished:
             audio_output = audio[self.future_size * self.vocoder_hop :]
         else:
-            audio_output = audio[self.future_size * self.vocoder_hop : -self.future_size * self.vocoder_hop]
+            audio_output = audio[
+                self.future_size
+                * self.vocoder_hop : -self.future_size
+                * self.vocoder_hop
+            ]
         return mel, audio_output
 
 
@@ -1582,9 +1726,9 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         self.config = vllm_config.model_config.hf_config
 
         # Initialize underlying HF Token2Wav model via registry
-        self.token2wav = _vllm_init_vllm_registered_model(
+        self.token2wav = _Vllm_init_vllm_registered_model(
             vllm_config=vllm_config,
-            prefix=_vllm_maybe_prefix(prefix, "token2wav_model"),
+            prefix=_Vllm_maybe_prefix(prefix, "token2wav_model"),
             hf_config=self.config,
             architectures=["Qwen2_5OmniToken2WavDiTModel"],
         )
@@ -1636,7 +1780,9 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
     ) -> Optional[SamplerOutput]:
         return None
 
-    def load_weights_without_buffers(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+    def load_weights_without_buffers(
+        self, weights: Iterable[Tuple[str, torch.Tensor]]
+    ) -> Set[str]:
         loader = _Vllm_AutoWeightsLoader(self)
         loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
         # Log load summary
@@ -1674,7 +1820,9 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         return registers
 
     # remove buffers from the weights and reload them after loading weights
-    def remove_buffers_from_weights(self, weights: Iterable[tuple[str, torch.Tensor]], buffers: dict):
+    def remove_buffers_from_weights(
+        self, weights: Iterable[Tuple[str, torch.Tensor]], buffers: dict
+    ):
         weights_to_load = []
         for key, value in weights:
             if key in buffers:
@@ -1694,7 +1842,9 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
                 loaded_buffers.add(name)
         return loaded_buffers
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]], spk_dict_path: str) -> set[str]:
+    def load_weights(
+        self, weights: Iterable[Tuple[str, torch.Tensor]], spk_dict_path: str
+    ) -> Set[str]:
         buffers = self.find_all_registers()
         weights_to_load = self.remove_buffers_from_weights(weights, buffers)
         loaded = self.load_weights_without_buffers(weights_to_load)
@@ -1722,7 +1872,9 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         )
 
     @torch.inference_mode()
-    def process_chunk_bigvgan_batch(self, mel_batch: torch.Tensor) -> Optional[torch.Tensor]:
+    def process_chunk_bigvgan_batch(
+        self, mel_batch: torch.Tensor
+    ) -> Optional[torch.Tensor]:
         # BigVGAN is not part of this wrapper; return None for parity.
         return None
 
@@ -1737,7 +1889,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         steps: int,
         prev_generated: torch.Tensor,
         finished: bool = False,
-    ) -> tuple[Optional[torch.Tensor], torch.Tensor]:
+    ) -> Tuple[Optional[torch.Tensor], torch.Tensor]:
         mel = self.token2wav(
             code=codec_all,
             conditioning=conditioning,
@@ -1755,9 +1907,9 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         y_all: torch.Tensor,
         i: int,
         steps: int,
-        prev_generated: Union[torch.Tensor, list[torch.Tensor]],
+        prev_generated: Union[torch.Tensor, List[torch.Tensor]],
         finished: bool = False,
-    ) -> tuple[Union[torch.Tensor, list[torch.Tensor]], torch.Tensor]:
+    ) -> Tuple[Union[torch.Tensor, List[torch.Tensor]], torch.Tensor]:
         _mel, out = self.process_little_chunk(
             conditioning=conditioning,
             reference_mel=reference_mel,
@@ -1765,7 +1917,9 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
             y_all=y_all,
             i=i,
             steps=steps,
-            prev_generated=(prev_generated if isinstance(prev_generated, torch.Tensor) else None),
+            prev_generated=(
+                prev_generated if isinstance(prev_generated, torch.Tensor) else None
+            ),
             finished=finished,
         )
         return _mel if _mel is not None else prev_generated, out

vllm_omni/model_executor/models/qwen2_old.py

@@ -1,9 +1,9 @@
-from collections.abc import Iterable
-from typing import Optional, Union
+from typing import Iterable, Optional, Set, Tuple, Union
 
 import torch
 from torch import nn
 from transformers import Qwen2Config
+
 from vllm.attention import Attention, AttentionType
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -11,13 +11,23 @@ from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.pooler import Pooler, PoolingType
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
-from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead, VocabParallelEmbedding
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader, maybe_remap_kv_scale_name
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
 from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
@@ -38,6 +48,7 @@ logger = init_logger(__name__)
 
 
 class Qwen2MLP(nn.Module):
+
     def __init__(
         self,
         hidden_size: int,
@@ -62,7 +73,10 @@ class Qwen2MLP(nn.Module):
             prefix=f"{prefix}.down_proj",
         )
         if hidden_act != "silu":
-            raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported for now.")
+            raise ValueError(
+                f"Unsupported activation: {hidden_act}. "
+                "Only silu is supported for now."
+            )
         self.act_fn = SiluAndMul()
 
     def forward(self, x):
@@ -73,6 +87,7 @@ class Qwen2MLP(nn.Module):
 
 
 class Qwen2Attention(nn.Module):
+
     def __init__(
         self,
         hidden_size: int,
@@ -83,7 +98,7 @@ class Qwen2Attention(nn.Module):
         head_dim: Optional[int] = None,
         cache_config: Optional[CacheConfig] = None,
         quant_config: Optional[QuantizationConfig] = None,
-        rope_scaling: Optional[tuple] = None,
+        rope_scaling: Optional[Tuple] = None,
         prefix: str = "",
         attn_type: str = AttentionType.DECODER,
     ) -> None:
@@ -158,6 +173,7 @@ class Qwen2Attention(nn.Module):
 
 
 class Qwen2DecoderLayer(nn.Module):
+
     def __init__(
         self,
         config: Qwen2Config,
@@ -201,14 +217,16 @@ class Qwen2DecoderLayer(nn.Module):
             prefix=f"{prefix}.mlp",
         )
         self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
 
     def forward(
         self,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
         residual: Optional[torch.Tensor],
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         # Self Attention
         if residual is None:
             residual = hidden_states
@@ -237,6 +255,7 @@ class Qwen2DecoderLayer(nn.Module):
     }
 )
 class Qwen2Model(nn.Module):
+
     def __init__(
         self,
         *,
@@ -251,20 +270,27 @@ class Qwen2Model(nn.Module):
         quant_config = vllm_config.quant_config
 
         # TODO (@robertgshaw2): see if this can be moved out
-        if cache_config.sliding_window is not None and hasattr(config, "max_window_layers"):
+        if cache_config.sliding_window is not None and hasattr(
+            config, "max_window_layers"
+        ):
             raise ValueError(
                 "Sliding window for some but all layers is not "
                 "supported. This model uses sliding window "
-                f"but `max_window_layers` = {config.max_window_layers} is less than "
-                f"`num_hidden_layers` = {config.num_hidden_layers}. Please open an issue "
-                "to discuss this feature."
+                "but `max_window_layers` = {} is less than "
+                "`num_hidden_layers` = {}. Please open an issue "
+                "to discuss this feature.".format(
+                    config.max_window_layers,
+                    config.num_hidden_layers,
+                )
             )
 
         self.config = config
         self.quant_config = quant_config
         self.vocab_size = config.vocab_size
 
-        if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank):
+        if get_pp_group().is_first_rank or (
+            config.tie_word_embeddings and get_pp_group().is_last_rank
+        ):
             self.embed_tokens = VocabParallelEmbedding(
                 config.vocab_size,
                 getattr(config, "embedding_size", config.hidden_size),
@@ -322,11 +348,13 @@ class Qwen2Model(nn.Module):
                 residual,
             )
         if not get_pp_group().is_last_rank:
-            return IntermediateTensors({"hidden_states": hidden_states, "residual": residual})
+            return IntermediateTensors(
+                {"hidden_states": hidden_states, "residual": residual}
+            )
         hidden_states, _ = self.norm(hidden_states, residual)
         return hidden_states
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
             ("qkv_proj", "q_proj", "q"),
@@ -336,15 +364,19 @@ class Qwen2Model(nn.Module):
             ("gate_up_proj", "up_proj", 1),
         ]
         params_dict = dict(self.named_parameters(remove_duplicate=False))
-        loaded_params: set[str] = set()
+        loaded_params: Set[str] = set()
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
                 continue
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
+            if self.quant_config is not None and (
+                scale_name := self.quant_config.get_cache_scale(name)
+            ):
                 # Loading kv cache quantization scales
                 param = params_dict[scale_name]
                 weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
+                loaded_weight = (
+                    loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
+                )
                 weight_loader(param, loaded_weight)
                 loaded_params.add(scale_name)
                 continue
@@ -401,7 +433,9 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         self.lora_config = lora_config
 
         self.quant_config = quant_config
-        self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
+        self.model = Qwen2Model(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
 
         if get_pp_group().is_last_rank:
             if config.tie_word_embeddings:
@@ -419,7 +453,9 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         self.logits_processor = LogitsProcessor(config.vocab_size)
         self.sampler = Sampler()
 
-        self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
 
     def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.get_input_embeddings(input_ids)
@@ -431,7 +467,9 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         intermediate_tensors: Optional[IntermediateTensors] = None,
         inputs_embeds: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, IntermediateTensors]:
-        hidden_states = self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
+        hidden_states = self.model(
+            input_ids, positions, intermediate_tensors, inputs_embeds
+        )
         return hidden_states
 
     def compute_logits(
@@ -449,7 +487,7 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         next_tokens = self.sampler(logits, sampling_metadata)
         return next_tokens
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
         loader = AutoWeightsLoader(
             self,
             skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
@@ -483,7 +521,9 @@ class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
         self.lora_config = lora_config
 
         self.quant_config = quant_config
-        self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
+        self.model = Qwen2Model(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
 
         # TODO: Replace this model class with as_embedding_model(
         # Qwen2ForCausalLM) after changing the default pooling method
@@ -514,7 +554,9 @@ class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
     ) -> Optional[PoolerOutput]:
         return self._pooler(hidden_states, pooling_metadata)
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         weights = self.hf_to_vllm_mapper.apply(weights)
-        weights = ((name, data) for name, data in weights if not name.startswith("lm_head."))
+        weights = (
+            (name, data) for name, data in weights if not name.startswith("lm_head.")
+        )
         self.model.load_weights(weights)

vllm_omni/model_executor/models/registry.py

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

vllm_omni/model_executor/models/utils.py

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

vllm_omni/model_executor/models/vision.py

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

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -8,24 +8,15 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_cls: vllm_omni.worker.gpu.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true  # need to discuss
       trust_remote_code: true
       engine_output_type: latent  # change the param name,such as pooling_output
       enable_prefix_caching: false
-    is_comprehension: true
     final_output: true
     final_output_type: text
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1
   - stage_id: 1
     runtime:
       process: true
@@ -34,7 +25,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_cls: vllm_omni.worker.gpu.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -43,15 +34,6 @@ stage_args:
       engine_output_type: latent
     engine_input_source: [0]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
-    default_sampling_params:
-      temperature: 0.9
-      top_p: 0.8
-      top_k: 40
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-      stop_token_ids: [8294]
   - stage_id: 2
     runtime:
       process: true
@@ -60,7 +42,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
+      worker_cls: vllm_omni.worker.gpu.gpu_diffusion_worker.GPUDiffusionWorker
       scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true
@@ -70,14 +52,6 @@ stage_args:
     engine_input_source: [1]
     final_output: true
     final_output_type: audio
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1
 
 # Top-level runtime config (concise): default windows and stage edges
 runtime:

vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py

@@ -1,20 +1,14 @@
 from typing import Union
 
 import torch
+
 from vllm.inputs import TextPrompt
-
 from vllm_omni.inputs.data import OmniTokensPrompt
-
 TALKER_CODEC_PAD_TOKEN_ID = 8292
 TALKER_CODEC_START_TOKEN_ID = 8293
 TALKER_CODEC_END_TOKEN_ID = 8294
-
-
 def thinker2talker(
-    stage_list,
-    engine_input_source,
-    prompt: Union[OmniTokensPrompt, TextPrompt] = None,
-    requires_multimodal_data: bool = False,
+    stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None
 ):
     if not engine_input_source:
         raise ValueError("engine_input_source cannot be empty")
@@ -25,10 +19,9 @@ def thinker2talker(
         raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
     thinker_outputs = stage_list[source_stage_id].engine_outputs
     talker_inputs = []
-    if not isinstance(prompt, list):
-        prompt = [prompt]
     multi_modal_data = {
-        thinker_output.request_id: p.get("multi_modal_data", None) for thinker_output, p in zip(thinker_outputs, prompt)
+        thinker_output.request_id: prompt.get("multi_modal_data", None)
+        for thinker_output, prompt in zip(thinker_outputs, prompt)
     }
 
     for i, thinker_output in enumerate(thinker_outputs):
@@ -36,25 +29,27 @@ def thinker2talker(
         prompt_token_ids = thinker_output.prompt_token_ids
         thinker_output_ids = output.token_ids
         prompt_token_ids_len = len(prompt_token_ids)
-        latent = output.multimodal_output["latent"]
-        thinker_hidden_states = latent.clone().detach().to(latent.device)
-        additional_information = {
-            "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to(torch.float32),
-            "prompt_embeds": thinker_hidden_states[:prompt_token_ids_len].to(torch.float32),
-            "prompt_token_ids": prompt_token_ids,
-            "thinker_output_token_ids": thinker_output_ids,
-            "thinker_result_shape": list(thinker_hidden_states[prompt_token_ids_len:].shape),
-            "prompt_embeds_shape": list(thinker_hidden_states[:prompt_token_ids_len].shape),
-        }
+        thinker_hidden_states = (
+            output.multimodal_output["latent"].clone().detach().cuda()
+        )
         talker_inputs.append(
             OmniTokensPrompt(
-                prompt_token_ids=[TALKER_CODEC_START_TOKEN_ID]
-                + [TALKER_CODEC_PAD_TOKEN_ID] * (len(prompt_token_ids))
-                + [TALKER_CODEC_END_TOKEN_ID],
-                additional_information=additional_information,
+                prompt_token_ids=[TALKER_CODEC_START_TOKEN_ID] + [TALKER_CODEC_PAD_TOKEN_ID] * (len(prompt_token_ids)) + [TALKER_CODEC_END_TOKEN_ID],
+                additional_information={
+                    "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to(
+                        torch.float32
+                    ),
+                    "prompt_embeds": thinker_hidden_states[:prompt_token_ids_len].to(
+                        torch.float32
+                    ),
+                    "prompt_token_ids": prompt_token_ids,
+                    "thinker_output_token_ids": thinker_output_ids,
+                    "thinker_result_shape": list(thinker_hidden_states[prompt_token_ids_len:].shape),
+                    "prompt_embeds_shape": list(thinker_hidden_states[:prompt_token_ids_len].shape),
+                },
                 multi_modal_data=(
                     multi_modal_data[thinker_output.request_id]
-                    if requires_multimodal_data and multi_modal_data is not None
+                    if multi_modal_data is not None
                     else None
                 ),
                 mm_processor_kwargs=None,

vllm_omni/outputs.py

@@ -2,6 +2,7 @@ from dataclasses import dataclass
 from typing import Optional
 
 import torch
+
 from vllm.outputs import RequestOutput
 from vllm.v1.outputs import ModelRunnerOutput
 

vllm_omni/patch.py

@@ -1,10 +1,11 @@
 import sys
 
 from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt
-from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding as _OriginalMRotaryEmbedding
+from vllm.model_executor.layers.rotary_embedding import (
+    MRotaryEmbedding as _OriginalMRotaryEmbedding,
+)
 from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest
 from vllm.v1.request import Request as _OriginalRequest
-
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
@@ -13,9 +14,15 @@ from vllm_omni.request import OmniRequest
 for module_name, module in sys.modules.items():
     if hasattr(module, "TokensPrompt") and module.TokensPrompt == _OriginalTokensPrompt:
         module.TokensPrompt = OmniTokensPrompt
-    if hasattr(module, "MRotaryEmbedding") and module.MRotaryEmbedding == _OriginalMRotaryEmbedding:
+    if (
+        hasattr(module, "MRotaryEmbedding")
+        and module.MRotaryEmbedding == _OriginalMRotaryEmbedding
+    ):
         module.MRotaryEmbedding = MRotaryEmbedding
     if hasattr(module, "Request") and module.Request == _OriginalRequest:
         module.Request = OmniRequest
-    if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest:
+    if (
+        hasattr(module, "EngineCoreRequest")
+        and module.EngineCoreRequest == _OriginalEngineCoreRequest
+    ):
         module.EngineCoreRequest = OmniEngineCoreRequest

vllm_omni/request.py

@@ -6,7 +6,11 @@ from vllm.v1.structured_output.request import StructuredOutputRequest
 if TYPE_CHECKING:
     from vllm.v1.core.kv_cache_utils import BlockHash
 
-from vllm_omni.engine import AdditionalInformationPayload, OmniEngineCoreRequest, PromptEmbedsPayload
+from vllm_omni.engine import (
+    AdditionalInformationPayload,
+    OmniEngineCoreRequest,
+    PromptEmbedsPayload,
+)
 
 
 class OmniRequest(Request):
@@ -21,7 +25,9 @@ class OmniRequest(Request):
         # Serialized prompt embeddings payload (optional)
         self.prompt_embeds: Optional[PromptEmbedsPayload] = prompt_embeds
         # Serialized additional information payload (optional)
-        self.additional_information: Optional[AdditionalInformationPayload] = additional_information
+        self.additional_information: Optional[AdditionalInformationPayload] = (
+            additional_information
+        )
 
     @classmethod
     def from_engine_core_request(
@@ -29,6 +35,7 @@ class OmniRequest(Request):
         request: OmniEngineCoreRequest,
         block_hasher: Optional[Callable[["Request"], list["BlockHash"]]],
     ) -> "Request":
+
         return cls(
             request_id=request.request_id,
             client_index=request.client_index,
@@ -40,7 +47,9 @@ class OmniRequest(Request):
             arrival_time=request.arrival_time,
             lora_request=request.lora_request,
             structured_output_request=(
-                StructuredOutputRequest(sampling_params=request.sampling_params) if request.sampling_params else None
+                StructuredOutputRequest(sampling_params=request.sampling_params)
+                if request.sampling_params
+                else None
             ),
             cache_salt=request.cache_salt,
             priority=request.priority,

vllm_omni/version.py

@@ -1,12 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-try:
-    from ._version import __version__, __version_tuple__
-except Exception as e:
-    import warnings
-
-    warnings.warn(f"Failed to read commit hash:\n{e}", RuntimeWarning, stacklevel=2)
-
-    __version__ = "dev"
-    __version_tuple__ = (0, 0, __version__)

vllm_omni/worker/cpu/cpu_ar_model_runner.py

@@ -0,0 +1,649 @@
+"""AR CPU Model Runner for vLLM-omni.
+
+Exposes per-request hidden representations via ModelRunnerOutput.pooler_output
+and also outputs sampled tokens.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional, Union
+
+import numpy as np
+import torch
+
+from vllm.forward_context import BatchDescriptor
+from vllm.logger import init_logger
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import AsyncModelRunnerOutput
+from vllm.v1.structured_output.utils import apply_grammar_bitmask
+from vllm.v1.utils import record_function_or_nullcontext
+from vllm.v1.worker.gpu_model_runner import (
+    EMPTY_MODEL_RUNNER_OUTPUT,
+    AsyncGPUModelRunnerOutput,
+    IntermediateTensors,
+    get_pp_group,
+    get_tp_group,
+    has_kv_transfer_group,
+    set_forward_context,
+)
+from vllm.v1.worker.ubatch_utils import UBatchSlices
+from vllm.v1.worker.utils import is_residual_scattered_for_sp
+from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
+from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker.cpu.cpu_model_runner import OmniCPUModelRunner
+
+logger = init_logger(__name__)
+
+
+class CPUARModelRunner(OmniCPUModelRunner):
+    """Autoregressive GPU model runner that returns hidden states per request.
+
+    This runner follows the same preparation and forward path as GPUModelRunner
+    (inputs assembly, multi-modal handling, TP/PP/DP integration, CUDA graphs),
+    and additionally performs lightweight sampling so that sampled tokens are
+    available in outputs. Hidden representations are taken at the same indices
+    that GPUModelRunner would use for sampling/logits (i.e. `logits_indices`).
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
+
+    def _preprocess(
+        self,
+        scheduler_output: "SchedulerOutput",
+        num_scheduled_tokens_np: np.ndarray,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        ubatch_slices: Optional[UBatchSlices] = None,
+        num_tokens_after_padding: Optional[torch.Tensor] = None,
+    ) -> tuple[
+        int,
+        int,
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        torch.Tensor,
+        Optional[IntermediateTensors],
+        dict[str, Any],
+        Optional[dict[str, dict]],
+    ]:
+
+        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+        if ubatch_slices:
+            assert num_tokens_after_padding is not None
+            num_input_tokens = int(num_tokens_after_padding[0].item() * 2)
+            self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
+        elif ubatch_slices is None:
+            num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens)
+            num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
+            num_input_tokens += num_pad
+
+        # _prepare_inputs may reorder the batch, so we must gather multi
+        # modal outputs after that to ensure the correct order
+        per_req_additional_information: Optional[dict[str, dict]] = None
+        if (
+            self.supports_mm_inputs
+            and get_pp_group().is_first_rank
+            and not self.model_config.is_encoder_decoder
+        ):
+            # Run the multimodal encoder if any.
+            self._execute_mm_encoder(scheduler_output)
+            mm_embeds = self._gather_mm_embeddings(scheduler_output)
+
+            # NOTE(woosuk): To unify token ids and soft tokens (vision
+            # embeddings), we always use embeddings (rather than token ids)
+            # as input to the multimodal model, even when the input is text.
+            inputs_embeds_scheduled = self.model.get_input_embeddings(
+                input_ids=self.input_ids.gpu[:num_scheduled_tokens],
+                multimodal_embeddings=mm_embeds or None,
+            )
+
+            # TODO(woosuk): Avoid the copy. Optimize.
+            self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
+
+            # Reset per-step additional information collector (deprecated concat path)
+            if hasattr(self, "_forward_additional_information"):
+                self._forward_additional_information = None
+            # New: per-request additional information for this step
+            per_req_additional_information = {}
+
+            # Overlay custom prompt_embeds per request for the prompt portion;
+            # collect additional_information (tensor/list) for prefill portion only
+            for req_index, req_id in enumerate(self.input_batch.req_ids):
+                req_state = self.requests[req_id]
+                pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
+                addi_cpu = getattr(req_state, "additional_information_cpu", None)
+                num_computed_tokens = int(
+                    self.input_batch.num_computed_tokens_cpu[req_index]
+                )
+                prompt_len = len(req_state.prompt_token_ids)
+                prompt_remaining = max(0, prompt_len - num_computed_tokens)
+                sched_tokens = int(num_scheduled_tokens_np[req_index])
+                overlay_len = min(sched_tokens, prompt_remaining)
+                if overlay_len <= 0:
+                    continue
+                if pe_cpu is not None:
+                    src = pe_cpu[
+                        num_computed_tokens : num_computed_tokens + overlay_len
+                    ].to(dtype=self.dtype, device=self.device, non_blocking=True)
+                    start_offset = int(self.query_start_loc.cpu[req_index])
+                    self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(
+                        src
+                    )
+                # Build per-request additional information (no cross-request concat)
+                if addi_cpu is not None and isinstance(addi_cpu, dict):
+                    req_info: dict[str, object] = {}
+                    for k, v in addi_cpu.items():
+                        if isinstance(v, torch.Tensor):
+                            # For prefill tokens, pass only the scheduled slice; for decode or no scheduled tokens, pass whole tensor
+                            if overlay_len > 0:
+                                try:
+                                    seg = v[
+                                        num_computed_tokens : num_computed_tokens + overlay_len
+                                    ].detach().to("cpu").contiguous()
+                                except Exception:
+                                    seg = v.detach().to("cpu").contiguous()
+                                req_info[k] = seg
+                            else:
+                                req_info[k] = v.detach().to("cpu").contiguous()
+                        elif isinstance(v, list):
+                            req_info[k] = v
+                        else:
+                            req_info[k] = v
+                    per_req_additional_information[req_id] = req_info
+
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = {
+                **self._init_model_kwargs(num_scheduled_tokens),
+                **self._extract_mm_kwargs(scheduler_output),
+            }
+        elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
+            # Get the input embeddings for the tokens that are not input embeds,
+            # then put them into the appropriate positions.
+            # TODO(qthequartermasterman): Since even when prompt embeds are
+            # enabled, (a) not all requests will use prompt embeds, and (b)
+            # after the initial prompt is processed, the rest of the generated
+            # tokens will be token ids, it is not desirable to have the
+            # embedding layer outside of the CUDA graph all the time. The v0
+            # engine avoids this by "double compiling" the CUDA graph, once
+            # with input_ids and again with inputs_embeds, for all num_tokens.
+            # If a batch only has token ids, then including the embedding layer
+            # in the CUDA graph will be more performant (like in the else case
+            # below).
+            token_ids_idx = (
+                self.is_token_ids.gpu[:num_scheduled_tokens]
+                .nonzero(as_tuple=False)
+                .squeeze(1)
+            )
+            # Some tokens ids may need to become embeds
+            if token_ids_idx.numel() > 0:
+                token_ids = self.input_ids.gpu[token_ids_idx]
+                tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
+                self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
+
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+            input_ids = None
+        else:
+            # For text-only models, we use token ids as input.
+            # While it is possible to use embeddings as input just like the
+            # multimodal models, it is not desirable for performance since
+            # then the embedding layer is not included in the CUDA graph.
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = None
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+        if self.uses_mrope:
+            positions = self.mrope_positions.gpu[:, :num_input_tokens]
+        else:
+            positions = self.positions.gpu[:num_input_tokens]
+
+        if get_pp_group().is_first_rank:
+            intermediate_tensors = None
+        else:
+            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                num_input_tokens, intermediate_tensors, True
+            )
+
+        if (
+            self.model_config.is_encoder_decoder
+            and scheduler_output.scheduled_encoder_inputs
+        ):
+            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
+            model_kwargs.update(encoder_inputs)
+
+        return (
+            num_scheduled_tokens,
+            num_input_tokens,
+            num_tokens_after_padding,
+            input_ids,
+            inputs_embeds,
+            positions,
+            intermediate_tensors,
+            model_kwargs,
+            per_req_additional_information,
+        )
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: "SchedulerOutput",
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+    ) -> Union[OmniModelRunnerOutput, AsyncModelRunnerOutput, IntermediateTensors]:
+        with record_function_or_nullcontext("Preprocess"):
+            with self.synchronize_input_prep():
+                # Update persistent batch states.
+                self._update_states(scheduler_output)
+
+                # Decode per-request prompt_embeds / additional_hidden_states payloads
+                # (if present) into CPU tensors
+                try:
+                    new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
+                    if new_reqs:
+                        for nr in new_reqs:
+                            req_id = getattr(nr, "req_id", None) or getattr(
+                                nr, "request_id", None
+                            )
+                            if req_id is None:
+                                continue
+                            # prompt_embeds
+                            payload_pe = getattr(nr, "prompt_embeds", None)
+                            if payload_pe is not None:
+                                if isinstance(payload_pe, torch.Tensor):
+                                    pe_cpu = payload_pe.detach().to("cpu").contiguous()
+                                elif isinstance(payload_pe, PromptEmbedsPayload):
+                                    dt = np.dtype(
+                                        getattr(payload_pe, "dtype", "float32")
+                                    )
+                                    arr = np.frombuffer(payload_pe.data, dtype=dt)
+                                    arr = arr.reshape(payload_pe.shape)
+                                    pe_cpu = torch.from_numpy(arr)
+                                else:
+                                    pe_cpu = None
+                                if pe_cpu is not None and req_id in self.requests:
+                                    setattr(
+                                        self.requests[req_id],
+                                        "prompt_embeds_cpu",
+                                        pe_cpu,
+                                    )
+                            # additional_information
+                            payload_info = getattr(nr, "additional_information", None)
+                            if payload_info is not None:
+                                info_dict = {}
+                                if isinstance(payload_info, dict):
+                                    # Already decoded
+                                    info_dict = payload_info
+                                elif isinstance(
+                                    payload_info, AdditionalInformationPayload
+                                ):
+                                    for k, entry in payload_info.entries.items():
+                                        if entry.tensor_data is not None:
+                                            dt = np.dtype(
+                                                getattr(
+                                                    entry, "tensor_dtype", "float32"
+                                                )
+                                            )
+                                            arr = np.frombuffer(
+                                                entry.tensor_data, dtype=dt
+                                            )
+                                            arr = arr.reshape(entry.tensor_shape)
+                                            info_dict[k] = torch.from_numpy(arr)
+                                        else:
+                                            info_dict[k] = entry.list_data
+                                if info_dict and req_id in self.requests:
+                                    setattr(
+                                        self.requests[req_id],
+                                        "additional_information_cpu",
+                                        info_dict,
+                                    )
+                except Exception as e:
+                    logger.error(
+                        f"Error decoding prompt_embeds / additional_information: {e}"
+                    )
+                    pass
+
+                if not scheduler_output.total_num_scheduled_tokens:
+                    if not has_kv_transfer_group():
+                        # Return empty ModelRunnerOutput if no work to do.
+                        return EMPTY_MODEL_RUNNER_OUTPUT
+                    return self.kv_connector_no_forward(
+                        scheduler_output, self.vllm_config
+                    )
+                if self.cache_config.kv_sharing_fast_prefill:
+                    assert not self.input_batch.num_prompt_logprobs, (
+                        "--kv-sharing-fast-prefill produces incorrect "
+                        "logprobs for prompt tokens, tokens, please disable "
+                        "it when the requests need prompt logprobs"
+                    )
+
+                # Prepare the decoder inputs.
+                (
+                    attn_metadata,
+                    logits_indices,
+                    spec_decode_metadata,
+                    num_scheduled_tokens_np,
+                    spec_decode_common_attn_metadata,
+                    max_query_len,
+                    ubatch_slices,
+                    num_tokens_after_padding,
+                ) = self._prepare_inputs(scheduler_output)
+
+            (
+                num_scheduled_tokens,
+                num_input_tokens,
+                num_tokens_across_dp,
+                input_ids,
+                inputs_embeds,
+                positions,
+                intermediate_tensors,
+                model_kwargs,
+                per_req_additional_information,
+            ) = self._preprocess(
+                scheduler_output,
+                num_scheduled_tokens_np,
+                intermediate_tensors,
+                ubatch_slices,
+                num_tokens_after_padding,
+            )
+
+            uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
+                num_scheduled_tokens == self.input_batch.num_reqs * max_query_len
+            )
+            batch_descriptor = BatchDescriptor(
+                num_tokens=num_input_tokens, uniform_decode=uniform_decode
+            )
+            cudagraph_runtime_mode, batch_descriptor = (
+                self.cudagraph_dispatcher.dispatch(batch_descriptor)
+            )
+
+        # This is currently to get around the assert in the DPMetadata
+        # where it wants `num_tokens_across_dp` to align with `num_tokens`
+        if ubatch_slices is not None:
+            num_input_tokens = ubatch_slices[0].num_tokens
+
+        # Run the model.
+        # Use persistent buffers for CUDA graphs.
+        with (
+            set_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_input_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+                cudagraph_runtime_mode=cudagraph_runtime_mode,
+                batch_descriptor=batch_descriptor,
+                ubatch_slices=ubatch_slices,
+            ),
+            record_function_or_nullcontext("Forward"),
+            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
+        ):
+
+            model_kwargs_extra = {}
+            # Pass per-request additional information map for this step (no concat)
+            if per_req_additional_information:
+                model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information
+            # Always pass per-request runtime additional_information (persisted in request state)
+            try:
+                per_req_runtime_info = []
+                for req_id in self.input_batch.req_ids:
+                    req_state = self.requests.get(req_id)
+                    info = (
+                        getattr(req_state, "additional_information_cpu", None)
+                        if req_state is not None
+                        else None
+                    )
+                    per_req_runtime_info.append(info if isinstance(info, dict) else {})
+                model_kwargs_extra["runtime_additional_information"] = (
+                    per_req_runtime_info
+                )
+                model_kwargs_extra["request_ids"] = self.input_batch.req_ids
+                # Pass each request's token span within the flattened sequence for this step, enabling the model to map decode/prefill by request
+                req_token_spans = []
+                for req_index in range(len(self.input_batch.req_ids)):
+                    start_offset = int(self.query_start_loc.cpu[req_index])
+                    sched_tokens = int(num_scheduled_tokens_np[req_index])
+                    req_token_spans.append((start_offset, start_offset + sched_tokens))
+                model_kwargs_extra["request_token_spans"] = req_token_spans
+            except Exception:
+                pass
+            model_output = self.model(
+                input_ids=input_ids,
+                positions=positions,
+                intermediate_tensors=intermediate_tensors,
+                inputs_embeds=inputs_embeds,
+                **model_kwargs,
+                sampling_metadata=self.input_batch.sampling_metadata,
+                logits_index=logits_indices,
+                sampler=self.sampler,
+                **model_kwargs_extra,
+            )
+
+        with record_function_or_nullcontext("Postprocess"):
+            if self.use_aux_hidden_state_outputs:
+                # True when EAGLE 3 is used.
+                hidden_states, aux_hidden_states = model_output
+            else:
+                # Common case.
+                hidden_states = model_output
+                aux_hidden_states = None
+
+            hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
+                hidden_states
+            )
+            # The model side may return per-request additional_information updates (model-agnostic channel).
+            # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order;
+            # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode.
+            try:
+                if (
+                    isinstance(multimodal_outputs, dict)
+                    and (
+                        "additional_information_update" in multimodal_outputs
+                        or "additional_information_update_by_req_id" in multimodal_outputs
+                    )
+                ):
+                    # Option A: list[dict] in batch order
+                    updates_list = multimodal_outputs.get(
+                        "additional_information_update"
+                    )
+                    if isinstance(updates_list, list):
+                        for idx, upd in enumerate(updates_list):
+                            if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids):
+                                continue
+                            req_id = self.input_batch.req_ids[idx]
+                            self._merge_additional_information_update(req_id, upd)
+                    # Option B: dict[str, dict] keyed by req_id
+                    updates_map = multimodal_outputs.get(
+                        "additional_information_update_by_req_id"
+                    )
+                    if isinstance(updates_map, dict):
+                        for req_id, upd in updates_map.items():
+                            if not isinstance(upd, dict):
+                                continue
+                            if req_id not in self.requests:
+                                continue
+                            self._merge_additional_information_update(req_id, upd)
+            except Exception as e:
+                logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}, with the multimodal_outputs as {multimodal_outputs}")
+            if not self.broadcast_pp_output:
+
+                # Common case.
+                if not get_pp_group().is_last_rank:
+                    # Return the intermediate tensors.
+                    assert isinstance(hidden_states, IntermediateTensors)
+                    hidden_states.kv_connector_output = kv_connector_output
+                    return hidden_states
+
+                if self.is_pooling_model:
+                    # Return the pooling output.
+                    output = self._pool(
+                        hidden_states, num_scheduled_tokens, num_scheduled_tokens_np
+                    )
+                    output.kv_connector_output = kv_connector_output
+                    return output
+
+                sample_hidden_states = hidden_states[logits_indices]
+                logits = self.model.compute_logits(sample_hidden_states)
+            else:
+                # Rare case.
+                assert not self.is_pooling_model
+
+                if not get_pp_group().is_last_rank:
+                    all_gather_tensors = {
+                        "residual": not is_residual_scattered_for_sp(
+                            self.vllm_config, num_input_tokens
+                        )
+                    }
+                    get_pp_group().send_tensor_dict(
+                        hidden_states.tensors,
+                        all_gather_group=get_tp_group(),
+                        all_gather_tensors=all_gather_tensors,
+                    )
+                    logits = None
+                else:
+                    sample_hidden_states = hidden_states[logits_indices]
+                    logits = self.model.compute_logits(sample_hidden_states)
+
+                model_output_broadcast_data = {}
+                if logits is not None:
+                    model_output_broadcast_data["logits"] = logits.contiguous()
+
+                model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
+                    model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
+                )
+                assert model_output_broadcast_data is not None
+                logits = model_output_broadcast_data["logits"]
+
+            # Apply structured output bitmasks if present
+            if scheduler_output.grammar_bitmask is not None:
+                apply_grammar_bitmask(
+                    scheduler_output, self.input_batch, logits, self.device
+                )
+
+        with record_function_or_nullcontext("Sample"):
+            sampler_output = self._sample(logits, spec_decode_metadata)
+
+        def propose_draft_token_ids(sampled_token_ids):
+            assert spec_decode_common_attn_metadata is not None
+            with record_function_or_nullcontext("Draft"):
+                self._draft_token_ids = self.propose_draft_token_ids(
+                    scheduler_output,
+                    sampled_token_ids,
+                    self.input_batch.sampling_metadata,
+                    hidden_states,
+                    sample_hidden_states,
+                    aux_hidden_states,
+                    spec_decode_metadata,
+                    spec_decode_common_attn_metadata,
+                )
+
+        use_padded_batch_for_eagle = (
+            self.speculative_config
+            and self.speculative_config.use_eagle()
+            and not self.speculative_config.disable_padded_drafter_batch
+        )
+        effective_drafter_max_model_len = self.max_model_len
+        if effective_drafter_max_model_len is None:
+            effective_drafter_max_model_len = self.model_config.max_model_len
+        if (
+            self.speculative_config
+            and self.speculative_config.draft_model_config is not None
+            and self.speculative_config.draft_model_config.max_model_len is not None
+        ):
+            effective_drafter_max_model_len = (
+                self.speculative_config.draft_model_config.max_model_len
+            )
+        input_fits_in_drafter = spec_decode_common_attn_metadata and (
+            spec_decode_common_attn_metadata.seq_lens.max()
+            + self.speculative_config.num_speculative_tokens
+            <= effective_drafter_max_model_len
+        )
+        if use_padded_batch_for_eagle and input_fits_in_drafter:
+            # EAGLE speculative decoding can use the GPU sampled tokens
+            # as inputs, and does not need to wait for bookkeeping to finish.
+            propose_draft_token_ids(sampler_output.sampled_token_ids)
+
+        with record_function_or_nullcontext("Bookkeep"):
+            (
+                num_nans_in_logits,
+                logprobs_lists,
+                valid_sampled_token_ids,
+                prompt_logprobs_dict,
+                req_ids_output_copy,
+                req_id_to_index_output_copy,
+                invalid_req_indices,
+            ) = self._bookkeeping_sync(
+                scheduler_output,
+                sampler_output,
+                logits,
+                hidden_states,
+                num_scheduled_tokens,
+            )
+
+        if (
+            self.speculative_config
+            and not use_padded_batch_for_eagle
+            and input_fits_in_drafter
+        ):
+            # ngram and other speculative decoding methods use the sampled
+            # tokens on the CPU, so they are run after bookkeeping.
+            propose_draft_token_ids(valid_sampled_token_ids)
+
+        with record_function_or_nullcontext("EPLB"):
+            self.eplb_step()
+
+        # Convert to per-request tensors on CPU
+        hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
+        pooler_output: list[Optional[torch.Tensor]] = []
+        prev_logits_index = 0
+        for logits_index in logits_indices:
+            pooler_output.append(
+                hidden_states_cpu[prev_logits_index : logits_index + 1]
+            )
+            prev_logits_index = logits_index + 1
+
+        output = OmniModelRunnerOutput(
+            req_ids=req_ids_output_copy,
+            req_id_to_index=req_id_to_index_output_copy,
+            sampled_token_ids=valid_sampled_token_ids,
+            logprobs=logprobs_lists,
+            prompt_logprobs_dict=prompt_logprobs_dict,
+            pooler_output=(
+                pooler_output
+                if self.vllm_config.model_config.engine_output_type != "text"
+                else None
+            ),
+            kv_connector_output=kv_connector_output,
+            num_nans_in_logits=num_nans_in_logits,
+        )
+
+        if not self.use_async_scheduling:
+            return output
+
+        return AsyncGPUModelRunnerOutput(
+            model_runner_output=output,
+            sampled_token_ids=sampler_output.sampled_token_ids,
+            invalid_req_indices=invalid_req_indices,
+            async_output_copy_stream=self.async_output_copy_stream,
+        )
+
+    def _merge_additional_information_update(self, req_id: str, upd: dict) -> None:
+        req_state = self.requests.get(req_id)
+        if req_state is None:
+            return
+        existing = getattr(req_state, "additional_information_cpu", {})
+        if not isinstance(existing, dict):
+            existing = {}
+        merged = dict(existing)
+        for k, v in upd.items():
+            if isinstance(v, torch.Tensor):
+                merged[k] = v.detach().to("cpu").contiguous()
+            elif isinstance(v, list):
+                new_list = []
+                for item in v:
+                    if isinstance(item, torch.Tensor):
+                        new_list.append(item.detach().to("cpu").contiguous())
+                    else:
+                        new_list.append(item)
+                merged[k] = new_list
+            else:
+                merged[k] = v
+        setattr(req_state, "additional_information_cpu", merged)