update

Signed-off-by: Roger Wang <hey@rogerw.io>

.buildkite/pipeline.yml

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

.buildkite/scripts/run_qwen_offline_example.sh

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

.github/CODEOWNERS

@@ -0,0 +1,17 @@
+# 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/hsliuustc0106/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/vllm-project/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/hsliuustc0106/vllm-omni/main/vllm_omni/collect_env.py
+      wget https://raw.githubusercontent.com/vllm-project/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/hsliuustc0106/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/vllm-project/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/hsliuustc0106/vllm-omni/main/vllm_omni/collect_env.py
+      wget https://raw.githubusercontent.com/vllm-project/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/hsliuustc0106/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/vllm-project/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/hsliuustc0106/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/vllm-project/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/hsliuustc0106/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/vllm-project/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/hsliuustc0106/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/vllm-project/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/hsliuustc0106/vllm-omni/main/vllm_omni/collect_env.py
+      wget https://raw.githubusercontent.com/vllm-project/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/hsliuustc0106/vllm-omni/issues?q=label%3ARFC+sort%3Aupdated-desc) for reference.
+      #### Please take a look at previous [RFCs](https://github.com/vllm-project/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/hsliuustc0106/vllm-omni/discussions
+    url: https://github.com/vllm-project/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/hsliuustc0106/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)
+**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)

.github/scripts/detect_changed_tests.sh

@@ -0,0 +1,89 @@
+#!/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

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

.readthedocs.yml

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

docs/DOCS_GUIDE.md

@@ -0,0 +1,139 @@
+# 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/overview.md

@@ -0,0 +1,68 @@
+# 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/contributing/README.md

@@ -0,0 +1,170 @@
+# 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/api_design_doc.md

@@ -1,13 +1,10 @@
-# vLLM-omni API Documentation
+# API Design Documentation
 
-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.
+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.
 
-## 📋 API Design Templates
+## 📋 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 API
 **Core module** provides fundamental scheduling, caching, and resource management functionality.
 
 **Key Components:**
@@ -16,7 +13,7 @@ This directory contains comprehensive API design documentation for all core modu
 - Resource allocation and coordination
 - Inter-module communication
 
-### 3. [Engine Module API](engine_module_api.md)
+### Engine Module API
 **Engine module** handles model loading, inference execution, and output processing.
 
 **Key Components:**
@@ -25,7 +22,7 @@ This directory contains comprehensive API design documentation for all core modu
 - Input preprocessing for various modalities
 - Output postprocessing and formatting
 
-### 4. [Executor Module API](executor_module_api.md)
+### Executor Module API
 **Executor module** coordinates and manages request execution across different engines and workers.
 
 **Key Components:**
@@ -34,7 +31,7 @@ This directory contains comprehensive API design documentation for all core modu
 - Worker management and task distribution
 - Error handling and recovery
 
-### 5. [Worker Module API](worker_module_api.md)
+### Worker Module API
 **Worker module** provides the actual execution environment for model inference.
 
 **Key Components:**
@@ -43,79 +40,6 @@ This directory contains comprehensive API design documentation for all core modu
 - 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
 
@@ -155,11 +79,11 @@ Each API template follows this standardized structure:
 2. **Set Up Development Environment**
    ```bash
    # Clone the repository
-   git clone https://github.com/hsliuustc0106/vllm-omni.git
+   git clone https://github.com/vllm-project/vllm-omni.git
    cd vllm-omni
 
    # Install dependencies
-   pip install -r requirements-dev.txt
+   pip install -e ".[dev]"
 
    # Set up pre-commit hooks
    pre-commit install
@@ -189,28 +113,12 @@ Each API template follows this standardized structure:
    - 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 master template as a guide
+   - Use the [api template](https://github.com/vllm-project/vllm-omni/docs/contributing/design_documents/api_design_template.md) as a guide
    - Maintain consistency across modules
    - Include all required sections
 

docs/contributing/design_documents/architecture_overview.md

@@ -1,4 +1,4 @@
-# vLLM-omni Architecture Design Document
+# Architecture Overview
 
 # 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: [hsliuustc@gmail.com](mailto:hsliuustc@gmail.com)
+Hongsheng Liu: [vllm-project@gmail.com](mailto:vllm-project@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/contributing/design_documents/index.md

@@ -0,0 +1,12 @@
+# 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/contributing/design_documents/vllm_omni_design.md

@@ -1,4 +1,4 @@
-# vLLM-omni Software Design Document
+# Architecture Design Principles
 
 ## 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 improt AsyncLLM
+from vllm.v1.engine.async_llm import 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)
-        negtive_prompt_embeds = self.pipe.embed_prompt(" ")
+        negative_prompt_embeds = self.pipe.embed_prompt(" ")
 
         width, height = aspect_ratios["16:9"]
 
         image = pipe(
             prompt_embeds=prompt_embeds,
-            negtive_prompt_embeds=negtive_prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
             width=width,
             height=height,
             num_inference_steps=50,

docs/index.md

@@ -0,0 +1,32 @@
+# 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

@@ -0,0 +1,233 @@
+# 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

@@ -0,0 +1,55 @@
+/* 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

@@ -0,0 +1,8 @@
+# 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

@@ -0,0 +1,64 @@
+# 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/offline_inference/qwen_2_5_omni.md

@@ -1,5 +1,8 @@
 # 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
@@ -49,7 +52,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 successfull.
+Just run the command below. If no error, it demonstrates that the installation is successful.
 ```bash
 python -c "import vllm._C"
 ```
@@ -59,34 +62,36 @@ Install additional requirements for vllm-omni
 ```bash
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm_omni
-uv pip install -r requirements.txt
+uv pip install -e .
 ```
 
 ## 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
 ```
-Modify PYTHONPATH in run_multiple_prompts.sh as your path of vllm_omni. Then run.
+Then run
 ```bash
-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
+bash run.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"
+    ``````

docs/user_guide/examples/online_serving/README.md

@@ -0,0 +1,35 @@
+# 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

@@ -0,0 +1,5 @@
+# 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

@@ -0,0 +1,7 @@
+# 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

@@ -0,0 +1,55 @@
+# 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

@@ -0,0 +1,42 @@
+# 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

@@ -0,0 +1,307 @@
+# 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/extract_prompts.py

@@ -1,30 +1,34 @@
 #!/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, "r", encoding="utf-8", errors="ignore") as f:
+    with open(args.input, encoding="utf-8", errors="ignore") as f:
         for line in f:
             if len(prompts) >= args.topk:
                 break
@@ -39,5 +43,3 @@ def main() -> None:
 
 if __name__ == "__main__":
     main()
-
-

examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh

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

examples/offline_inference/qwen2_5_omni/run_single_prompt.sh

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

examples/offline_inference/qwen_2_5_omni/end2end.py

@@ -1,253 +0,0 @@
-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/processing_omni.py

@@ -1,391 +0,0 @@
-from __future__ import annotations
-
-import base64
-import logging
-import math
-import os
-import time
-import warnings
-from functools import lru_cache
-from io import BytesIO
-
-import requests
-import torch
-import torchvision
-from packaging import version
-from PIL import Image
-from torchvision import io, transforms
-from torchvision.transforms import InterpolationMode
-
-logger = logging.getLogger(__name__)
-
-IMAGE_FACTOR = 28
-MIN_PIXELS = 4 * 28 * 28
-MAX_PIXELS = 16384 * 28 * 28
-MAX_RATIO = 200
-
-VIDEO_MIN_PIXELS = 128 * 28 * 28
-VIDEO_MAX_PIXELS = 768 * 28 * 28
-VIDEO_TOTAL_PIXELS = 24576 * 28 * 28
-FRAME_FACTOR = 2
-FPS = 2.0
-FPS_MIN_FRAMES = 4
-FPS_MAX_FRAMES = 768
-
-temporal_patch_size = 2
-spatial_patch_size = 14
-spatial_merge_size = 2
-
-
-def round_by_factor(number: int, factor: int) -> int:
-    """Returns the closest integer to 'number' that is divisible by 'factor'."""
-    return round(number / factor) * factor
-
-
-def ceil_by_factor(number: int, factor: int) -> int:
-    """Returns the smallest integer greater than or equal to 'number' that is
-    divisible by 'factor'."""
-    return math.ceil(number / factor) * factor
-
-
-def floor_by_factor(number: int, factor: int) -> int:
-    """Returns the largest integer less than or equal to 'number' that is
-    divisible by 'factor'."""
-    return math.floor(number / factor) * factor
-
-
-def smart_resize(
-    height: int,
-    width: int,
-    factor: int = IMAGE_FACTOR,
-    min_pixels: int = MIN_PIXELS,
-    max_pixels: int = MAX_PIXELS,
-) -> tuple[int, int]:
-    """
-    Rescales the image so that the following conditions are met:
-
-    1. Both dimensions (height and width) are divisible by 'factor'.
-
-    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
-
-    3. The aspect ratio of the image is maintained as closely as possible.
-    """
-    if max(height, width) / min(height, width) > MAX_RATIO:
-        raise ValueError(
-            f"absolute aspect ratio must be smaller than {MAX_RATIO}, "
-            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

@@ -1,9 +0,0 @@
-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

@@ -1,9 +0,0 @@
-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

@@ -1,324 +0,0 @@
-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

@@ -0,0 +1,77 @@
+# 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

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

examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py

@@ -0,0 +1,228 @@
+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

@@ -0,0 +1,187 @@
+#!/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

@@ -0,0 +1,180 @@
+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.1.0"
+version = "0.11.0rc1"
 description = "vLLM-omni: Multi-modality models inference and serving with non-autoregressive structures"
 readme = "README.md"
-requires-python = ">=3.8"
+requires-python = ">=3.9,<3.14"  # Align with vLLM v0.11 supported Python versions
 license = {text = "Apache-2.0"}
 authors = [
-    {name = "vLLM-omni Team", email = "hsliuustc@gmail.com"}
+    {name = "vLLM-omni Team", email = "vllm-project@gmail.com"}
 ]
 keywords = ["vllm", "multimodal", "diffusion", "transformer", "inference", "serving"]
 classifiers = [
@@ -19,96 +19,108 @@ 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 = [
-    "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",
+    # 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
 ]
 
 [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",
-]
-docs = [
-    "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",
+    "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",
+]
+
+
 [project.urls]
-Homepage = "https://github.com/hsliuustc0106/vllm-omni"
-Repository = "https://github.com/hsliuustc0106/vllm-omni"
+Homepage = "https://github.com/vllm-project/vllm-omni"
+Repository = "https://github.com/vllm-project/vllm-omni"
 Documentation = "https://vllm-omni.readthedocs.io"
-"Bug Tracker" = "https://github.com/hsliuustc0106/vllm-omni/issues"
+"Bug Tracker" = "https://github.com/vllm-project/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.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.setuptools.package-data]
+"vllm_omni.model_executor.stage_configs" = ["*.yaml"]
 
-[tool.isort]
-profile = "black"
-multi_line_output = 3
-line_length = 88
-known_first_party = ["vllm_omni"]
+[tool.ruff]
+line-length = 120
+exclude = [
+    ".eggs",
+    ".git",
+    ".hg",
+    ".mypy_cache",
+    ".tox",
+    ".venv",
+    "build",
+    "dist",
+    "vllm_omni.egg-info",
+]
 
-[tool.flake8]
-max-line-length = 88
-extend-ignore = ["E203", "W503"]
+[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.mypy]
-python_version = "3.8"
+python_version = "3.12, 3.13"
 warn_return_any = true
 warn_unused_configs = true
 disallow_untyped_defs = true
@@ -141,3 +153,11 @@ markers = [
     "benchmark: Benchmark tests",
     "slow: Slow tests",
 ]
+
+[tool.typos.default]
+extend-ignore-identifiers-re = [
+    ".*_thw",
+    ".*thw",
+    "ein",
+    ".*arange"
+]

requirements.txt

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

scripts/README.md

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

scripts/build_wheel.sh

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

scripts/test_serving.sh

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

vllm_omni/__init__.py

@@ -12,16 +12,28 @@ Architecture:
   processing
 """
 
-from . import patch  # noqa: F401
+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 .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", # TODO: add AsyncOmniLLM back when it is implemented
+    "AsyncOmniLLM",
     # 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,15 +4,10 @@ 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
@@ -97,17 +92,13 @@ 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(
@@ -117,9 +108,7 @@ 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
         )
@@ -177,9 +166,7 @@ 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
@@ -196,13 +183,9 @@ 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
@@ -239,11 +222,7 @@ 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)
@@ -252,9 +231,7 @@ 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(  # type: ignore[union-attr]  # noqa: E501
-                    req_id, new_token_ids
-                )
+                request.structured_output_request.grammar.accept_tokens(req_id, new_token_ids)  # type: ignore[union-attr]  # noqa: E501
 
             # 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:
@@ -263,7 +240,6 @@ 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(
@@ -298,10 +274,7 @@ 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:
@@ -312,14 +285,10 @@ 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,6 +3,7 @@ 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,22 +39,14 @@ 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,7 +5,6 @@ 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,6 +2,8 @@ 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
 
 
@@ -31,3 +33,31 @@ 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,7 +1,6 @@
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import torch
-
 from vllm.logger import init_logger
 from vllm.outputs import PoolingRequestOutput
 from vllm.sampling_params import RequestOutputKind
@@ -10,20 +9,16 @@ 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,
@@ -44,7 +39,6 @@ class OmniRequestState(RequestState):
         queue: Optional[Any],
         log_stats: bool,
     ) -> "OmniRequestState":
-
         if sampling_params := request.sampling_params:
             if not sampling_params.detokenize:
                 tokenizer = None
@@ -75,9 +69,7 @@ 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,
@@ -93,9 +85,7 @@ 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:
@@ -137,15 +127,11 @@ 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,
@@ -154,9 +140,7 @@ 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
@@ -196,14 +180,12 @@ 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(
@@ -255,9 +237,7 @@ 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
@@ -270,9 +250,7 @@ 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
@@ -280,12 +258,8 @@ 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(
@@ -310,18 +284,13 @@ 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:
@@ -335,9 +304,7 @@ 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
@@ -352,9 +319,7 @@ 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:
@@ -385,9 +350,7 @@ 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
 
@@ -414,9 +377,7 @@ 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
 
@@ -444,9 +405,7 @@ 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,6 +17,7 @@ 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,
@@ -25,6 +26,8 @@ from vllm_omni.engine import (
 )
 from vllm_omni.inputs.preprocess import OmniInputPreprocessor
 
+logger = init_logger(__name__)
+
 
 class OmniProcessor(Processor):
     @staticmethod
@@ -75,20 +78,14 @@ 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} "
-                f"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} is out of range [0, {data_parallel_size}).")
 
         if arrival_time is None:
             arrival_time = time.time()
@@ -139,19 +136,9 @@ 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
@@ -160,13 +147,9 @@ 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:
@@ -201,12 +184,9 @@ 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
@@ -233,13 +213,9 @@ 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

@@ -0,0 +1,742 @@
+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

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

vllm_omni/entrypoints/cli/main.py

@@ -2,24 +2,55 @@
 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" 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
+    if "--omni" not in sys.argv:
         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,93 +3,65 @@ Omni serve command for vLLM-omni.
 """
 
 import argparse
-from typing import List
+
+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.
+"""
 
 
-class OmniServeCommand:
-    """Command handler for vLLM-omni serve command."""
+class OmniServeCommand(CLISubcommand):
+    """The `serve` subcommand for the vLLM CLI."""
 
-    def __init__(self):
-        self.parser = self._create_parser()
+    name = "serve"
 
-    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"
+    @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]",
         )
 
-        # 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)",
-        )
-
-        # Server arguments
-        parser.add_argument(
-            "--port",
-            type=int,
-            default=8000,
-            help="Port to run the server on",
-        )
-
-        parser.add_argument(
-            "--host",
-            type=str,
-            default="0.0.0.0",
-            help="Host to run the server on",
-        )
-
-        # Stage configuration arguments
-        parser.add_argument("--ar-stage", type=str, help="AR stage model path")
-
-        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",
+        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="Use diffusers pipeline for DiT stage",
+            help="Enable vLLM-omni mode",
         )
-
-        # Other arguments
-        parser.add_argument(
-            "--log-stats", action="store_true", help="Enable logging statistics"
+        serve_parser.add_argument(
+            "--stage-configs-path",
+            type=str,
+            default=None,
+            help="Path to the stage configs file. If not specified, the stage configs will be loaded from the model.",
         )
+        return serve_parser
 
-        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
+def cmd_init() -> list[CLISubcommand]:
+    return [OmniServeCommand()]

vllm_omni/entrypoints/log_utils.py

@@ -1,15 +1,14 @@
 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: Optional[str], num_stages: int) -> None:
+def remove_old_logs(log_file: str | None, num_stages: int) -> None:
     try:
         if not log_file:
             return
@@ -49,7 +48,7 @@ def remove_old_logs(log_file: Optional[str], num_stages: int) -> None:
         pass
 
 
-def configure_orchestrator_logger(logger: logging.Logger, log_file: Optional[str]) -> None:
+def configure_orchestrator_logger(logger: logging.Logger, log_file: str | None) -> None:
     try:
         if not log_file:
             return
@@ -64,9 +63,9 @@ def configure_orchestrator_logger(logger: logging.Logger, log_file: Optional[str
         pass
 
 
-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
+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
     try:
         if enable_stats and log_file:
             stats_file = f"{log_file}.orchestrator.stats.jsonl"
@@ -77,7 +76,7 @@ def init_stats_paths(enable_stats: bool, log_file: Optional[str]) -> tuple[Optio
     return stats_file, overall_file
 
 
-def _safe_append_jsonl(path: Optional[str], record: Dict[str, Any]) -> None:
+def _safe_append_jsonl(path: str | None, record: dict[str, Any]) -> None:
     if not path:
         return
     try:
@@ -86,122 +85,204 @@ def _safe_append_jsonl(path: Optional[str], record: Dict[str, Any]) -> None:
         pass
 
 
-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_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_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_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_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_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_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_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_summary(stats_file: Optional[str], summary: Dict[str, Any]) -> None:
+def log_orchestrator_summary(stats_file: str | None, summary: dict[str, Any]) -> None:
     _safe_append_jsonl(stats_file, {"type": "orchestrator_summary", **summary})
 
 
-def log_overall_summary(overall_stats_file: Optional[str], summary: Dict[str, Any]) -> None:
+def log_overall_summary(overall_stats_file: str | None, summary: dict[str, Any]) -> None:
     _safe_append_jsonl(overall_stats_file, {"type": "overall_summary", **summary})
 
 
-def log_overall_record(overall_stats_file: Optional[str], record: Dict[str, Any]) -> None:
+def log_overall_record(overall_stats_file: str | None, record: dict[str, Any]) -> None:
     _safe_append_jsonl(overall_stats_file, record)
 
 
-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_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_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_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_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 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 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:
+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:
     """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,
@@ -218,13 +299,16 @@ def compute_and_log_stage_request_stats(stats_file: Optional[str], stage_id: int
 
 # ----------------- 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))
@@ -239,23 +323,32 @@ def record_stage_metrics(per_request: Dict[int, Dict[str, Any]],
         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) -> Optional[tuple[int, float, float]]:
+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:
     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)
@@ -287,20 +380,29 @@ def aggregate_rx_and_maybe_total(transfer_edge_req: Dict[tuple[int, int, int], D
         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)
@@ -318,7 +420,7 @@ def record_sender_transfer_agg(transfer_agg: Dict[tuple[int, int], Dict[str, flo
         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:
@@ -332,29 +434,35 @@ 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))
@@ -367,50 +475,57 @@ def build_transfer_summary(transfer_agg: Dict[tuple[int, int], Dict[str, float]]
         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: Optional[str], overall_stats_file: Optional[str],
-                 wall_start_ts: float) -> None:
+    def __init__(
+        self,
+        num_stages: int,
+        enable_stats: bool,
+        stats_file: str | None,
+        overall_stats_file: str | None,
+        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[Optional[float]] = [None for _ in range(self.num_stages)]
-        self.stage_last_ts: List[Optional[float]] = [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[float | None] = [None for _ in range(self.num_stages)]
+        self.stage_last_ts: list[float | None] = [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,
@@ -466,8 +581,15 @@ 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()
@@ -491,8 +613,7 @@ 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:
@@ -526,30 +647,34 @@ 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),
@@ -566,5 +691,3 @@ 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,11 +1,10 @@
-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
@@ -26,19 +25,17 @@ 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 (
-    remove_old_logs,
+    OrchestratorMetrics,
     configure_orchestrator_logger,
     init_stats_paths,
-    OrchestratorMetrics,
-)
-from vllm_omni.entrypoints.stage_utils import (
-    maybe_load_from_ipc as _load,
-    encode_for_ipc as _encode,
-    serialize_obj as _ser,
+    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__)
@@ -52,23 +49,25 @@ class OmniLLM:
     - max_inflight=1 per stage (serial within a stage), but pipeline across stages
     """
 
-    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):
+    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)
         # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent.
-        if stage_configs is None:
+        if stage_configs_path is None:
             self.stage_configs = load_stage_configs_from_model(model)
         else:
-            self.stage_configs = stage_configs
-
+            self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
 
         # Optional file handler for orchestrator
         self._log_file = log_file
@@ -79,8 +78,15 @@ 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
@@ -104,7 +110,6 @@ 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):
@@ -131,7 +136,10 @@ 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()
@@ -147,17 +155,26 @@ class OmniLLM:
     def generate(
         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]:
+        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,
     ) -> 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)):
@@ -176,7 +193,6 @@ 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
@@ -187,10 +203,20 @@ 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))
@@ -214,7 +240,11 @@ 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):
@@ -225,11 +255,17 @@ 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
 
@@ -241,8 +277,17 @@ 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):
@@ -253,15 +298,30 @@ 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:
@@ -272,7 +332,7 @@ class OmniLLM:
                         # Measure transfer size and time (encode + enqueue)
                         size_bytes = 0
                         try:
-                            size_bytes = len(_ser(next_inputs))
+                            size_bytes = len(_set(next_inputs))
                         except Exception:
                             size_bytes = 0
                         t0 = time.time()
@@ -282,27 +342,51 @@ 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)
@@ -341,7 +425,9 @@ 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:
@@ -352,9 +438,7 @@ 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),
@@ -384,13 +468,9 @@ 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."""
@@ -404,9 +484,7 @@ 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"]
@@ -414,8 +492,7 @@ 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,
                 )
@@ -428,16 +505,10 @@ 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
@@ -447,11 +518,7 @@ 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
@@ -467,9 +534,7 @@ 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,
@@ -491,6 +556,4 @@ 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,64 +1,72 @@
 """
 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, 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 typing import Any, Optional, Union
 
 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
-        if hasattr(stage_config, "engine_input_source"):
-            self.engine_input_source = stage_config.engine_input_source
-        else:
-            self.engine_input_source = []
+        self.requires_multimodal_data = getattr(stage_config.runtime, "requires_multimodal_data", False)
+        self.engine_input_source = getattr(stage_config, "engine_input_source", [])
         self.engine_output_type = stage_config.engine_args.engine_output_type
         self.engine_outputs = None
+        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
 
-        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
-
+        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))
         # Runtime orchestration state (added)
         self._in_q: Optional[mp.Queue] = None
         self._out_q: Optional[mp.Queue] = None
@@ -75,6 +83,22 @@ 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
@@ -88,6 +112,7 @@ 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,
@@ -100,16 +125,37 @@ 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,
         }
-        self._proc = ctx.Process(
-            target=_stage_worker,
-            args=(model, stage_payload, self._in_q, self._out_q, self._log_file, batch_timeout),
-        )
+        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.start()
 
     def stop_stage_worker(self) -> None:
@@ -129,11 +175,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()
@@ -142,7 +188,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 = []
@@ -150,9 +196,11 @@ 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: prompt.get("multi_modal_data", None)
-                for source_output, prompt in zip(source_outputs, prompt)
+                source_output.request_id: p.get("multi_modal_data", None)
+                for source_output, p in zip(source_outputs, prompt)
             }
 
             for source_output in source_outputs:
@@ -160,7 +208,7 @@ class OmniStage:
                     prompt_token_ids=source_output.outputs[0].token_ids,
                     multi_modal_data=(
                         multi_modal_data[source_output.request_id]
-                        if multi_modal_data
+                        if self.requires_multimodal_data and multi_modal_data
                         else None
                     ),
                 )
@@ -170,13 +218,13 @@ class OmniStage:
         else:
             engine_input_source = self.engine_input_source
             return self.custom_process_input_func(
-                stage_list, engine_input_source, prompt
+                stage_list, engine_input_source, prompt, self.requires_multimodal_data
             )
 
 
 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,
@@ -184,15 +232,16 @@ def _stage_worker(
 ) -> None:
     """Stage worker entry: device setup, LLM init, batching, SHM IPC."""
     import logging as _logging
-    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,
+    import time as _time
+
+    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,
     )
-    import queue as _queue
-    import time as _time
+    from vllm_omni.entrypoints.omni_llm import OmniStageLLM
+
     # no inline JSONL/serialization imports; logging handled by utilities
 
     stage_id = stage_payload["stage_id"]
@@ -207,11 +256,15 @@ 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:
@@ -233,7 +286,11 @@ 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
@@ -251,7 +308,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():
@@ -263,11 +320,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:
@@ -278,9 +335,7 @@ 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)
@@ -291,13 +346,21 @@ 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)
@@ -305,8 +368,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:
@@ -326,12 +389,25 @@ 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:
@@ -354,25 +430,281 @@ 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({
+                        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(
+                        {
                             "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:
-                    out_q.put({
+                        }
+                    )
+            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(
+                    {
                         "request_id": rid,
                         "stage_id": stage_id,
                         "engine_outputs": r_outputs,
@@ -383,13 +715,16 @@ def _stage_worker(
                             "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 batch %s: %s", stage_id, batch_request_ids, e)
-            for rid in batch_request_ids:
-                out_q.put({
+            _logging.getLogger(__name__).exception("[Stage-%s] Failed on request %s: %s", stage_id, rid, e)
+            out_q.put(
+                {
                     "request_id": rid,
                     "stage_id": stage_id,
                     "error": str(e),
-                })
+                }
+            )

vllm_omni/entrypoints/openai/api_server.py

@@ -0,0 +1,313 @@
+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

@@ -0,0 +1,779 @@
+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
 
-from typing import Any, Dict, Optional, Tuple, Union
-
-import logging
 import json
+import logging
 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: Optional[Union[str, int]]) -> None:
+def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None:
     """Configure per-stage CUDA visibility and current device.
 
     Behavior
@@ -26,8 +26,8 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
     - Otherwise: set CUDA_VISIBLE_DEVICES to the provided single device string.
     """
     try:
-        selected_physical: Optional[int] = None
-        logical_idx: Optional[int] = None
+        selected_physical: int | None = None
+        logical_idx: int | None = None
 
         if isinstance(devices, str) and "," in devices:
             os.environ["CUDA_VISIBLE_DEVICES"] = devices
@@ -37,7 +37,9 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
                     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)
@@ -51,43 +53,70 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
                     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)
 
@@ -97,7 +126,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.
@@ -110,11 +139,11 @@ def shm_write_bytes(payload: bytes) -> Dict[str, Any]:
     try:
         shm.close()
     except Exception as e:
-        logger.debug("Ensure parent dir failed for %s: %s", path, e)
+        logger.debug("Failed to close shared memory: %s", 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)
@@ -141,7 +170,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
@@ -158,7 +187,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).
@@ -169,7 +198,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
@@ -180,7 +209,9 @@ 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:
@@ -188,6 +219,7 @@ def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, sh
       - 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]
@@ -211,13 +243,13 @@ def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, sh
     }
 
 
-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
@@ -227,7 +259,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,10 +1,14 @@
+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__)
+
 # Get the project root directory (2 levels up from this file)
 PROJECT_ROOT = Path(__file__).parent.parent.parent
 

vllm_omni/inputs/data.py

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

vllm_omni/inputs/parse.py

@@ -17,15 +17,9 @@ 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,13 +1,16 @@
 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(
@@ -17,9 +20,7 @@ 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")
 
@@ -62,7 +63,7 @@ class OmniInputPreprocessor(InputPreprocessor):
 
         Returns:
 
-        * [`SingletonInputs`][vllm.inputs.data.SingletonInputs] instance
+        * Input container compatible with vLLM's singleton prompt handling.
         """
         parsed = parse_singleton_prompt_omni(prompt)
 

vllm_omni/model_executor/layers/mrope.py

@@ -4,7 +4,6 @@ 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
 
 
@@ -249,9 +248,7 @@ 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
@@ -260,9 +257,7 @@ 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],
@@ -275,27 +270,10 @@ 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":
@@ -311,36 +289,17 @@ 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:
@@ -372,9 +331,7 @@ 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()
@@ -424,15 +381,11 @@ 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
                 )
@@ -440,29 +393,15 @@ 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()
@@ -512,9 +451,7 @@ 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)
@@ -534,21 +471,13 @@ 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])
@@ -571,58 +500,40 @@ 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,
@@ -633,38 +544,23 @@ 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
@@ -673,9 +569,7 @@ 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
@@ -692,25 +586,10 @@ 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)
@@ -731,14 +610,7 @@ 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(
@@ -748,7 +620,6 @@ 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,
@@ -779,25 +650,19 @@ 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,7 +4,6 @@ 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/qwen2_5_omni.py

@@ -1,26 +1,21 @@
 import glob
 import os
+from collections.abc import Iterable
 from functools import cached_property
-from typing import Dict, Iterable, NamedTuple, Optional, Set, Tuple, Union
+from typing import NamedTuple, Optional, Union
 
 import numpy as np
 import torch
 import torch.nn as nn
+from transformers import PretrainedConfig
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
     Qwen2_5OmniConfig,
     Qwen2_5OmniTalkerConfig,
     Qwen2_5OmniThinkerConfig,
 )
-
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
-from vllm.model_executor.models.qwen2_5_omni_thinker import (
-    Qwen2_5OmniConditionalGenerationMixin,
-    Qwen2_5OmniThinkerDummyInputsBuilder,
-    Qwen2_5OmniThinkerMultiModalProcessor,
-    Qwen2_5OmniThinkerProcessingInfo,
-)
+from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP
 from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
 
 # from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav
@@ -29,10 +24,16 @@ 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.model_loader.weight_utils import download_weights_from_hf_specific
+from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
+    Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniThinkerDummyInputsBuilder,
+    Qwen2_5OmniThinkerMultiModalProcessor,
+    Qwen2_5OmniThinkerProcessingInfo,
 )
-from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
+from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
+from vllm_omni.model_executor.models.vision import get_llm_pos_ids_for_vision
 
 TALKER_CODEC_EOS_TOKEN_ID = 8294
 TALKER_CODEC_BOS_TOKEN_ID = 8293
@@ -55,9 +56,8 @@ logger = init_logger(__name__)
     dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
 class Qwen2_5OmniForConditionalGeneration(
-    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin, SupportsMRoPE
 ):
-
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.have_multimodal_outputs = True
@@ -117,17 +117,15 @@ 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 --------------------
@@ -176,11 +174,7 @@ 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):
@@ -229,9 +223,7 @@ 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
@@ -259,11 +251,7 @@ 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,
             )
 
@@ -273,13 +261,17 @@ 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)
@@ -293,13 +285,17 @@ 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]
@@ -332,9 +328,19 @@ 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)):
@@ -354,7 +360,9 @@ 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):
@@ -378,7 +386,11 @@ 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:]
@@ -403,7 +415,8 @@ 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(
@@ -421,16 +434,12 @@ 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(
@@ -441,9 +450,7 @@ 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))
                         )
@@ -454,6 +461,170 @@ class Qwen2_5OmniForConditionalGeneration(
             multimodal_outputs=None,
         )
 
+    def get_mrope_input_positions(
+        self,
+        input_tokens: list[int],
+        hf_config: PretrainedConfig,
+        image_grid_thw: list[list[int]] | torch.Tensor,
+        video_grid_thw: list[list[int]] | torch.Tensor,
+        second_per_grid_ts: list[float] | None = None,
+        context_len: int = 0,
+        seq_len: int | None = None,
+        audio_feature_lengths: torch.Tensor | None = None,
+        use_audio_in_video: bool = False,
+    ) -> tuple[torch.Tensor, int]:
+        """Get mrope input positions and delta value (Qwen2.5-Omni version).
+
+        Differences from MRotaryEmbedding:
+            1. Add audio support (and related `audio_feature_lengths`).
+            2. Add `use_audio_in_video` option to read audio from video inputs.
+                In this case, audio and vision position ids will be split into
+                chunks and interleaved.
+
+        Example:
+
+            (V_i are vision position ids, A_i are audio position ids)
+
+            |V_1 ...    V_n|A_1 ...   A_n|V_n+1 ... V_2n|A_n+1 ... A_2n|...
+            |vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |...
+        """
+
+        # TODO(fyabc): refactor and share more code with
+        #  _vl_get_input_positions_tensor.
+
+        thinker_config = hf_config.thinker_config
+        audio_token_id = thinker_config.audio_token_index
+        image_token_id = thinker_config.image_token_index
+        video_token_id = thinker_config.video_token_index
+        audio_start_token_id = thinker_config.audio_start_token_id
+        audio_end_token_id = thinker_config.audio_end_token_id
+        vision_start_token_id = thinker_config.vision_start_token_id
+        vision_end_token_id = thinker_config.vision_end_token_id
+        seconds_per_chunk = thinker_config.seconds_per_chunk
+        spatial_merge_size = thinker_config.vision_config.spatial_merge_size
+        tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25)
+
+        if isinstance(image_grid_thw, list):
+            image_grid_thw = torch.tensor(image_grid_thw)
+        if isinstance(video_grid_thw, list):
+            video_grid_thw = torch.tensor(video_grid_thw)
+
+        src_item = input_tokens
+        audio_seqlens = audio_feature_lengths
+        if not second_per_grid_ts:
+            second_per_grid_ts = [1] * video_grid_thw.shape[0]
+        audio_idx = 0
+        video_idx = 0
+        image_idx = 0
+        new_src_item: list[int] = []
+        llm_pos_ids_list: list[torch.Tensor] = []
+
+        idx = 0
+        while idx < len(src_item):
+            new_src_item_len = len(new_src_item)
+            start_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+            if src_item[idx] not in [audio_token_id, video_token_id, image_token_id]:
+                if use_audio_in_video and idx > 0:
+                    if src_item[idx] == vision_end_token_id and src_item[idx - 1] == audio_end_token_id:
+                        # processing the <|audio_eos|> before <|vision_eos|>
+                        start_idx -= 1
+                    elif src_item[idx] == audio_start_token_id and src_item[idx - 1] == vision_start_token_id:
+                        # processing the <|audio_bos|> after <|vision_eos|>
+                        start_idx -= 1
+                new_src_item.append(src_item[idx])
+                llm_pos_ids = torch.tensor([start_idx], dtype=torch.long).expand(3, -1)
+                llm_pos_ids_list.append(llm_pos_ids)
+            elif src_item[idx] == audio_token_id:
+                assert audio_seqlens is not None
+                audio_seqlen = audio_seqlens[audio_idx]
+                place_num = ((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1
+                new_src_item.extend([audio_token_id] * place_num)
+                llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx
+                llm_pos_ids_list.append(llm_pos_ids)
+                audio_idx += 1
+            elif src_item[idx] == image_token_id:
+                grid_t = image_grid_thw[image_idx][0]
+                grid_hs = image_grid_thw[:, 1]
+                grid_ws = image_grid_thw[:, 2]
+                t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long()
+                llm_pos_ids = get_llm_pos_ids_for_vision(
+                    start_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+                )
+                llm_pos_ids_list.append(llm_pos_ids)
+                vision_seqlen = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
+                new_src_item.extend([image_token_id] * vision_seqlen)
+                image_idx += 1
+            elif src_item[idx] == video_token_id and not use_audio_in_video:
+                grid_t = video_grid_thw[video_idx][0]
+                grid_hs = video_grid_thw[:, 1]
+                grid_ws = video_grid_thw[:, 2]
+                t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long()
+                llm_pos_ids = get_llm_pos_ids_for_vision(
+                    start_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+                )
+                llm_pos_ids_list.append(llm_pos_ids)
+                vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
+                new_src_item.extend([video_token_id] * vision_seqlen)
+                video_idx += 1
+            else:
+                # read audio from video
+                assert audio_seqlens is not None
+                audio_seqlen = audio_seqlens[audio_idx]
+                vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
+                grid_t = video_grid_thw[video_idx][0]
+                grid_h = video_grid_thw[video_idx][1]
+                grid_w = video_grid_thw[video_idx][2]
+                grid_hs = video_grid_thw[:, 1]
+                grid_ws = video_grid_thw[:, 2]
+                t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
+                t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long()
+                t_index_split_chunk = split_list_into_ranges(t_index, t_ntoken_per_chunk)
+                place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2
+                pure_audio_len = place_num - 2
+                added_audio_len = 0
+                audio_llm_pos_ids_list: list[torch.Tensor] = []
+                for t_chunk in t_index_split_chunk:
+                    vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
+                    new_src_item.extend([video_token_id] * vision_ntoken_per_chunk)
+                    vision_llm_pos_ids_list = get_llm_pos_ids_for_vision(
+                        start_idx,
+                        video_idx,
+                        spatial_merge_size,
+                        t_chunk,
+                        grid_hs,
+                        grid_ws,
+                    ).split(1, dim=1)
+                    llm_pos_ids_list.extend(vision_llm_pos_ids_list)
+                    new_src_item.extend(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) * [audio_token_id])
+                    audio_start_idx = (
+                        start_idx if len(audio_llm_pos_ids_list) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1
+                    )
+                    if min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) > 0:
+                        audio_llm_pos_ids_list = (
+                            torch.arange(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)).expand(3, -1)
+                            + audio_start_idx
+                        ).split(1, dim=1)
+                    else:
+                        audio_llm_pos_ids_list = []
+                    added_audio_len += min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)
+                    llm_pos_ids_list.extend(audio_llm_pos_ids_list)
+                if added_audio_len < pure_audio_len:
+                    new_src_item.extend((pure_audio_len - added_audio_len) * [audio_token_id])
+                    audio_llm_pos_ids_list = (
+                        torch.arange(pure_audio_len - added_audio_len).expand(3, -1) + llm_pos_ids_list[-1].max() + 1
+                    ).split(1, dim=1)
+                    llm_pos_ids_list.extend(audio_llm_pos_ids_list)
+                audio_idx += 1
+                video_idx += 1
+            # move to the next token
+            idx += len(new_src_item) - new_src_item_len
+
+        llm_positions = torch.cat(llm_pos_ids_list, dim=1)
+        mrope_position_delta = torch.cat(llm_pos_ids_list, dim=1).max() + 1 - len(src_item)
+        llm_positions = llm_positions[:, context_len:seq_len]
+
+        return llm_positions, mrope_position_delta
+
     def generate_audio(self, code, voice_type):
         token2wav_dev = self._module_device(self.token2wav)
         if isinstance(code, torch.Tensor):
@@ -570,7 +741,6 @@ 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
@@ -593,9 +763,7 @@ 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,
             ]
@@ -606,9 +774,7 @@ 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
 
@@ -617,14 +783,12 @@ 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
@@ -640,7 +804,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"):
+    def generate_speech(self, text_tokens: torch.Tensor, voice_type: str = "default") -> torch.Tensor:
         """
         Generate speech from text tokens using the talker and token2wav models.
         This method is kept for backward compatibility and direct speech generation.
@@ -653,9 +817,7 @@ 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)
@@ -682,9 +844,7 @@ 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
@@ -703,12 +863,8 @@ 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):
@@ -720,24 +876,14 @@ 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:
@@ -771,9 +917,7 @@ 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)
@@ -788,9 +932,7 @@ 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 "
@@ -802,9 +944,7 @@ 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
@@ -833,7 +973,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 = []
@@ -862,20 +1002,14 @@ 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")
@@ -896,9 +1030,7 @@ 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/qwen2_5_omni_talker.py

@@ -1,24 +1,17 @@
+from collections.abc import Iterable
 from functools import cached_property
-from typing import Iterable, Optional, Set, Tuple, Union
+from typing import Optional, 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,
@@ -72,9 +65,7 @@ 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
 
@@ -92,9 +83,7 @@ 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)
@@ -122,7 +111,6 @@ 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(
@@ -145,9 +133,7 @@ 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:
@@ -166,8 +152,16 @@ 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]:
-        return self.language_model.compute_logits(hidden_states)
+        logits = self.language_model.compute_logits(hidden_states)
+        logits = self.bad_word_processor(logits)
+        return logits
 
     def sample(
         self,
@@ -176,7 +170,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."],
@@ -212,31 +206,15 @@ 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/qwen2_5_omni_thinker.py

@@ -1,30 +1,18 @@
+from collections import defaultdict
 from collections.abc import Iterable, Mapping, Sequence
-from copy import copy
 from functools import partial
-from typing import Any, Optional, Union
+from typing import Annotated, Any, Callable, Literal, Optional, Union
 
 import torch
 import torch.nn as nn
 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,
-    SupportsMultiModal,
-    SupportsPP,
-)
+from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMRoPE, SupportsMultiModal, SupportsPP
 from vllm.model_executor.models.qwen2_5_vl import (
     Qwen2_5_VisionTransformer,
     Qwen2_5_VLImageEmbeddingInputs,
@@ -36,7 +24,6 @@ from vllm.model_executor.models.qwen2_5_vl import (
     Qwen2_5_VLVideoPixelInputs,
 )
 from vllm.model_executor.models.qwen2_audio import (
-    Qwen2AudioInputs,
     Qwen2AudioProcessingInfo,
     _get_feat_extract_output_lengths,
 )
@@ -54,7 +41,7 @@ from vllm.multimodal.inputs import (
     ModalityData,
     MultiModalDataDict,
     MultiModalFieldConfig,
-    MultiModalKwargs,
+    MultiModalKwargsItems,
     NestedTensors,
 )
 from vllm.multimodal.parse import (
@@ -66,13 +53,18 @@ 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:
@@ -83,6 +75,27 @@ except (ImportError, ModuleNotFoundError):
 logger = init_logger(__name__)
 
 
+class Qwen2_5OmniAudioFeatureInputs(TensorSchema):
+    """
+    Dimensions:
+        - na: Number of audios
+        - nmb: Number of mel bins
+        - msl: Maximum sequence length
+        - tsl: Total sequence length
+    """
+
+    type: Literal["audio_features"]
+    input_features: Annotated[
+        Union[torch.Tensor, list[torch.Tensor]],
+        TensorShape("nmb", "tsl"),
+    ]
+
+    feature_attention_mask: Annotated[
+        torch.Tensor,
+        TensorShape("na", "msl"),
+    ]
+
+
 def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
     audio_feature_lengths = hf_inputs.get("audio_feature_lengths", torch.empty((0,)))
 
@@ -93,25 +106,53 @@ 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"),
     )
 
 
-class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
+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]],
@@ -127,10 +168,7 @@ 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
 
@@ -151,10 +189,7 @@ class Qwen2_5OmniThinkerProcessingInfo(
         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)
@@ -166,11 +201,7 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(
         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,
@@ -191,36 +222,25 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(
             * 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(
-            target_sr=feature_extractor.sampling_rate
+            spatial_merge_size=self.info.get_hf_config().vision_config.spatial_merge_size,
+            target_sr=feature_extractor.sampling_rate,
         )
 
     def _call_hf_processor(
@@ -252,15 +272,18 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         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(
@@ -268,70 +291,177 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         hf_inputs: BatchFeature,
         hf_processor_mm_kwargs: Mapping[str, object],
     ) -> Mapping[str, MultiModalFieldConfig]:
-        return _qwen2_5_omni_thinker_field_config(hf_inputs)
+        return create_qwen2_5_omni_thinker_field_factory(self.info.get_hf_config().vision_config.spatial_merge_size)(
+            hf_inputs
+        )
+
+    def _apply_prompt_updates(
+        self,
+        token_ids: list[int],
+        mm_prompt_updates: MultiModalPromptUpdates,
+    ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
+        tokenizer = self.info.get_tokenizer()
+
+        new_token_ids, match_result = self._apply_token_matches(
+            token_ids,
+            mm_prompt_updates,
+        )
+
+        # If the search text does not represent a special token,
+        # it may have different token IDs in the prompt, because
+        # the tokens may go across the boundaries of the search text.
+        # ----
+        # e.g. when searching for "foo" in "food", if "food" itself makes
+        # up a token, then the token ID of "foo" will not appear at all
+        # ----
+        # Since it is inefficient to search for all possible tokenizations
+        # of the search text in the prompt, we instead perform string-based
+        # updates on the decoded token IDs, then encode them back.
+        if not all(all(update_idx is not None for update_idx in update_idxs) for update_idxs in match_result.values()):
+            new_text, match_result = self._apply_text_matches(
+                decode_tokens(tokenizer, token_ids),
+                mm_prompt_updates,
+            )
+
+            new_token_ids = encode_tokens(
+                tokenizer,
+                new_text,
+                add_special_tokens=False,
+            )
+
+        matched_updates = defaultdict[str, list[Sequence[ResolvedPromptUpdate]]](list)
+        unmatched_audio_items = []
+
+        for modality, update_idxs in match_result.items():
+            for item_idx, update_idx in enumerate(update_idxs):
+                if update_idx is None:
+                    # Check if this is audio that might be embedded in video
+                    if modality == "audio":
+                        unmatched_audio_items.append(item_idx)
+                        continue
+                    else:
+                        assert False, f"Failed to apply prompt replacement for mm_items[{modality!r}][{item_idx}]"
+
+                matched_updates[modality].append([mm_prompt_updates[modality][item_idx][update_idx]])
+
+        # If there are unmatched audio items, check if we can derive them from video
+        if unmatched_audio_items:
+            # Check if video exists in the original updates (not just matched ones)
+            if "video" in mm_prompt_updates:
+                num_videos = len(mm_prompt_updates.get("video", []))
+                num_audios = len(mm_prompt_updates.get("audio", []))
+                # If counts match, audio might be embedded in video - skip assertion
+                # The placeholders will be derived from video in _maybe_apply_prompt_updates
+                if num_audios == num_videos and "video" in matched_updates:
+                    pass  # Will be handled by deriving from video placeholders
+                else:
+                    # Audio items exist but can't be matched
+                    # Check if video was also matched
+                    if "video" not in matched_updates:
+                        # Neither audio nor video matched - this is a real error
+                        assert False, (
+                            f"Failed to apply prompt replacement for {len(unmatched_audio_items)} "
+                            f"audio item(s). Audio items cannot be matched in prompt."
+                        )
+                    elif num_audios != num_videos:
+                        # Audio and video counts don't match
+                        assert False, (
+                            f"Failed to apply prompt replacement for {len(unmatched_audio_items)} "
+                            f"audio item(s). Audio items cannot be matched in prompt and "
+                            f"audio count ({num_audios}) does not match video count ({num_videos})."
+                        )
+                    else:
+                        # Video matched but audio didn't, and counts match - allow it
+                        pass
+            else:
+                # Audio items can't be matched and there are no video items
+                assert False, (
+                    f"Failed to apply prompt replacement for {len(unmatched_audio_items)} "
+                    f"audio item(s). Audio items cannot be matched in prompt."
+                )
+
+        placeholders = self._find_mm_placeholders(
+            new_token_ids,
+            dict(matched_updates),
+        )
+
+        return new_token_ids, placeholders
 
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
         prompt_ids: list[int],
-        mm_kwargs: MultiModalKwargs,
+        mm_kwargs: MultiModalKwargsItems,
+        mm_prompt_updates: MultiModalPromptUpdates,
         is_update_applied: bool,
     ) -> tuple[list[int], str, Mapping[str, list[PlaceholderFeaturesInfo]]]:
         """
         Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.
         """
-        unbound_prompt_updates = self._get_prompt_updates(
-            mm_items,
-            hf_processor_mm_kwargs,
-            mm_kwargs,
-        )
-        mm_prompt_updates = self._bind_and_group_updates(unbound_prompt_updates)
-
         mm_item_counts = mm_items.get_all_counts()
         self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
 
-        use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)
+        use_audio_in_video = False
+        if "video" in mm_kwargs:
+            video_items = [item for item in mm_kwargs["video"] if item is not None]
+            # only check video items (if there are any)
+            if video_items:
+                use_audio_in_video = all(item["use_audio_in_video"].data for item in video_items)
 
         if is_update_applied:
             mm_placeholders = self._find_mm_placeholders(
-                mm_prompt_updates,
                 prompt_ids,
-                mm_item_counts,
+                mm_prompt_updates,
             )
             self._validate_mm_placeholders(
-                mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video
+                mm_placeholders,
+                mm_item_counts,
             )
 
-            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,
-                mm_item_counts,
-            )
+            # 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, use_audio_in_video=use_audio_in_video
+                mm_placeholders,
+                mm_item_counts,
             )
 
         tokenizer = self.info.get_tokenizer()
         prompt = decode_tokens(tokenizer, prompt_ids)
 
-        if use_audio_in_video:
-            mm_kwargs["use_audio_in_video"] = True
-
         return prompt_ids, prompt, mm_placeholders
 
     def _get_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
         hf_processor_mm_kwargs: Mapping[str, Any],
-        out_mm_kwargs: MultiModalKwargs,
+        out_mm_kwargs: MultiModalKwargsItems,
     ) -> Sequence[PromptUpdate]:
         processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
         tokenizer = self.info.get_tokenizer()
@@ -345,20 +475,17 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         image_token_id = vocab[image_token]
         video_token_id = vocab[video_token]
 
-        audio_feature_lengths = out_mm_kwargs.get("audio_feature_lengths")
-        feature_attention_mask = out_mm_kwargs.get("feature_attention_mask")
+        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")
         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.
@@ -372,14 +499,13 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
                 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_kwargs[f"{modality}_grid_thw"][item_idx]
+            grid_thw = out_mm_data[f"{modality}_grid_thw"][item_idx]
             assert isinstance(grid_thw, torch.Tensor)
             merge_length = image_processor.merge_size**2
 
@@ -392,10 +518,8 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         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_kwargs["video_grid_thw"][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_in_video_item_idx += 1
 
@@ -405,13 +529,15 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             else:
                 video_second_per_grid_t = 1.0
 
-            return MRotaryEmbedding.omni_get_updates_use_audio_in_video(
+            placeholder = MRotaryEmbedding.omni_get_updates_use_audio_in_video(
                 thinker_config=thinker_config,
                 audio_len=audio_num_features,
                 video_grid_thw=video_grid_thw,
                 video_second_per_grid_t=video_second_per_grid_t,
             )
 
+            return PromptUpdateDetails.select_token_id(placeholder, embed_token_id=video_token_id)
+
         video_replacement_fn = (
             get_replacement_qwen2_use_audio_in_video
             if use_audio_in_video
@@ -436,6 +562,50 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             ),
         ]
 
+    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]],
@@ -494,55 +664,33 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
 
         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}. " f"Got type: {type(mm_input)}"
-            )
+            raise ValueError(f"Incorrect type of {name}. Got type: {type(mm_input)}")
         if isinstance(mm_input, torch.Tensor):
+            if dim == 0:
+                return mm_input.reshape(-1, *mm_input.shape[2:])
             return torch.concat(list(mm_input), dim=dim)
         else:
             return torch.concat(mm_input, dim=dim)
 
-    def _parse_and_validate_audio_input(
-        self, **kwargs: object
-    ) -> Optional[Qwen2AudioInputs]:
+    def _parse_and_validate_audio_input(self, **kwargs: object) -> Optional[Qwen2_5OmniAudioFeatureInputs]:
         input_audio_features = kwargs.pop("input_audio_features", None)
         audio_feature_lengths = kwargs.pop("audio_feature_lengths", None)
         feature_attention_mask = kwargs.pop("feature_attention_mask", None)
         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(
-                "Incorrect type of audio input features. "
-                f"Got type: {type(input_audio_features)}"
-            )
-        return Qwen2AudioInputs(
+            raise ValueError(f"Incorrect type of audio input features. Got type: {type(input_audio_features)}")
+        return Qwen2_5OmniAudioFeatureInputs(
+            type="audio_features",
             input_features=input_audio_features,
             audio_feature_lengths=audio_feature_lengths,
             feature_attention_mask=feature_attention_mask,
@@ -560,42 +708,24 @@ 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(
-                    "Incorrect type of image pixel values. "
-                    f"Got type: {type(pixel_values)}"
-                )
+                raise ValueError(f"Incorrect type of image pixel values. Got type: {type(pixel_values)}")
 
             return Qwen2_5_VLImagePixelInputs(
-                type="pixel_values",
-                pixel_values=pixel_values,
-                image_grid_thw=image_grid_thw,
+                type="pixel_values", pixel_values=pixel_values, image_grid_thw=image_grid_thw
             )
 
         if image_embeds is not None:
-            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(
-                    "Incorrect type of image embeddings. "
-                    f"Got type: {type(image_embeds)}"
-                )
+                raise ValueError(f"Incorrect type of image embeddings. Got type: {type(image_embeds)}")
             return Qwen2_5_VLImageEmbeddingInputs(
-                type="image_embeds",
-                image_embeds=image_embeds,
-                image_grid_thw=image_grid_thw,
+                type="image_embeds", image_embeds=image_embeds, image_grid_thw=image_grid_thw
             )
 
     def _parse_and_validate_video_input(
@@ -610,12 +740,8 @@ 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",
@@ -624,48 +750,35 @@ 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(
-                    "Incorrect type of video embeddings. "
-                    f"Got type: {type(video_embeds)}"
-                )
+                raise ValueError(f"Incorrect type of video embeddings. Got type: {type(video_embeds)}")
             return Qwen2_5_VLVideoEmbeddingInputs(
-                type="video_embeds",
-                video_embeds=video_embeds,
-                video_grid_thw=video_grid_thw,
+                type="video_embeds", video_embeds=video_embeds, video_grid_thw=video_grid_thw
             )
 
     def _process_audio_input(
         self,
-        audio_input: Qwen2AudioInputs,
+        audio_input: Qwen2_5OmniAudioFeatureInputs,
         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(
@@ -673,12 +786,9 @@ class Qwen2_5OmniConditionalGenerationMixin:
             feature_lens=audio_feature_lengths,
             aftercnn_lens=audio_feat_lengths,
         )
-        audio_features = audio_outputs.last_hidden_state
-        return audio_features.split(audio_output_lengths.tolist())
+        return audio_outputs.last_hidden_state.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)
 
@@ -720,7 +830,7 @@ class Qwen2_5OmniConditionalGenerationMixin:
     dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
 class Qwen2_5OmniThinkerForConditionalGeneration(
-    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin, SupportsMRoPE
 ):
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
@@ -758,17 +868,24 @@ 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."
             )
 
-        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"),
-        )
+        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.quant_config = quant_config
         self.language_model = init_vllm_registered_model(
             vllm_config=vllm_config,
@@ -777,9 +894,7 @@ 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 = {}
@@ -787,34 +902,18 @@ 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 []
@@ -845,7 +944,6 @@ 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(
@@ -897,9 +995,7 @@ 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(
@@ -921,10 +1017,7 @@ 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/qwen2_5_omni_token2wav.py

@@ -3,7 +3,8 @@
 ############################
 
 import math
-from typing import Iterable, List, Optional, Set, Tuple, Union
+from collections.abc import Iterable
+from typing import Optional, Union
 
 import numpy as np
 import torch
@@ -16,24 +17,17 @@ 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
@@ -218,9 +212,7 @@ 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):
@@ -282,9 +274,7 @@ 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,
@@ -312,13 +302,10 @@ 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()
 
@@ -412,26 +399,14 @@ 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
 
@@ -465,13 +440,9 @@ 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
 
 
@@ -490,9 +461,7 @@ 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
 
 
@@ -572,9 +541,7 @@ 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,
@@ -600,9 +567,7 @@ 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(
@@ -616,9 +581,7 @@ 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
@@ -648,9 +611,7 @@ 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)
@@ -661,29 +622,21 @@ 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(
@@ -696,9 +649,7 @@ 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
 
@@ -741,14 +692,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, half_width, kernel_size):
+def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
     """Generates a 1D Kaiser-windowed sinc filter.
 
     Args:
@@ -773,9 +724,7 @@ def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
     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:
@@ -785,9 +734,7 @@ def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
 
     # 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
@@ -802,27 +749,19 @@ 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),
@@ -855,9 +794,7 @@ 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),
@@ -958,15 +895,10 @@ 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):
@@ -999,9 +931,7 @@ 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 = [
@@ -1016,9 +946,7 @@ 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)
 
@@ -1030,16 +958,12 @@ 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),
@@ -1079,9 +1003,7 @@ 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
@@ -1109,11 +1031,7 @@ 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,
@@ -1139,9 +1057,7 @@ 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:
@@ -1162,15 +1078,10 @@ 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,
@@ -1201,9 +1112,7 @@ 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)
@@ -1223,16 +1132,12 @@ 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]
@@ -1257,12 +1162,8 @@ 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,
@@ -1302,17 +1203,11 @@ 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")
@@ -1339,10 +1234,7 @@ 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(
@@ -1354,13 +1246,9 @@ 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]
@@ -1370,14 +1258,14 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
     @torch.no_grad()
     def fast_block_sample(
         self,
-        conditioning_vector,
-        reference_mel_spectrogram,
-        quantized_code,
+        conditioning_vector: torch.Tensor,
+        reference_mel_spectrogram: torch.Tensor,
+        quantized_code: torch.Tensor,
         y0: torch.Tensor,
-        num_steps=10,
-        guidance_scale=0.5,
+        num_steps: int = 10,
+        guidance_scale: float = 0.5,
         sway_coefficient: Optional[float] = -1.0,
-    ):
+    ) -> torch.Tensor:
         """
         Block-wise ODE sampling starting from provided initial state y0.
 
@@ -1391,9 +1279,7 @@ 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")
@@ -1420,10 +1306,7 @@ 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(
@@ -1435,13 +1318,9 @@ 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]
@@ -1477,8 +1356,7 @@ 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(
@@ -1501,9 +1379,7 @@ 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
 
@@ -1585,7 +1461,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(
@@ -1593,16 +1469,8 @@ 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(
@@ -1632,9 +1500,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
@@ -1659,7 +1527,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.
 
@@ -1676,21 +1544,13 @@ 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:
@@ -1698,11 +1558,7 @@ 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
 
 
@@ -1726,9 +1582,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"],
         )
@@ -1780,9 +1636,7 @@ 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
@@ -1820,9 +1674,7 @@ 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:
@@ -1842,9 +1694,7 @@ 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)
@@ -1872,9 +1722,7 @@ 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
 
@@ -1889,7 +1737,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,
@@ -1907,9 +1755,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,
@@ -1917,9 +1765,7 @@ 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_5_omni/qwen2_old.py

@@ -1,9 +1,9 @@
-from typing import Iterable, Optional, Set, Tuple, Union
+from collections.abc import Iterable
+from typing import Optional, 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,23 +11,13 @@ 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,
@@ -48,7 +38,6 @@ logger = init_logger(__name__)
 
 
 class Qwen2MLP(nn.Module):
-
     def __init__(
         self,
         hidden_size: int,
@@ -73,10 +62,7 @@ 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):
@@ -87,7 +73,6 @@ class Qwen2MLP(nn.Module):
 
 
 class Qwen2Attention(nn.Module):
-
     def __init__(
         self,
         hidden_size: int,
@@ -98,7 +83,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:
@@ -173,7 +158,6 @@ class Qwen2Attention(nn.Module):
 
 
 class Qwen2DecoderLayer(nn.Module):
-
     def __init__(
         self,
         config: Qwen2Config,
@@ -217,16 +201,14 @@ 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
@@ -255,7 +237,6 @@ class Qwen2DecoderLayer(nn.Module):
     }
 )
 class Qwen2Model(nn.Module):
-
     def __init__(
         self,
         *,
@@ -270,27 +251,20 @@ 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 "
-                "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,
-                )
+                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."
             )
 
         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),
@@ -348,13 +322,11 @@ 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"),
@@ -364,19 +336,15 @@ 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
@@ -433,9 +401,7 @@ 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:
@@ -453,9 +419,7 @@ 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)
@@ -467,9 +431,7 @@ 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(
@@ -487,7 +449,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),
@@ -521,9 +483,7 @@ 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
@@ -554,9 +514,7 @@ 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,31 +1,32 @@
-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_old", "Qwen2ForCausalLM"),  # need to discuss
+    "Qwen2ForCausalLM_old": ("qwen2_5_omni", "qwen2_old", "Qwen2ForCausalLM"),  # need to discuss
 }
 
 _VLLM_OMNI_MODELS = {
@@ -45,10 +46,10 @@ OmniModelRegistry = _ModelRegistry(
         },
         **{
             model_arch: _LazyRegisteredModel(
-                module_name=f"vllm_omni.model_executor.models.{mod_relname}",
+                module_name=f"vllm_omni.model_executor.models.{mod_folder}.{mod_relname}",
                 class_name=cls_name,
             )
-            for model_arch, (mod_relname, cls_name) in _OMNI_MODELS.items()
+            for model_arch, (mod_folder, mod_relname, cls_name) in _OMNI_MODELS.items()
         },
     }
 )

vllm_omni/model_executor/models/utils.py

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

vllm_omni/model_executor/models/vision.py

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

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -15,8 +15,17 @@ stage_args:
       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,6 +43,15 @@ 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
@@ -52,6 +70,14 @@ 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,14 +1,20 @@
 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
+    stage_list,
+    engine_input_source,
+    prompt: Union[OmniTokensPrompt, TextPrompt] = None,
+    requires_multimodal_data: bool = False,
 ):
     if not engine_input_source:
         raise ValueError("engine_input_source cannot be empty")
@@ -19,9 +25,10 @@ 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: prompt.get("multi_modal_data", None)
-        for thinker_output, prompt in zip(thinker_outputs, prompt)
+        thinker_output.request_id: p.get("multi_modal_data", None) for thinker_output, p in zip(thinker_outputs, prompt)
     }
 
     for i, thinker_output in enumerate(thinker_outputs):
@@ -29,27 +36,25 @@ def thinker2talker(
         prompt_token_ids = thinker_output.prompt_token_ids
         thinker_output_ids = output.token_ids
         prompt_token_ids_len = len(prompt_token_ids)
-        thinker_hidden_states = (
-            output.multimodal_output["latent"].clone().detach().cuda()
-        )
+        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),
+        }
         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={
-                    "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),
-                },
+                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,
                 multi_modal_data=(
                     multi_modal_data[thinker_output.request_id]
-                    if multi_modal_data is not None
+                    if requires_multimodal_data and multi_modal_data is not None
                     else None
                 ),
                 mm_processor_kwargs=None,

vllm_omni/outputs.py

@@ -2,7 +2,6 @@ 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,11 +1,10 @@
 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
@@ -14,15 +13,9 @@ 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,11 +6,7 @@ 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):
@@ -25,9 +21,7 @@ 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(
@@ -35,7 +29,6 @@ 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,
@@ -47,9 +40,7 @@ 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

@@ -0,0 +1,12 @@
+# 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__)