Compare commits

..
Author SHA1 Message Date
汪志鹏 12494a2b8c Remove mm_prefix_lm patch because vllm==0.18.0 already support (#2047)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
2026-03-21 21:38:59 +08:00
Zhou Taichang c27f1c4a00 Merge branch 'main' into dev/rebase-v0.18.0 2026-03-21 19:13:12 +08:00
Zhou Taichang e80548767b Merge branch 'main' into dev/rebase-v0.18.0 2026-03-21 15:56:29 +08:00
Zhou Taichang 826b3b361c Merge branch 'main' into dev/rebase-v0.18.0 2026-03-20 20:39:41 +08:00
Zhou Taichang 869a593483 Merge branch 'main' into dev/rebase-v0.18.0 2026-03-20 19:55:26 +08:00
gcanlin 9538e3f30b [NPU] Upgrade to v0.18.0
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-03-20 09:53:34 +00:00
tzhouam 521f178619 [Refactor] Remove unused parameters from AsyncOmni class constructor in async_omni.py
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-20 09:04:35 +00:00
Zhou Taichang 66c31a3502 Merge branch 'main' into dev/rebase-v0.18.0
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
2026-03-20 16:42:17 +08:00
tzhouam 6ff31eb6ec Update rope parameters in Qwen3OmniMoeTalker to include rope_theta for improved configuration handling.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-20 03:42:07 +00:00
tzhouam b8f736c5a9 Remove obsolete image and YAML configuration files for Qwen2.5 and Qwen3 models from the e2e stage configurations.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-20 02:11:18 +00:00
tzhouam 0d23f68997 revert changes in qwen 3 tts
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-20 02:11:10 +00:00
tzhouam 3c52432100 Merge branch 'main' into dev/vllm-align 2026-03-20 01:38:43 +00:00
tzhouam 367ce6328c Merge remote-tracking branch 'remotes/origin/main' into dev/vllm-align 2026-03-19 09:38:17 +00:00
tzhouam 37e055ed9e Merge remote-tracking branch 'remotes/origin/main' into dev/vllm-align 2026-03-19 08:49:35 +00:00
tzhouam f83df010ef Remove redundant assignment of rope_theta in Qwen3OmniMoeTalker initialization for cleaner configuration setup.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-19 08:35:56 +00:00
tzhouam 22b6703773 Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-03-19 05:16:40 +00:00
tzhouam 034d579d77 Update vLLM precompiled wheel version and refactor various components for improved functionality and error handling. Changes include modifying the get_name method in DiffusionQuantizationConfig to use class-level attributes, enhancing metadata handling in StageEngineCoreClient, and adjusting the AsyncOmni class constructor for better parameter management. Additionally, improved error handling in omni_snapshot_download and optimized initial chunk size computation in talker2code2wav_async_chunk for dynamic load adaptation.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-19 05:14:34 +00:00
tzhouam 4f75cccacc Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-03-19 03:11:40 +00:00
tzhouam 342a209924 Merge remote-tracking branch 'origin/main' into dev/vllm-align
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-19 02:06:26 +00:00
Taichang Zhou dc6f2fb81f Enhance cu_seqlens computation in Qwen3Omni_VisionTransformer by switching to tensor operations for improved performance and compatibility. Added error handling to fallback on a vectorized implementation if necessary. This update aligns with upstream changes and optimizes the handling of sequence lengths during attention execution. 2026-03-18 23:27:55 +08:00
Taichang Zhou 7a68bd37b7 Refactor cu_seqlens handling in Qwen3Omni_VisionTransformer to align with upstream vLLM changes. The sequence length helper now retains cu_seqlens in numpy format and converts it for attention execution, improving compatibility and performance. 2026-03-18 23:23:41 +08:00
Taichang Zhou 8a31e7865d Increase timeout for Diffusion Model Test and Diffusion Images API LoRA E2E in Buildkite configurations to 30 minutes. Also, introduce a new timeout constant for diffusion initialization in relevant test files to enhance stability during testing. Update OmniServer initialization to utilize the new timeout settings. 2026-03-17 23:43:46 +08:00
tzhouam 004c60a031 Refactor shutdown and cleanup logic in OmniRunner and OmniStage to improve process termination and resource management. The changes ensure that stage workers are stopped before closing queues, preventing potential orphaned subprocesses. Additionally, enhanced error handling during process termination provides clearer logging for failed shutdown attempts. This update aims to increase the reliability and stability of the shutdown process across the application.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-17 09:33:28 +00:00
tzhouam cc4ade4ba7 Refactor shutdown logic in OmniLLM and stage worker to ensure proper engine closure. The shutdown method now checks for the presence of engine_core before invoking shutdown, enhancing resource management. Additionally, the stage worker now explicitly calls stage_engine.close() upon receiving a shutdown signal. This improves the overall stability and reliability of the shutdown process.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-17 08:19:56 +00:00
tzhouam 35b3dd638a Refactor Dockerfile.ci to install vLLM from a dynamically retrieved precompiled wheel URL, improving the installation process by ensuring compatibility with PEP 440 local version identifiers. This change enhances the reliability of vLLM installations. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-17 06:08:59 +00:00
tzhouam 83effe39ac Update Dockerfile.ci to include an additional extra index URL for pip installation, enhancing package resolution for vLLM. This change ensures compatibility with both precompiled wheels and packages from PyPI.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-17 04:07:50 +00:00
tzhouam b7a36d16d4 Merge branch 'dev/vllm-align' of https://github.com/vllm-project/vllm-omni into dev/vllm-align
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-17 03:54:31 +00:00
tzhouam 722ad6bbae rebase: align vllm-omni with vLLM 0a0a1a198be8
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-17 03:26:42 +00:00
Zhou Taichang ac937b82e6 Refactor Dockerfile to install vLLM from source
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
2026-03-17 11:08:45 +08:00
tzhouam 0eccb77c4f Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-03-17 02:48:50 +00:00
tzhouam b0e36f9f5d Update Dockerfile.ci to force reinstall vLLM from a precompiled wheel, ensuring the latest version is used. This change enhances the installation process by guaranteeing that any existing installations are replaced. Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-17 02:04:51 +00:00
tzhouam 83b05ce3e4 Refactor Dockerfile.ci to install vLLM using a precompiled wheel from a specified commit, replacing the previous source installation method. This update simplifies the installation process and ensures correct dependency management. Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-17 01:35:12 +00:00
Zhou Taichang 22a2bd3406 Upgrade flashinfer-cubin and flashinfer-jit-cache versions
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
2026-03-16 19:31:03 +08:00
tzhouam 8e81e60d32 Refactor Dockerfile.ci to install vLLM from source using precompiled artifacts at a specified commit. This change replaces the previous wheel installation method and ensures the correct versioning of dependencies. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-16 10:01:18 +00:00
tzhouam 57ac7eb155 rebase: align vllm-omni with vLLM 2754231ba3a7
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-16 09:28:56 +00:00
tzhouam de6bcc734c Update Dockerfile.ci to use a new vllm wheel URL and refactor MooncakeConnector import logic in monkey_patch.py for improved compatibility. Adjust OpenAIServingChat to call the harmony request method from the correct context.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-16 08:33:28 +00:00
tzhouam a26e568dfc Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-03-16 07:30:43 +00:00
tzhouam 9c22e3afb0 Refactor import statement for OpenAIServingEmbedding in api_server.py to use aliasing for clarity
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-16 06:13:07 +00:00
tzhouam aae101f495 Refactor import statement for OpenAIServingEmbedding in api_server.py
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-16 03:39:51 +00:00
tzhouam 3a508d7b7a rebase: align vllm-omni with vLLM e9163b536e72
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-16 02:32:39 +00:00
tzhouam f137a1e3ad Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-03-16 01:49:46 +00:00
tzhouam 97899c6c98 Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-03-13 19:03:57 +00:00
tzhouam d3db476f2d Merge branch 'main' into dev/vllm-align 2026-03-13 06:51:22 +00:00
tzhouam 62c8695dfd Update GPUGenerationModelRunner to clarify KV connector metadata handling during draft model runs
- Changed variable name to better reflect its purpose in managing KV connector metadata.
- Updated comments for improved clarity regarding the handling of speculative configurations.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-13 03:16:33 +00:00
tzhouam d911475b22 Fix output queue handling in OmniStage to return None if queue is closed during shutdown
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-13 01:59:04 +00:00
tzhouam e17fb0e97e Refactor OpenAI serving components and enhance multimodal handling
- Updated import for OpenAIServingEmbedding in api_server.py.
- Enhanced omni_init_app_state to initialize renderer for engine_client.
- Added handle_oov_mm_token parameter to multiple model classes for better multimodal token handling.
- Improved comments for clarity in various model files.
- Adjusted GPUARModelRunner to ensure proper handling of late interaction runner attributes.

Signed-off-by: [Your Name] <[Your Email]>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-13 01:55:17 +00:00
tzhouam 72cb4f661b Adjust GPU memory utilization and device settings in qwen2_5_omni_ci.yaml for improved performance
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-12 13:57:52 +00:00
tzhouam 11b6c16007 rebase: align vllm-omni with vLLM 84e436ed1c94
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-11 19:24:44 +00:00
tzhouam 08a2673c66 rebase: align vllm-omni with vLLM 84e436ed1c94
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-11 07:29:52 +00:00
tzhouam e47c5dd7fb Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-03-11 02:50:45 +00:00
tzhouam 91135f13aa Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-03-11 02:43:59 +00:00
tzhouam ad12921cc2 Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-03-11 01:44:32 +00:00
tzhouam 20c102d3cf debug: empty commit to test CI pipeline (vLLM ddbb0d230a35)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 10:04:29 +00:00
tzhouam a9862f17f3 Revert "Refactor Dockerfile.ci to consolidate package installation commands into a single RUN statement, improving readability and efficiency of the CI environment setup."
This reverts commit 1c0a71e30d.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 09:41:45 +00:00
tzhouam 1c0a71e30d Refactor Dockerfile.ci to consolidate package installation commands into a single RUN statement, improving readability and efficiency of the CI environment setup.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 09:13:24 +00:00
tzhouam cee2b4bd11 debug: empty commit to test CI pipeline (vLLM ddbb0d230a35)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 08:58:06 +00:00
tzhouam b0cf788d8f debug: empty commit to test CI pipeline (vLLM ddbb0d230a35)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 08:56:23 +00:00
tzhouam dd0ec557b1 debug: empty commit to test CI pipeline (vLLM ddbb0d230a35)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 08:55:56 +00:00
tzhouam 4014e76795 debug: empty commit to test CI pipeline (vLLM ddbb0d230a35)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 08:54:54 +00:00
tzhouam 5c845ce50b debug: empty commit to test CI pipeline (vLLM ddbb0d230a35)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 08:54:06 +00:00
tzhouam 9dc5d15f1d debug: empty commit to test CI pipeline (vLLM ddbb0d230a35)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 08:09:53 +00:00
tzhouam cf0a7a5475 Update Dockerfile.ci to install additional dependencies for flashinfer and numpy, and create a symlink for python3. This enhances the CI environment setup.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 08:05:28 +00:00
tzhouam 8ac369c84a debug: empty commit to test CI pipeline (vLLM 156e33553ccd)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 07:13:09 +00:00
tzhouam cfbdb57ec6 debug: empty commit to test CI pipeline (vLLM 156e33553ccd)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 07:11:31 +00:00
tzhouam 7706132d8a debug: empty commit to test CI pipeline (vLLM 156e33553ccd)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 07:08:42 +00:00
tzhouam e150a1b28b debug: empty commit to test CI pipeline (vLLM 156e33553ccd)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 06:55:27 +00:00
tzhouam 3630bdfaa7 Update Dockerfile.ci to install vLLM from a new commit and simplify worker_type retrieval in omni_stage.py. Removed unnecessary dependency installations for flashinfer and numpy.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 01:55:12 +00:00
tzhouam dd0c893e79 Align dev/vllm-align with upstream vLLM changes
- Update Dockerfile.ci to install vLLM from specific commit wheel,
  add flashinfer/cublas/numpy dependencies
- Fix worker_type leak in omni_stage by using pop instead of get

Made-with: Cursor

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-10 01:46:53 +00:00
425 changed files with 7052 additions and 40645 deletions

No files matched your search

-20
View File
@@ -1,20 +0,0 @@
steps:
- label: "Build and upload wheel"
key: "build-wheel"
agents:
queue: cpu_queue_release
commands:
- "curl -LsSf https://astral.sh/uv/install.sh | sh"
- 'export PATH="$HOME/.local/bin:$PATH"'
- "uv venv --python=3.12 && source .venv/bin/activate"
- "uv pip install --upgrade build"
- "python3 -m build"
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
- label: "Generate and upload wheel indices"
depends_on: "build-wheel"
allow_dependency_failure: true
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
+1 -1
View File
@@ -10,7 +10,7 @@ steps:
DOCKER_BUILDKIT: "1"
# Buildkite will automatically replace this with the actual commit hash
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
VLLM_VERSION: "v0.18.0"
VLLM_VERSION: "v0.17.0"
priority: 100
timeout_in_minutes: 60
soft_fail: true
+4 -4
View File
@@ -13,7 +13,7 @@ steps:
- label: "Upload Ready Pipeline"
depends_on: image-build
key: upload-ready-pipeline
if: build.branch != "main" && build.pull_request.labels includes "ready"
if: build.branch != "main"
commands:
- buildkite-agent pipeline upload .buildkite/test-ready.yml
agents:
@@ -23,17 +23,17 @@ steps:
- label: "Upload Merge Pipeline"
depends_on: image-build
key: upload-merge-pipeline
if: build.branch == "main" && build.env("NIGHTLY") != "1"
if: build.branch == "main"
commands:
- buildkite-agent pipeline upload .buildkite/test-merge.yml
agents:
queue: "cpu_queue_premerge"
# L4 Test — main+NIGHTLY=1 (scheduled), or PR with label nightly-test (e.g. add label then Rebuild)
# L4 Test
- label: "Upload Nightly Pipeline"
depends_on: image-build
key: upload-nightly-pipeline
if: '(build.branch == "main" && build.env("NIGHTLY") == "1") || (build.branch != "main" && build.pull_request.labels includes "nightly-test")'
if: build.env("NIGHTLY") == "1"
commands:
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
agents:
@@ -1,87 +0,0 @@
#!/usr/bin/env bash
set -ex
# Generate and upload wheel indices for all vllm-omni wheels in the commit directory.
# This script should run once after all wheels have been built and uploaded.
# All paths are under the omni/ prefix in the vllm-wheels S3 bucket.
# ======== setup ========
BUCKET="vllm-wheels"
INDICES_OUTPUT_DIR="indices"
PYTHON="${PYTHON_PROG:-python3}"
SUBPATH="omni/$BUILDKITE_COMMIT"
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
# detect if python3.12+ is available
has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
if [[ "$has_new_python" -eq 0 ]]; then
# use new python from docker
docker pull python:3-slim
PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3"
fi
echo "Using python interpreter: $PYTHON"
echo "Python version: $($PYTHON --version)"
# ======== generate and upload indices ========
# list all wheels in the commit directory
echo "Existing wheels on S3:"
aws s3 ls "$S3_COMMIT_PREFIX" || echo "(no objects found)"
obj_json="objects.json"
aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$SUBPATH/" --delimiter / --output json > "$obj_json"
mkdir -p "$INDICES_OUTPUT_DIR"
# HACK: we do not need regex module here, but it is required by pre-commit hook
# To avoid any external dependency, we simply replace it back to the stdlib re module
sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
# Generate indices -- the version is just the commit hash (not omni/{commit})
# because relative paths are computed between the index and wheel directories,
# both of which live under the omni/ prefix in S3.
$PYTHON .buildkite/scripts/generate-nightly-index.py \
--version "$BUILDKITE_COMMIT" \
--current-objects "$obj_json" \
--output-dir "$INDICES_OUTPUT_DIR" \
--comment "commit $BUILDKITE_COMMIT"
# copy indices to /omni/{commit}/ unconditionally
echo "Uploading indices to $S3_COMMIT_PREFIX"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
# copy to /omni/nightly/ when NIGHTLY=1
if [[ "${NIGHTLY:-}" == "1" ]]; then
echo "Uploading indices to overwrite /omni/nightly/"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/nightly/"
fi
# detect version from any wheel in the commit directory
first_wheel_key=$($PYTHON -c "import json; obj=json.load(open('$obj_json')); print(next((c['Key'] for c in obj.get('Contents', []) if c['Key'].endswith('.whl')), ''))")
if [[ -z "$first_wheel_key" ]]; then
echo "Error: No wheels found in $S3_COMMIT_PREFIX"
exit 1
fi
first_wheel=$(basename "$first_wheel_key")
aws s3 cp "s3://$BUCKET/${first_wheel_key}" "/tmp/${first_wheel}"
version=$(unzip -p "/tmp/${first_wheel}" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
rm -f "/tmp/${first_wheel}"
echo "Version in wheel: $version"
pure_version="${version%%+*}"
echo "Pure version (without variant): $pure_version"
# re-generate and copy to /omni/{version}/ only if it does not have "dev" in the version
if [[ "$version" != *"dev"* ]]; then
echo "Re-generating indices for /omni/$pure_version/"
rm -rf "${INDICES_OUTPUT_DIR:?}"
mkdir -p "$INDICES_OUTPUT_DIR"
# wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
$PYTHON .buildkite/scripts/generate-nightly-index.py \
--version "$pure_version" \
--wheel-dir "$BUILDKITE_COMMIT" \
--current-objects "$obj_json" \
--output-dir "$INDICES_OUTPUT_DIR" \
--comment "version $pure_version"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/$pure_version/"
fi
@@ -1,193 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import json
import sys
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from typing import Any
from urllib.parse import quote
import regex as re
def normalize_package_name(name: str) -> str:
"""Normalize package name per PEP 503."""
return re.sub(r"[-_.]+", "-", name).lower()
if not sys.version_info >= (3, 12):
raise RuntimeError("This script requires Python 3.12 or higher.")
INDEX_HTML_TEMPLATE = """<!DOCTYPE html>
<html>
<!-- {comment} -->
<meta name="pypi:repository-version" content="1.0">
<body>
{items}
</body>
</html>
"""
@dataclass
class WheelFileInfo:
package_name: str
version: str
build_tag: str | None
python_tag: str
abi_tag: str
platform_tag: str
filename: str
def parse_from_filename(file: str) -> WheelFileInfo:
"""
Parse wheel filename per PEP 427:
{package_name}-{version}(-{build_tag})?-{python_tag}-{abi_tag}-{platform_tag}.whl
"""
wheel_file_re = re.compile(
r"^(?P<package_name>.+)-(?P<version>[^-]+?)(-(?P<build_tag>[^-]+))?-(?P<python_tag>[^-]+)-(?P<abi_tag>[^-]+)-(?P<platform_tag>[^-]+)\.whl$"
)
match = wheel_file_re.match(file)
if not match:
raise ValueError(f"Invalid wheel file name: {file}")
return WheelFileInfo(
package_name=match.group("package_name"),
version=match.group("version"),
build_tag=match.group("build_tag"),
python_tag=match.group("python_tag"),
abi_tag=match.group("abi_tag"),
platform_tag=match.group("platform_tag"),
filename=file,
)
def generate_project_list(package_names: list[str], comment: str = "") -> str:
"""Generate top-level PEP 503 project list HTML."""
href_tags = []
for name in sorted(package_names):
href_tags.append(f' <a href="{name}/">{name}/</a><br/>')
return INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
def generate_package_index(
wheel_files: list[WheelFileInfo],
wheel_base_dir: Path,
index_base_dir: Path,
comment: str = "",
) -> tuple[str, str]:
"""Generate package index HTML and metadata JSON linking to wheel files."""
href_tags = []
metadata = []
for file in sorted(wheel_files, key=lambda x: x.filename):
relative_path = wheel_base_dir.relative_to(index_base_dir, walk_up=True) / file.filename
# handle '+' in URL; avoid double-encoding '/' and '%2B' (AWS S3 behavior)
file_path_quoted = quote(relative_path.as_posix(), safe=":%/")
href_tags.append(f' <a href="{file_path_quoted}">{file.filename}</a><br/>')
file_meta = asdict(file)
file_meta["path"] = file_path_quoted
metadata.append(file_meta)
index_str = INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
metadata_str = json.dumps(metadata, indent=2)
return index_str, metadata_str
def generate_index(
whl_files: list[str],
wheel_base_dir: Path,
index_base_dir: Path,
comment: str = "",
):
"""
Generate PEP 503 index for all wheel files.
Output structure:
index_base_dir/
index.html # project list linking to vllm-omni/
vllm-omni/
index.html # package index linking to wheel files
metadata.json # machine-readable metadata
"""
parsed_files = [parse_from_filename(f) for f in whl_files]
if not parsed_files:
print("No wheel files found, skipping index generation.")
return
comment_str = f" ({comment})" if comment else ""
comment_tmpl = f"Generated on {datetime.now().isoformat()}{comment_str}"
# Group by normalized package name
packages: dict[str, list[WheelFileInfo]] = {}
for file in parsed_files:
name = normalize_package_name(file.package_name)
packages.setdefault(name, []).append(file)
print(f"Found packages: {list(packages.keys())}")
# Generate per-package index
for package, files in packages.items():
package_dir = index_base_dir / package
package_dir.mkdir(parents=True, exist_ok=True)
index_str, metadata_str = generate_package_index(files, wheel_base_dir, package_dir, comment)
with open(package_dir / "index.html", "w") as f:
f.write(index_str)
with open(package_dir / "metadata.json", "w") as f:
f.write(metadata_str)
# Generate top-level project list
project_list_str = generate_project_list(sorted(packages.keys()), comment_tmpl)
with open(index_base_dir / "index.html", "w") as f:
f.write(project_list_str)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Generate PEP 503 wheel index from S3 object listing.")
parser.add_argument("--version", type=str, required=True, help="Version string (e.g., commit hash)")
parser.add_argument("--current-objects", type=str, required=True, help="Path to JSON from S3 list-objects-v2")
parser.add_argument("--output-dir", type=str, required=True, help="Directory to write index files")
parser.add_argument("--wheel-dir", type=str, default=None, help="Wheel directory (defaults to --version)")
parser.add_argument("--comment", type=str, default="", help="Comment for generated HTML")
args = parser.parse_args()
version = args.version
if "\\" in version or "/" in version:
raise ValueError("Version string must not contain slashes or backslashes.")
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
with open(args.current_objects) as f:
current_objects: dict[str, list[dict[str, Any]]] = json.load(f)
wheel_files = [
item["Key"].split("/")[-1] for item in current_objects.get("Contents", []) if item["Key"].endswith(".whl")
]
print(f"Found {len(wheel_files)} wheel files for version {version}: {wheel_files}")
# For release versions, filter to only matching non-dev wheels
PY_VERSION_RE = re.compile(r"^\d+\.\d+\.\d+([a-zA-Z0-9.+-]*)?$")
if PY_VERSION_RE.match(version):
wheel_files = [f for f in wheel_files if version in f and "dev" not in f]
print(f"Non-nightly version detected, wheel files used: {wheel_files}")
else:
print("Nightly version detected, keeping all wheel files.")
wheel_dir = (args.wheel_dir or version).strip().rstrip("/")
wheel_base_dir = Path(output_dir).parent / wheel_dir
index_base_dir = Path(output_dir)
generate_index(
whl_files=wheel_files,
wheel_base_dir=wheel_base_dir,
index_base_dir=index_base_dir,
comment=args.comment.strip(),
)
print(f"Successfully generated index in {output_dir}")
+11 -14
View File
@@ -64,22 +64,15 @@ while true; do
done
echo "--- Pulling container"
## Temporary change to use AMD Docker Hub to store the vllm-ci image
# to bypass the rate limit issue with ECR Public Gallery.
# TODO: @tjtanaa point back to ECR Public Gallery
# once the amd agents are configured to use ECR Public Gallery.
# image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}-rocm-omni"
image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
# TODO: @tjtanaa uncomment this once the amd agents are configured to use ECR Public Gallery.
# # Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
# sudo apt-get update && sudo apt-get install -y awscli
## Use safe docker login helper to prevent race conditions
# source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
# safe_docker_login_ecr_public
## Pull the container from AMD Docker Hub
# Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
sudo apt-get update && sudo apt-get install -y awscli
# Use safe docker login helper to prevent race conditions
source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
safe_docker_login_ecr_public
# Pull the container from ECR Public Gallery
docker pull "${image_name}"
@@ -133,6 +126,8 @@ if [[ $commands == *"--shard-id="* ]]; then
-e VLLM_ROCM_USE_AITER=0 \
-e HIP_VISIBLE_DEVICES="${GPU}" \
-e HF_TOKEN \
-e AWS_ACCESS_KEY_ID \
-e AWS_SECRET_ACCESS_KEY \
-v "${HF_CACHE}:${HF_MOUNT}" \
-e "HF_HOME=${HF_MOUNT}" \
-e "PYTHONPATH=${MYPYTHONPATH}" \
@@ -165,6 +160,8 @@ else
-e MIOPEN_DEBUG_CONV_GEMM=0 \
-e VLLM_ROCM_USE_AITER=0 \
-e HF_TOKEN \
-e AWS_ACCESS_KEY_ID \
-e AWS_SECRET_ACCESS_KEY \
-v "${HF_CACHE}:${HF_MOUNT}" \
-e "HF_HOME=${HF_MOUNT}" \
-e "PYTHONPATH=${MYPYTHONPATH}" \
@@ -1,33 +0,0 @@
#!/usr/bin/env bash
set -ex
# Upload a single wheel to S3 under the omni/ prefix.
# Index generation is handled separately by generate-and-upload-nightly-index.sh.
BUCKET="vllm-wheels"
SUBPATH="omni/$BUILDKITE_COMMIT"
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
# ========= collect & upload the wheel ==========
# python3 -m build outputs to dist/ by default
wheel_files=(dist/*.whl)
# Check that exactly one wheel is found
if [[ ${#wheel_files[@]} -ne 1 ]]; then
echo "Error: Expected exactly one wheel file in dist/, but found ${#wheel_files[@]}"
exit 1
fi
wheel="${wheel_files[0]}"
echo "Uploading wheel: $wheel"
# Extract the version from the wheel
version=$(unzip -p "$wheel" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
echo "Version in wheel: $version"
# Upload wheel to S3
aws s3 cp "$wheel" "$S3_COMMIT_PREFIX"
echo "Wheel uploaded to $S3_COMMIT_PREFIX. Index generation is handled by a separate step."
+4 -22
View File
@@ -134,33 +134,15 @@ steps:
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
- label: "Qwen3-TTS CustomVoice E2E Test"
- label: "Qwen3-TTS E2E Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
'
- label: "Qwen3-TTS Base E2E Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_base.py
'
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
- label: "Diffusion Image Edit Test"
agent_pool: mi325_1
+1 -2
View File
@@ -124,8 +124,7 @@ steps:
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
- label: "Diffusion Image Edit Test"
agent_pool: mi325_1
+39 -89
View File
@@ -112,7 +112,6 @@ steps:
depends_on: upload-merge-pipeline
commands:
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
- pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
@@ -188,13 +187,12 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test"
timeout_in_minutes: 30
timeout_in_minutes: 15
depends_on: upload-merge-pipeline
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
- pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
@@ -208,52 +206,6 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Qwen3-TTS CustomVoice E2E Test"
depends_on: upload-merge-pipeline
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
'
agents:
queue: "gpu_1_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Qwen3-TTS Base E2E Test"
depends_on: upload-merge-pipeline
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_base.py
'
agents:
queue: "gpu_1_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test with H100"
timeout_in_minutes: 30
depends_on: upload-merge-pipeline
@@ -262,7 +214,6 @@ steps:
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
- pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
@@ -385,42 +336,41 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Voxtral-TTS E2E Test"
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
# - label: "Voxtral TTS E2E Test"
# timeout_in_minutes: 20
# depends_on: upload-merge-pipeline
# commands:
# - |
# timeout 20m bash -c '
# export VLLM_LOGGING_LEVEL=DEBUG
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
# pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
# '
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# resources:
# limits:
# nvidia.com/gpu: 1
# volumeMounts:
# - name: devshm
# mountPath: /dev/shm
# - name: hf-cache
# mountPath: /root/.cache/huggingface
# env:
# - name: HF_HOME
# value: /root/.cache/huggingface
# nodeSelector:
# node.kubernetes.io/instance-type: gpu-h100-sxm
# volumes:
# - name: devshm
# emptyDir:
# medium: Memory
# - name: hf-cache
# hostPath:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
+69 -293
View File
@@ -2,7 +2,7 @@ steps:
- label: ":full_moon: Omni Model Test with H100"
timeout_in_minutes: 90
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
@@ -11,11 +11,10 @@ steps:
EXIT1=$$?
pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
EXIT2=$$?
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
EXIT3=$$?
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
EXIT4=$$?
exit $$((EXIT1 | EXIT2 | EXIT3 | EXIT4))
# TODO: uncomment when mistralai/tts-model is public
# pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
# EXIT3=$$?
exit $$((EXIT1 | EXIT2))
agents:
queue: "mithril-h100-pool"
plugins:
@@ -50,15 +49,13 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Omni Model Test with L4"
timeout_in_minutes: 90
- label: ":full_moon: Omni Model Test"
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/examples/ -m "advanced_model and L4 and omni" --run-level "advanced_model"
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
- pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
@@ -74,96 +71,57 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: ":full_moon: Diffusion Model Test with H100"
timeout_in_minutes: 120
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion Model Wan22 completed Test with H100"
timeout_in_minutes: 90
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -m "advanced_model" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion Model Test with L4"
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model"
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
timeout_in_minutes: 30
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
- |
huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk
agents:
queue: "gpu_4_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -176,62 +134,18 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: ":full_moon: Documentation Example Code Test with H100"
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/examples/online_serving/test_text_to_image.py tests/examples/offline_inference/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Omni Model Perf Test & Test Case Statistics"
key: nightly-omni-performance
key: nightly-performance
timeout_in_minutes: 180
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export BENCHMARK_DIR=tests/dfx/perf/results
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
- buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
- python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
- export BENCHMARK_DIR=tests
- pytest -s -v tests/perf/scripts/run_benchmark.py
- buildkite-agent artifact upload "tests/*.json"
- python tools/nightly/buildkite_testcase_statistics.py -o tests/buildkite_testcase_statistics.html
- buildkite-agent artifact upload "tests/*.html"
agents:
queue: "mithril-h100-pool"
plugins:
@@ -247,93 +161,6 @@ steps:
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: GEBench Accuracy Test with H100"
key: nightly-gebench-accuracy
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: GEdit-Bench Accuracy Test with H100"
key: nightly-gedit-bench-accuracy
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
@@ -355,70 +182,19 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Qwen-Image Diffusion Perf Test with H100"
key: nightly-qwen-image-performance
timeout_in_minutes: 300
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export CACHE_DIT_VERSION=1.3.0
- pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
- buildkite-agent artifact upload "tests/dfx/perf/results/benchmark_results_*.json"
- buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 4
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
# No need to run this step for PRs with label nightly-test
- label: ":email: Nightly Collection & Email"
key: nightly-perf-distribution
depends_on:
- nightly-omni-performance
- nightly-qwen-image-performance
- nightly-performance
if: build.env("NIGHTLY") == "1"
commands:
- pip install openpyxl
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
- export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-qwen-image-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-omni-performance
- export DEFAULT_INPUT_DIR=tests
- export DEFAULT_OUTPUT_DIR=tests
- buildkite-agent artifact download "tests/*.json" . --step nightly-performance
- buildkite-agent artifact download "tests/*.html" . --step nightly-performance
- python tools/nightly/generate_nightly_perf_excel.py
- python tools/nightly/generate_nightly_perf_html.py
- python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html"
- buildkite-agent artifact upload "tests/dfx/perf/results/*.xlsx"
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
- python tools/nightly/send_nightly_email.py --report-file "tests/*.xlsx, tests/*.html"
- buildkite-agent artifact upload "tests/*.xlsx"
agents:
queue: "cpu_queue_premerge"
+42 -119
View File
@@ -50,39 +50,6 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Batching Test"
depends_on: upload-ready-pipeline
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Custom Pipeline Test"
depends_on: upload-ready-pipeline
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
agents:
queue: "gpu_1_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Model CPU offloading Test"
depends_on: upload-ready-pipeline
commands:
@@ -201,7 +168,7 @@ steps:
timeout 17m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "core_model" --run-level "core_model"
pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
'
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
@@ -255,61 +222,18 @@ steps:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
- label: "MiMo-Audio E2E Test with H100"
depends_on: upload-ready-pipeline
commands:
- |
timeout 30m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Qwen3-TTS E2E Test"
depends_on: upload-ready-pipeline
commands:
- |
timeout 20m bash -c '
huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py -k "not NoAsyncChunk"
'
agents:
queue: "gpu_1_queue"
queue: "gpu_4_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -322,45 +246,44 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Voxtral-TTS E2E Test"
timeout_in_minutes: 20
depends_on: upload-ready-pipeline
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
# TODO: uncomment when mistralai/tts-model is public
# - label: "Voxtral-TTS E2E Test"
# depends_on: upload-ready-pipeline
# commands:
# - |
# timeout 20m bash -c '
# export VLLM_LOGGING_LEVEL=DEBUG
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
# pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
# '
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# resources:
# limits:
# nvidia.com/gpu: 1
# volumeMounts:
# - name: devshm
# mountPath: /dev/shm
# - name: hf-cache
# mountPath: /root/.cache/huggingface
# env:
# - name: HF_HOME
# value: /root/.cache/huggingface
# nodeSelector:
# node.kubernetes.io/instance-type: gpu-h100-sxm
# volumes:
# - name: devshm
# emptyDir:
# medium: Memory
# - name: hf-cache
# hostPath:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
# - label: "Diffusion Image Edit Test with H100 (1 GPU)"
# depends_on: upload-ready-pipeline
+5 -6
View File
@@ -3,7 +3,7 @@
Last synced: 2025-12-15
Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests
#}
{% set docker_image_amd = "rocm/vllm-ci:$BUILDKITE_COMMIT-rocm-omni" %}
{% set docker_image_amd = "public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT-rocm-omni" %}
{% set default_working_dir = "/app/vllm-omni" %}
- group: "AMD Tests"
@@ -13,6 +13,7 @@
depends_on: ~
soft_fail: false
commands:
- "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
- "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target test --build-arg ARG_PYTORCH_ROCM_ARCH=gfx942 --progress plain ."
- "docker push {{ docker_image_amd }}"
key: "amd-build"
@@ -21,15 +22,13 @@
retry:
automatic:
- exit_status: -1 # Agent was lost
limit: 2
limit: 1
- exit_status: -10 # Agent was lost
limit: 2
- exit_status: 128 # Git connectivity issues
limit: 2
limit: 1
- exit_status: 1 # Machine occasionally fail
limit: 1
agents:
queue: amd-cpu
queue: cpu_queue_premerge
{% for step in steps %}
{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
+3 -3
View File
@@ -39,10 +39,10 @@ jobs:
python-version: ["3.11", "3.12"]
steps:
- uses: actions/checkout@v6
- uses: actions/checkout@v4
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v6
uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
@@ -57,7 +57,7 @@ jobs:
bash scripts/build_wheel.sh --python python
- name: Upload Wheel Artifact
uses: actions/upload-artifact@v7
uses: actions/upload-artifact@v4
with:
name: vllm-omni-wheel-py${{ matrix.python-version }}
path: dist/*.whl
+3 -3
View File
@@ -24,12 +24,12 @@ jobs:
github.event.comment.author_association == 'OWNER')
steps:
- name: Checkout repository
uses: actions/checkout@v6.0.2
uses: actions/checkout@v4.2.2
with:
fetch-depth: 0
- name: Set up Python
uses: actions/setup-python@v6.2.0
uses: actions/setup-python@v5.3.0
with:
python-version: '3.11'
cache: 'pip'
@@ -54,7 +54,7 @@ jobs:
- name: Upload review logs
if: always()
uses: actions/upload-artifact@v7.0.0
uses: actions/upload-artifact@v4.4.3
with:
name: pr-review-logs-${{ github.event.issue.number || github.event.pull_request.number }}
path: /tmp/pr_review_*.log
+2 -2
View File
@@ -18,8 +18,8 @@ jobs:
pre-commit:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- uses: actions/setup-python@v6
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: "3.12"
cache: 'pip'
-2
View File
@@ -15,8 +15,6 @@ This document explains the benchmark architecture across all benchmark assets in
## Directory layout
- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
- `benchmarks/accuracy/` — accuracy benchmark integrations that adapt external
benchmark suites to vLLM-Omni serving and evaluation flows.
- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
## Reference workflows
-1
View File
@@ -1 +0,0 @@
"""Benchmark helpers and runnable benchmark entrypoints."""
-27
View File
@@ -1,27 +0,0 @@
# Accuracy Benchmarks
This directory hosts accuracy benchmark integrations that run entirely through a
local `vllm-omni serve` deployment.
Current integrations:
- `text_to_image/`: GEBench generation + local judge scoring flow.
- `image_to_image/`: GEdit-Bench generation + local VIEScore-style scoring flow.
Design notes:
- Generation is executed through the OpenAI-compatible endpoints exposed by
`vllm-omni serve`.
- Evaluation is also executed through a local OpenAI-compatible judge model
served by `vllm-omni`.
- Both generation and judge requests accept either `http://host:port` or
`http://host:port/v1`.
- Output directory layout intentionally stays close to the upstream repos.
Test guidance:
- Local static/self-checks live in `tests/benchmarks/test_accuracy_bench_utils.py`.
- End-to-end generation/evaluation should be validated in a remote GPU
environment. In the current repo marker system there is `L4` but no `L5`
marker, so benchmark smoke tests should be wired as `advanced_model +
benchmark + L4` when GPU capacity is available.
-1
View File
@@ -1 +0,0 @@
"""Accuracy benchmark integrations for vLLM-Omni."""
-199
View File
@@ -1,199 +0,0 @@
from __future__ import annotations
import base64
import io
import json
from pathlib import Path
from typing import Any
import requests
from PIL import Image
def ensure_dir(path: Path) -> Path:
path.mkdir(parents=True, exist_ok=True)
return path
def load_json(path: Path) -> dict[str, Any]:
with path.open("r", encoding="utf-8") as handle:
return json.load(handle)
def write_json(path: Path, payload: dict[str, Any]) -> None:
ensure_dir(path.parent)
with path.open("w", encoding="utf-8") as handle:
json.dump(payload, handle, indent=2, ensure_ascii=False)
def save_image(path: Path, image: Image.Image) -> None:
ensure_dir(path.parent)
image.save(path)
def find_first_image(folder: Path, stem: str | None = None) -> Path | None:
patterns = [f"{stem}.*"] if stem else ["*.png", "*.jpg", "*.jpeg", "*.webp"]
for pattern in patterns:
for candidate in sorted(folder.glob(pattern)):
if candidate.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"}:
return candidate
return None
def extract_json_object(raw_text: str) -> dict[str, Any]:
raw_text = raw_text.strip()
delimiter = "||V^=^V||"
if raw_text.count(delimiter) >= 2:
start = raw_text.find(delimiter) + len(delimiter)
end = raw_text.rfind(delimiter)
raw_text = raw_text[start:end].strip()
start = raw_text.find("{")
end = raw_text.rfind("}")
if start == -1 or end == -1 or end < start:
raise ValueError(f"Could not find JSON object in: {raw_text[:200]}")
return json.loads(raw_text[start : end + 1])
def build_openai_url(base_url: str, api_path: str) -> str:
base = base_url.rstrip("/")
normalized_path = api_path if api_path.startswith("/") else f"/{api_path}"
if base.endswith(normalized_path):
return base
if base.endswith("/v1"):
return f"{base}{normalized_path}"
return f"{base}/v1{normalized_path}"
def pil_to_base64(image: Image.Image, image_format: str = "PNG") -> str:
buffer = io.BytesIO()
image.save(buffer, format=image_format)
return base64.b64encode(buffer.getvalue()).decode("utf-8")
def pil_to_data_url(image: Image.Image, image_format: str = "PNG") -> str:
return f"data:image/{image_format.lower()};base64,{pil_to_base64(image, image_format=image_format)}"
def decode_base64_image(encoded: str) -> Image.Image:
image = Image.open(io.BytesIO(base64.b64decode(encoded)))
image.load()
return image.convert("RGB")
def pil_to_png_bytes(image: Image.Image) -> bytes:
buffer = io.BytesIO()
image.save(buffer, format="PNG")
return buffer.getvalue()
class VllmOmniImageClient:
"""Thin OpenAI-compatible image client for vLLM-Omni serving."""
def __init__(self, base_url: str, api_key: str = "EMPTY", timeout: int = 600):
self.base_url = base_url.rstrip("/")
self.api_key = api_key
self.timeout = timeout
@property
def _headers(self) -> dict[str, str]:
return {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
}
def generate_text_to_image(
self,
*,
model: str,
prompt: str,
width: int,
height: int,
num_inference_steps: int = 20,
guidance_scale: float | None = None,
seed: int | None = None,
output_compression: int | None = None,
) -> Image.Image:
payload: dict[str, Any] = {
"model": model,
"prompt": prompt,
"n": 1,
"size": f"{width}x{height}",
"response_format": "b64_json",
"num_inference_steps": num_inference_steps,
}
if guidance_scale is not None:
payload["guidance_scale"] = guidance_scale
if seed is not None:
payload["seed"] = seed
if output_compression is not None:
payload["output_compression"] = output_compression
response = requests.post(
build_openai_url(self.base_url, "/images/generations"),
json=payload,
headers=self._headers,
timeout=self.timeout,
)
response.raise_for_status()
return decode_base64_image(response.json()["data"][0]["b64_json"])
def generate_image_edit(
self,
*,
model: str,
prompt: str,
images: Image.Image | list[Image.Image],
width: int,
height: int,
num_inference_steps: int = 20,
guidance_scale: float | None = None,
seed: int | None = None,
negative_prompt: str | None = None,
output_compression: int | None = None,
) -> Image.Image:
if not isinstance(images, list):
images = [images]
data: dict[str, Any] = {
"model": model,
"prompt": prompt,
"n": 1,
"size": f"{width}x{height}",
"response_format": "b64_json",
"num_inference_steps": str(num_inference_steps),
}
if guidance_scale is not None:
data["guidance_scale"] = str(guidance_scale)
if seed is not None:
data["seed"] = str(seed)
if negative_prompt:
data["negative_prompt"] = negative_prompt
if output_compression is not None:
data["output_compression"] = str(output_compression)
files = [
(
"image[]" if len(images) > 1 else "image",
(f"image_{index}.png", pil_to_png_bytes(image), "image/png"),
)
for index, image in enumerate(images)
]
edit_paths = ["/images/edits", "/images/edit"]
last_response: requests.Response | None = None
for api_path in edit_paths:
response = requests.post(
build_openai_url(self.base_url, api_path),
data=data,
files=files,
headers={"Authorization": f"Bearer {self.api_key}"},
timeout=self.timeout,
)
last_response = response
if response.status_code != 404:
response.raise_for_status()
return decode_base64_image(response.json()["data"][0]["b64_json"])
assert last_response is not None
last_response.raise_for_status()
raise ValueError("No image payload returned from image edit endpoint")
@@ -1,103 +0,0 @@
# GEdit-Bench on vLLM-Omni
This integration adapts the upstream `stepfun-ai/Step1X-Edit/GEdit-Bench`
evaluation flow into `vllm-omni/benchmarks/accuracy/image_to_image`.
Upstream mapping:
- `run_gedit_score.py` -> `run_gedit_bench.py evaluate`
- `calculate_statistics.py` -> `run_gedit_bench.py summarize`
- upstream output layout under `results/<method>/fullset/...` is preserved
What changed:
- The upstream repo mainly ships evaluation scripts. This integration adds a
generation runner that uses the local `vllm-omni` OpenAI-compatible
`/v1/images/edits` endpoint to produce benchmark outputs in the expected
directory structure.
- The evaluator keeps the same VIEScore-style decomposition:
- `semantics_score`
- `quality_score`
- `overall_score = sqrt(semantics_score * quality_score)`
- Judge calls are routed to a local OpenAI-compatible model served by
`vllm-omni`, not a remote provider.
Dataset:
- Default `--dataset-ref` is `stepfun-ai/GEdit-Bench`
- You can also pass a local dataset directory previously saved with
Hugging Face `datasets`
Example usage:
```bash
python benchmarks/accuracy/image_to_image/run_gedit_bench.py generate \
--output-root benchmarks/accuracy/image_to_image/results \
--base-url http://127.0.0.1:8000 \
--model Qwen/Qwen-Image-Edit \
--model-name qwen_image_edit \
--dataset-ref stepfun-ai/GEdit-Bench \
--task-type all \
--instruction-language en
```
```bash
python benchmarks/accuracy/image_to_image/run_gedit_bench.py evaluate \
--output-root benchmarks/accuracy/image_to_image/results \
--model-name qwen_image_edit \
--save-dir benchmarks/accuracy/image_to_image/scores \
--dataset-ref stepfun-ai/GEdit-Bench \
--judge-base-url http://127.0.0.1:8000 \
--judge-model Qwen/Qwen2.5-VL-7B-Instruct \
--judge-api-key EMPTY
```
```bash
python benchmarks/accuracy/image_to_image/run_gedit_bench.py summarize \
--csv-path benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_vie_score.csv \
--language en
```
Example summary output:
```json
{
"language": "all",
"languages": {
"en": {
"overall": {"count": 110, "Q_SC": 6.58, "Q_PQ": 5.89, "Q_O": 5.86},
"intersection": {"count": 78, "Q_SC": 6.50, "Q_PQ": 5.66, "Q_O": 5.65}
},
"cn": {
"overall": {"count": 110, "Q_SC": 6.90, "Q_PQ": 5.78, "Q_O": 6.11},
"intersection": {"count": 63, "Q_SC": 7.22, "Q_PQ": 5.59, "Q_O": 6.28}
}
}
}
```
Example generated images to inspect:
- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/background_change/en/<sample>.png`
- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/text_change/en/<sample>.png`
- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/subject-replace/cn/<sample>.png`
Example score artifacts to inspect together with the images:
- `benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_vie_score.csv`
- `benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_summary.json`
What to expect:
- `Q_SC` measures instruction following and content preservation.
- `Q_PQ` measures image naturalness and artifact quality.
- `Q_O` is the combined overall score; higher is better.
- `overall.count` is the number of evaluated samples for that language, while `intersection.count` is the subset with `intersection_exist == True`.
Notes:
- This flow requires the optional Hugging Face `datasets` package.
- `generate` writes `generation_manifest.json` with local output coverage.
- The current repo marker set exposes `L4` but not `L5`, so if you promote an
end-to-end smoke test into CI, use the existing `advanced_model`, `benchmark`,
and `L4` markers or introduce a new repo-wide marker explicitly first.
@@ -1 +0,0 @@
"""GEdit-Bench integration for vLLM-Omni."""
@@ -1,718 +0,0 @@
from __future__ import annotations
import argparse
import csv
import json
import logging
import math
import statistics
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from io import BytesIO
from pathlib import Path
from typing import Any
import requests
from PIL import Image
from tqdm.auto import tqdm
from benchmarks.accuracy.common import (
VllmOmniImageClient,
build_openai_url,
extract_json_object,
write_json,
)
GROUPS = [
"background_change",
"color_alter",
"material_alter",
"motion_change",
"ps_human",
"style_change",
"subject-add",
"subject-remove",
"subject-replace",
"text_change",
"tone_transfer",
]
DEFAULT_SAMPLES_PER_GROUP = 10
logger = logging.getLogger(__name__)
def infer_model_name(model: str) -> str:
normalized = model.rstrip("/\\")
name = Path(normalized).name
return name or normalized
def resolve_model_name(*, model_name: str | None, model: str | None = None, output_root: Path | None = None) -> str:
if model_name:
return model_name
if model:
return infer_model_name(model)
if output_root is not None:
candidates = sorted(path.name for path in output_root.iterdir() if path.is_dir())
if len(candidates) == 1:
return candidates[0]
if not candidates:
raise ValueError(f"Could not infer model-name from empty output root: {output_root}")
raise ValueError(
f"Could not infer model-name from output root {output_root}; multiple candidates found: {candidates}"
)
raise ValueError("model-name is required when it cannot be inferred from model or output-root")
def parse_score_payload(raw_text: str) -> dict[str, Any]:
try:
parsed = extract_json_object(raw_text)
except Exception:
stripped = raw_text.strip()
if stripped.startswith("[") and stripped.endswith("]"):
parsed = {"score": json.loads(stripped), "reasoning": ""}
else:
raise
score = parsed.get("score", [])
if not isinstance(score, list):
score = [score]
parsed["score"] = [int(value) for value in score]
return parsed
def _utc_timestamp() -> str:
return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
by_task: dict[str, list[dict[str, Any]]] = defaultdict(list)
by_language: dict[str, list[dict[str, Any]]] = defaultdict(list)
for record in records:
by_task[record["task_type"]].append(record)
by_language[record["instruction_language"]].append(record)
return {
"count": len(records),
"by_task": {
task_type: {
"count": len(rows),
"samples": sorted(row["key"] for row in rows),
}
for task_type, rows in sorted(by_task.items())
},
"by_language": {
language: {
"count": len(rows),
"samples": sorted(row["key"] for row in rows),
}
for language, rows in sorted(by_language.items())
},
}
def select_balanced_gedit_rows(
rows: list[dict[str, Any]],
*,
task_type: str = "all",
instruction_language: str = "all",
samples_per_group: int | None,
) -> list[dict[str, Any]]:
filtered_rows = []
for row in rows:
if task_type != "all" and row["task_type"] != task_type:
continue
if instruction_language != "all" and row["instruction_language"] != instruction_language:
continue
filtered_rows.append(row)
if samples_per_group is None:
return filtered_rows
def _select_rows_for_group(group_rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
if instruction_language != "all":
return group_rows[:samples_per_group]
rows_by_language: dict[str, list[dict[str, Any]]] = defaultdict(list)
for row in group_rows:
rows_by_language[str(row.get("instruction_language", "")).strip()].append(row)
ordered_languages = ["en", "cn"]
per_language_quota = samples_per_group // max(len(ordered_languages), 1)
remainder = samples_per_group % max(len(ordered_languages), 1)
selected_rows: list[dict[str, Any]] = []
leftovers: list[dict[str, Any]] = []
for index, language in enumerate(ordered_languages):
language_rows = rows_by_language.get(language, [])
quota = per_language_quota + (1 if index < remainder else 0)
selected_rows.extend(language_rows[:quota])
leftovers.extend(language_rows[quota:])
if len(selected_rows) < samples_per_group:
selected_rows.extend(leftovers[: samples_per_group - len(selected_rows)])
return selected_rows
if task_type != "all":
return _select_rows_for_group(filtered_rows)
grouped_rows: dict[str, list[dict[str, Any]]] = defaultdict(list)
for row in filtered_rows:
grouped_rows[row["task_type"]].append(row)
selected: list[dict[str, Any]] = []
for group in GROUPS:
selected.extend(_select_rows_for_group(grouped_rows.get(group, [])))
return selected
def summarize_gedit_rows(rows: list[dict[str, Any]], language: str = "all") -> dict[str, Any]:
return summarize_gedit_rows_with_backbone(rows, language=language)
def _mean_or_none(values: list[float]) -> float | None:
return statistics.fmean(values) if values else None
def _to_q_metrics(section: dict[str, float | None]) -> dict[str, float | None]:
return {
"count": section.get("count"),
"Q_SC": section.get("avg_semantics"),
"Q_PQ": section.get("avg_quality"),
"Q_O": section.get("avg_overall"),
}
def _summarize_gedit_rows_single_language(
rows: list[dict[str, Any]],
*,
language: str,
) -> dict[str, Any]:
filtered_rows = [row for row in rows if str(row.get("instruction_language", "")).strip() == language]
def _to_bool(value: Any) -> bool:
if isinstance(value, bool):
return value
return str(value).strip().lower() in {"1", "true", "yes"}
per_group: dict[str, dict[str, float | None]] = {}
intersection_group: dict[str, dict[str, float | None]] = {}
for group in GROUPS:
group_rows = [row for row in filtered_rows if row.get("task_type") == group]
semantics = [float(row["semantics_score"]) for row in group_rows]
quality = [float(row["quality_score"]) for row in group_rows]
overall = [math.sqrt(float(row["semantics_score"]) * float(row["quality_score"])) for row in group_rows]
per_group[group] = {
"count": len(group_rows),
"avg_semantics": _mean_or_none(semantics),
"avg_quality": _mean_or_none(quality),
"avg_overall": _mean_or_none(overall),
}
intersection_rows = [row for row in group_rows if _to_bool(row.get("intersection_exist", False))]
intersection_semantics = [float(row["semantics_score"]) for row in intersection_rows]
intersection_quality = [float(row["quality_score"]) for row in intersection_rows]
intersection_overall = [
math.sqrt(float(row["semantics_score"]) * float(row["quality_score"])) for row in intersection_rows
]
intersection_group[group] = {
"count": len(intersection_rows),
"avg_semantics": _mean_or_none(intersection_semantics),
"avg_quality": _mean_or_none(intersection_quality),
"avg_overall": _mean_or_none(intersection_overall),
}
overall_section = {
"count": len(filtered_rows),
"avg_semantics": _mean_or_none(
[score["avg_semantics"] for score in per_group.values() if score["avg_semantics"] is not None]
),
"avg_quality": _mean_or_none(
[score["avg_quality"] for score in per_group.values() if score["avg_quality"] is not None]
),
"avg_overall": _mean_or_none(
[score["avg_overall"] for score in per_group.values() if score["avg_overall"] is not None]
),
}
intersection_section = {
"count": sum(score["count"] for score in intersection_group.values()),
"avg_semantics": _mean_or_none(
[score["avg_semantics"] for score in intersection_group.values() if score["avg_semantics"] is not None]
),
"avg_quality": _mean_or_none(
[score["avg_quality"] for score in intersection_group.values() if score["avg_quality"] is not None]
),
"avg_overall": _mean_or_none(
[score["avg_overall"] for score in intersection_group.values() if score["avg_overall"] is not None]
),
}
return {
"language": language,
"by_group": {group: _to_q_metrics(section) for group, section in per_group.items()},
"overall": _to_q_metrics(overall_section),
"intersection": _to_q_metrics(intersection_section),
}
def summarize_gedit_rows_with_backbone(
rows: list[dict[str, Any]],
*,
language: str = "all",
) -> dict[str, Any]:
if language == "all":
return {
"language": "all",
"languages": {
single_language: _summarize_gedit_rows_single_language(
rows,
language=single_language,
)
for single_language in ["en", "cn"]
},
}
return _summarize_gedit_rows_single_language(rows, language=language)
def _require_datasets():
try:
from datasets import load_dataset, load_from_disk
except ImportError as exc:
raise ImportError("GEdit-Bench requires the optional `datasets` package.") from exc
return load_dataset, load_from_disk
def _load_gedit_dataset(dataset_ref: str):
load_dataset, load_from_disk = _require_datasets()
dataset_path = Path(dataset_ref)
if dataset_path.exists():
if (dataset_path / "state.json").exists() and (dataset_path / "dataset_info.json").exists():
return load_from_disk(str(dataset_path))
return load_dataset(str(dataset_path))
return load_dataset(dataset_ref)
def _resolve_gedit_split(dataset_obj: Any) -> Any:
if isinstance(dataset_obj, dict):
if "train" in dataset_obj:
return dataset_obj["train"]
return dataset_obj
try:
return dataset_obj["train"]
except Exception:
return dataset_obj
def _to_pil_image(value: Any) -> Image.Image:
if isinstance(value, Image.Image):
return value.convert("RGB")
if isinstance(value, dict) and "bytes" in value:
image = Image.open(BytesIO(value["bytes"]))
image.load()
return image.convert("RGB")
raise TypeError(f"Unsupported image payload type: {type(value)!r}")
class LocalVIEScorer:
def __init__(self, *, base_url: str, api_key: str, model: str, timeout: int = 600):
self.base_url = base_url.rstrip("/")
self.api_key = api_key
self.model = model
self.timeout = timeout
self.sc_prompt = (
"You are evaluating image editing quality.\n"
"Two images are provided: the source image and the edited image.\n"
'Return JSON only in the format {"score": [edit_success, content_preservation], '
'"reasoning": "..."}.\n'
"Each score must be an integer from 0 to 10.\n"
"Editing instruction: <instruction>"
)
self.pq_prompt = (
"You are evaluating image quality.\n"
'Return JSON only in the format {"score": [naturalness, artifact_free], "reasoning": "..."}.\n'
"Each score must be an integer from 0 to 10."
)
def _request(self, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
from benchmarks.accuracy.common import pil_to_data_url
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
for image in images:
content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
response = requests.post(
build_openai_url(self.base_url, "/chat/completions"),
json={
"model": self.model,
"messages": [{"role": "user", "content": content}],
"temperature": 0,
},
headers={
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
},
timeout=self.timeout,
)
response.raise_for_status()
message_content = response.json()["choices"][0]["message"]["content"]
if isinstance(message_content, list):
text = "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
else:
text = str(message_content)
return parse_score_payload(text)
def evaluate(self, source_image: Image.Image, edited_image: Image.Image, instruction: str) -> dict[str, float]:
sc_payload = self._request(self.sc_prompt.replace("<instruction>", instruction), [source_image, edited_image])
pq_payload = self._request(self.pq_prompt, [edited_image])
semantics = float(min(sc_payload["score"])) if sc_payload["score"] else 0.0
quality = float(min(pq_payload["score"])) if pq_payload["score"] else 0.0
overall = math.sqrt(semantics * quality)
return {
"semantics_score": semantics,
"quality_score": quality,
"overall_score": overall,
}
class GEditBenchRunner:
def __init__(
self,
*,
dataset_ref: str,
output_root: Path,
base_url: str,
model: str,
api_key: str = "EMPTY",
width: int = 512,
height: int = 512,
num_inference_steps: int = 20,
guidance_scale: float | None = None,
seed: int | None = 42,
):
self.dataset_ref = dataset_ref
self.output_root = output_root
self.model = model
self.width = width
self.height = height
self.num_inference_steps = num_inference_steps
self.guidance_scale = guidance_scale
self.seed = seed
self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
def generate(
self,
*,
model_name: str,
task_type: str = "all",
instruction_language: str = "all",
workers: int = 1,
max_samples: int | None = None,
samples_per_group: int | None = None,
) -> list[dict[str, Any]]:
dataset = _resolve_gedit_split(_load_gedit_dataset(self.dataset_ref))
rows = select_balanced_gedit_rows(
list(dataset),
task_type=task_type,
instruction_language=instruction_language,
samples_per_group=samples_per_group,
)
if max_samples is not None:
rows = rows[:max_samples]
outputs: list[dict[str, Any]] = []
total = len(rows)
if workers <= 1:
with tqdm(total=total, desc="GEdit generate", unit="sample") as progress:
for item in rows:
result = self._safe_generate_one(model_name, item)
if result:
outputs.append(result)
progress.update(1)
return outputs
with tqdm(total=total, desc="GEdit generate", unit="sample") as progress:
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = [executor.submit(self._safe_generate_one, model_name, item) for item in rows]
for future in as_completed(futures):
result = future.result()
if result:
outputs.append(result)
progress.update(1)
return outputs
def _safe_generate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
try:
return self._generate_one(model_name, item)
except Exception:
logger.exception("Failed to generate GEdit-Bench sample %s", item.get("key", "<unknown>"))
return None
def _generate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
output_path = (
self.output_root
/ model_name
/ "fullset"
/ item["task_type"]
/ item["instruction_language"]
/ f"{item['key']}.png"
)
output_path.parent.mkdir(parents=True, exist_ok=True)
if output_path.exists():
return {
"task_type": item["task_type"],
"instruction_language": item["instruction_language"],
"key": item["key"],
"output_path": str(output_path),
}
source_image = _to_pil_image(item["input_image_raw"])
edited_image = self.client.generate_image_edit(
model=self.model,
prompt=item["instruction"],
images=source_image,
width=self.width,
height=self.height,
num_inference_steps=self.num_inference_steps,
guidance_scale=self.guidance_scale,
seed=self.seed,
)
edited_image.save(output_path)
return {
"task_type": item["task_type"],
"instruction_language": item["instruction_language"],
"key": item["key"],
"output_path": str(output_path),
}
class GEditBenchEvaluator:
def __init__(self, *, dataset_ref: str, output_root: Path, scorer: LocalVIEScorer):
self.dataset_ref = dataset_ref
self.output_root = output_root
self.scorer = scorer
def evaluate(
self,
*,
model_name: str,
save_dir: Path,
task_type: str = "all",
instruction_language: str = "all",
workers: int = 1,
max_samples: int | None = None,
samples_per_group: int | None = None,
) -> dict[str, Any]:
dataset = _resolve_gedit_split(_load_gedit_dataset(self.dataset_ref))
rows = select_balanced_gedit_rows(
list(dataset),
task_type=task_type,
instruction_language=instruction_language,
samples_per_group=samples_per_group,
)
if max_samples is not None:
rows = rows[:max_samples]
results: list[dict[str, Any]] = []
total = len(rows)
if workers <= 1:
with tqdm(total=total, desc="GEdit evaluate", unit="sample") as progress:
for item in rows:
result = self._safe_evaluate_one(model_name, item)
if result:
results.append(result)
progress.update(1)
else:
with tqdm(total=total, desc="GEdit evaluate", unit="sample") as progress:
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = [executor.submit(self._safe_evaluate_one, model_name, item) for item in rows]
for future in as_completed(futures):
result = future.result()
if result:
results.append(result)
progress.update(1)
save_dir.mkdir(parents=True, exist_ok=True)
base_name = f"{model_name}_{task_type}_{instruction_language}"
timestamp = _utc_timestamp()
csv_path = save_dir / f"{base_name}_vie_score.csv"
timestamped_csv_path = save_dir / f"{base_name}_vie_score_{timestamp}.csv"
def _write_csv(path: Path) -> None:
with path.open("w", encoding="utf-8", newline="") as handle:
writer = csv.DictWriter(
handle,
fieldnames=[
"key",
"task_type",
"edited_image",
"instruction",
"semantics_score",
"quality_score",
"overall_score",
"intersection_exist",
"instruction_language",
],
)
writer.writeheader()
for row in results:
writer.writerow(row)
_write_csv(csv_path)
_write_csv(timestamped_csv_path)
summary = summarize_gedit_rows_with_backbone(
results,
language=instruction_language,
)
summary_path = save_dir / f"{base_name}_summary.json"
timestamped_summary_path = save_dir / f"{base_name}_summary_{timestamp}.json"
write_json(summary_path, summary)
write_json(timestamped_summary_path, summary)
return {
"results": results,
"summary": summary,
"csv_path": str(csv_path),
"summary_path": str(summary_path),
"timestamped_csv_path": str(timestamped_csv_path),
"timestamped_summary_path": str(timestamped_summary_path),
}
def _safe_evaluate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
try:
return self._evaluate_one(model_name, item)
except Exception:
logger.exception("Failed to evaluate GEdit-Bench sample %s", item.get("key", "<unknown>"))
return None
def _evaluate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
edited_image_path = (
self.output_root
/ model_name
/ "fullset"
/ item["task_type"]
/ item["instruction_language"]
/ f"{item['key']}.png"
)
if not edited_image_path.exists():
return None
source_image = _to_pil_image(item["input_image_raw"])
edited_image = Image.open(edited_image_path).convert("RGB")
scores = self.scorer.evaluate(source_image, edited_image, item["instruction"])
return {
"key": item["key"],
"task_type": item["task_type"],
"edited_image": str(edited_image_path),
"instruction": item["instruction"],
"semantics_score": scores["semantics_score"],
"quality_score": scores["quality_score"],
"overall_score": scores["overall_score"],
"intersection_exist": item.get("Intersection_exist", False),
"instruction_language": item["instruction_language"],
}
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Run the GEdit-Bench integration against a local vLLM-Omni server.")
subparsers = parser.add_subparsers(dest="command", required=True)
generate = subparsers.add_parser("generate")
generate.add_argument("--dataset-ref", type=str, default="stepfun-ai/GEdit-Bench")
generate.add_argument("--output-root", type=Path, required=True)
generate.add_argument("--base-url", type=str, required=True)
generate.add_argument("--model", type=str, required=True)
generate.add_argument("--model-name", type=str, default=None)
generate.add_argument("--api-key", type=str, default="EMPTY")
generate.add_argument("--task-type", choices=["all", *GROUPS], default="all")
generate.add_argument("--instruction-language", choices=["all", "en", "cn"], default="all")
generate.add_argument("--width", type=int, default=512)
generate.add_argument("--height", type=int, default=512)
generate.add_argument("--num-inference-steps", type=int, default=20)
generate.add_argument("--guidance-scale", type=float, default=None)
generate.add_argument("--seed", type=int, default=42)
generate.add_argument("--workers", type=int, default=1)
generate.add_argument("--max-samples", type=int, default=None)
generate.add_argument("--samples-per-group", type=int, default=None)
evaluate = subparsers.add_parser("evaluate")
evaluate.add_argument("--dataset-ref", type=str, default="stepfun-ai/GEdit-Bench")
evaluate.add_argument("--output-root", type=Path, required=True)
evaluate.add_argument("--model-name", type=str, default=None)
evaluate.add_argument("--save-dir", type=Path, required=True)
evaluate.add_argument("--task-type", choices=["all", *GROUPS], default="all")
evaluate.add_argument("--instruction-language", choices=["all", "en", "cn"], default="all")
evaluate.add_argument("--judge-base-url", type=str, required=True)
evaluate.add_argument("--judge-model", type=str, required=True)
evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
evaluate.add_argument("--workers", type=int, default=1)
evaluate.add_argument("--max-samples", type=int, default=None)
evaluate.add_argument("--samples-per-group", type=int, default=None)
summarize = subparsers.add_parser("summarize")
summarize.add_argument("--csv-path", type=Path, required=True)
summarize.add_argument("--language", choices=["all", "en", "cn"], default="all")
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
if args.command == "generate":
model_name = resolve_model_name(model_name=args.model_name, model=args.model)
runner = GEditBenchRunner(
dataset_ref=args.dataset_ref,
output_root=args.output_root,
base_url=args.base_url,
model=args.model,
api_key=args.api_key,
width=args.width,
height=args.height,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
)
records = runner.generate(
model_name=model_name,
task_type=args.task_type,
instruction_language=args.instruction_language,
workers=args.workers,
max_samples=args.max_samples,
samples_per_group=args.samples_per_group,
)
payload = {"records": records, "summary": summarize_generated_records(records)}
write_json(args.output_root / model_name / "generation_manifest.json", payload)
return 0
if args.command == "evaluate":
model_name = resolve_model_name(model_name=args.model_name, output_root=args.output_root)
scorer = LocalVIEScorer(
base_url=args.judge_base_url,
api_key=args.judge_api_key,
model=args.judge_model,
)
evaluator = GEditBenchEvaluator(dataset_ref=args.dataset_ref, output_root=args.output_root, scorer=scorer)
evaluator.evaluate(
model_name=model_name,
save_dir=args.save_dir,
task_type=args.task_type,
instruction_language=args.instruction_language,
workers=args.workers,
max_samples=args.max_samples,
samples_per_group=args.samples_per_group,
)
return 0
if args.command == "summarize":
with args.csv_path.open("r", encoding="utf-8", newline="") as handle:
rows = list(csv.DictReader(handle))
summary = summarize_gedit_rows_with_backbone(rows, language=args.language)
print(json.dumps(summary, indent=2, ensure_ascii=False))
return 0
parser.error(f"Unknown command: {args.command}")
return 1
@@ -1,13 +0,0 @@
# ruff: noqa: E402, I001
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[3]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from benchmarks.accuracy.image_to_image.gedit_bench import main
if __name__ == "__main__":
raise SystemExit(main())
-103
View File
@@ -1,103 +0,0 @@
# GEBench on vLLM-Omni
This integration adapts the upstream `stepfun-ai/GEBench` scripts into
`vllm-omni/benchmarks/accuracy/text_to_image`.
Upstream mapping:
- `scripts/generate.py` -> `run_gebench.py generate`
- `scripts/evaluate.py` -> `run_gebench.py evaluate`
- upstream prompt / judge logic -> `gbench.py`
What changed:
- Generation calls the local OpenAI-compatible `vllm-omni` endpoints:
- `/v1/images/generations` for text-only frame generation
- `/v1/images/edits` for image-conditioned GUI transition generation
- Evaluation still keeps the GEBench scoring dimensions:
- `goal`
- `logic`
- `cons`
- `ui`
- `qual`
- Judge calls are also routed to a local OpenAI-compatible model served by
`vllm-omni` instead of a remote service.
Expected dataset layout:
- Clone the benchmark dataset from Hugging Face into a local directory:
```bash
git clone https://huggingface.co/datasets/stepfun-ai/GEBench /path/to/GEBench
```
Example usage:
```bash
python benchmarks/accuracy/text_to_image/run_gebench.py generate \
--dataset-root /path/to/GEBench \
--output-root benchmarks/accuracy/text_to_image/outputs \
--base-url http://127.0.0.1:8000 \
--model Tongyi-MAI/Z-Image-Turbo \
--data-type type3
```
```bash
python benchmarks/accuracy/text_to_image/run_gebench.py evaluate \
--dataset-root /path/to/GEBench \
--output-root benchmarks/accuracy/text_to_image/outputs \
--data-type type3 \
--judge-base-url http://127.0.0.1:8000 \
--judge-model Qwen/Qwen2.5-VL-7B-Instruct \
--judge-api-key EMPTY
```
```bash
python benchmarks/accuracy/text_to_image/run_gebench.py summarize \
--output-root benchmarks/accuracy/text_to_image/outputs
```
Example summary output:
```json
{
"generation": {
"count": 20,
"by_type": {
"type3": {"count": 10},
"type4": {"count": 10}
}
},
"evaluation": {
"count": 20,
"overall_mean": 0.52,
"by_type": {
"type3": {"count": 10, "overall_mean": 0.50, "overall_mean_100": 50.0},
"type4": {"count": 10, "overall_mean": 0.54, "overall_mean_100": 54.0}
}
}
}
```
Example generated images to inspect:
- `benchmarks/accuracy/text_to_image/outputs/03_trajectory_text_fictionalapp/<lang_device>/<sample>/frame0.png`
- `benchmarks/accuracy/text_to_image/outputs/03_trajectory_text_fictionalapp/<lang_device>/<sample>/frame5.png`
- `benchmarks/accuracy/text_to_image/outputs/04_trajectory_text_realapp/<lang_device>/<sample>/frame0.png`
- `benchmarks/accuracy/text_to_image/outputs/04_trajectory_text_realapp/<lang_device>/<sample>/frame5.png`
What to expect:
- `overall_mean` is normalized to `0.0 ~ 1.0`; higher is better.
- `frame0.png` is the initial GUI frame; `frame5.png` is the final trajectory frame most often used for quick inspection.
- For full debugging, inspect the whole `frame0.png` ... `frame5.png` sequence for one sample directory.
Notes:
- GEBench upstream leaves type3/type4 generation unfinished. This integration
fills that gap with a trajectory runner that generates `frame0.png` followed
by `frame1.png` ... `frame5.png`.
- Type1/2/5 require an image-edit capable model exposed through
`vllm-omni serve`.
- `summarize` will report both generated coverage and any existing evaluation
summary files.
@@ -1 +0,0 @@
"""GEBench integration for vLLM-Omni."""
-927
View File
@@ -1,927 +0,0 @@
from __future__ import annotations
import argparse
import json
import statistics
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
import requests
from PIL import Image
from benchmarks.accuracy.common import (
VllmOmniImageClient,
build_openai_url,
ensure_dir,
extract_json_object,
find_first_image,
load_json,
pil_to_data_url,
save_image,
write_json,
)
TYPE_TO_FOLDER = {
"type1": "01_single_step",
"type2": "02_multi_step",
"type3": "03_trajectory_text_fictionalapp",
"type4": "04_trajectory_text_realapp",
"type5": "05_grounding_data",
}
SCORE_KEYS = ("goal", "logic", "cons", "ui", "qual")
DEFAULT_SAMPLES_PER_TYPE = 10
def _utc_timestamp() -> str:
return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
def _write_json_with_timestamp(path: Path, payload: dict[str, Any]) -> Path:
write_json(path, payload)
timestamped_path = path.with_name(f"{path.stem}_{_utc_timestamp()}{path.suffix}")
write_json(timestamped_path, payload)
return timestamped_path
@dataclass(frozen=True)
class GEBenchSampleSpec:
sample_path: Path
metadata: dict[str, Any]
sample_name: str
lang_device: str
def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
for record in records:
by_type[record["data_type"]].append(record)
return {
"count": len(records),
"by_type": {
data_type: {
"count": len(rows),
}
for data_type, rows in sorted(by_type.items())
},
}
def summarize_gebench_results(results: list[dict[str, Any]]) -> dict[str, Any]:
by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
for result in results:
by_type[result["data_type"]].append(result)
summary: dict[str, Any] = {
"count": len(results),
"overall_mean": statistics.fmean(r["overall"] for r in results) if results else 0.0,
"by_type": {},
}
for data_type, rows in by_type.items():
score_means: dict[str, float] = {}
all_score_keys = {key for row in rows for key in row.get("scores", {}).keys()}
for score_key in all_score_keys:
values = [row["scores"][score_key] for row in rows if score_key in row.get("scores", {})]
score_means[score_key] = statistics.fmean(values) if values else 0.0
overall_mean = statistics.fmean(row["overall"] for row in rows)
summary["by_type"][data_type] = {
"count": len(rows),
"overall_mean": overall_mean,
"overall_mean_100": overall_mean * 100.0,
"score_means": score_means,
}
return summary
def select_balanced_gebench_samples(
sample_paths_by_type: dict[str, list[Any]],
*,
samples_per_type: int | None,
) -> dict[str, list[Any]]:
if samples_per_type is None:
return {data_type: list(paths) for data_type, paths in sample_paths_by_type.items()}
return {data_type: list(paths)[:samples_per_type] for data_type, paths in sample_paths_by_type.items()}
def collect_gebench_generation_summary(output_root: Path) -> dict[str, Any]:
records: list[dict[str, Any]] = []
for data_type, folder_name in TYPE_TO_FOLDER.items():
type_root = output_root / folder_name
if not type_root.exists():
continue
for lang_dir in sorted(path for path in type_root.iterdir() if path.is_dir()):
for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir()):
expected = sample_dir / "frame5.png" if data_type in {"type2", "type3", "type4"} else None
if expected is None:
expected = find_first_image(sample_dir)
elif not expected.exists():
expected = None
if expected is None:
continue
records.append(
{
"data_type": data_type,
"sample_name": f"{lang_dir.name}/{sample_dir.name}",
"output_path": str(expected),
}
)
return summarize_generated_records(records)
def _normalize_score_key(key: str) -> str:
mapping = {
"goal": "goal",
"logic": "logic",
"cons": "cons",
"consistency": "cons",
"ui": "ui",
"qual": "qual",
"quality": "qual",
}
return mapping.get(key.lower(), key.lower())
def _normalize_scores(raw_scores: dict[str, Any]) -> dict[str, int]:
scores: dict[str, int] = {}
for key, value in raw_scores.items():
normalized = _normalize_score_key(key)
if normalized not in SCORE_KEYS:
continue
scalar = value.get("s", 0) if isinstance(value, dict) else value
try:
scores[normalized] = int(scalar)
except (TypeError, ValueError):
scores[normalized] = 0
for key in SCORE_KEYS:
scores.setdefault(key, 0)
return scores
def _compute_overall(scores: dict[str, int]) -> float:
return sum(scores.values()) / (len(SCORE_KEYS) * 5.0)
def _iter_sample_paths(dataset_root: Path, data_type: str) -> list[Path]:
data_dir = dataset_root / TYPE_TO_FOLDER[data_type]
if not data_dir.exists():
raise FileNotFoundError(f"Dataset path not found: {data_dir}")
samples: list[Path] = []
for lang_dir in sorted(path for path in data_dir.iterdir() if path.is_dir()):
for child in sorted(lang_dir.iterdir()):
if child.is_dir():
samples.append(child)
elif child.suffix.lower() == ".json":
samples.append(child)
return samples
def _sample_name_from_metadata(metadata: dict[str, Any], sample_path: Path, item_index: int | None = None) -> str:
sample_id = metadata.get("id") or metadata.get("sample_id") or metadata.get("name")
if sample_id:
return str(sample_id)
if item_index is not None:
return f"{sample_path.stem}_{item_index:04d}"
return sample_path.stem if sample_path.is_file() else sample_path.name
def _expand_sample_path(sample_path: Path) -> list[GEBenchSampleSpec]:
if sample_path.is_dir():
metadata = _load_metadata(sample_path)
return [
GEBenchSampleSpec(
sample_path=sample_path,
metadata=metadata,
sample_name=_sample_name(sample_path),
lang_device=_lang_device(sample_path, metadata),
)
]
payload = load_json(sample_path)
if isinstance(payload, dict):
return [
GEBenchSampleSpec(
sample_path=sample_path,
metadata=payload,
sample_name=_sample_name_from_metadata(payload, sample_path),
lang_device=_lang_device(sample_path, payload),
)
]
if isinstance(payload, list):
specs: list[GEBenchSampleSpec] = []
for item_index, item in enumerate(payload):
if not isinstance(item, dict):
continue
specs.append(
GEBenchSampleSpec(
sample_path=sample_path,
metadata=item,
sample_name=_sample_name_from_metadata(item, sample_path, item_index),
lang_device=_lang_device(sample_path, item),
)
)
return specs
raise TypeError(f"Unsupported metadata payload type for sample {sample_path}: {type(payload)!r}")
def _iter_sample_specs(dataset_root: Path, data_type: str) -> list[GEBenchSampleSpec]:
sample_specs: list[GEBenchSampleSpec] = []
for sample_path in _iter_sample_paths(dataset_root, data_type):
sample_specs.extend(_expand_sample_path(sample_path))
return sample_specs
def _load_metadata(sample_path: Path) -> dict[str, Any]:
if sample_path.is_file():
return load_json(sample_path)
for candidate in ("meta_data.json", "metadata.json"):
meta_path = sample_path / candidate
if meta_path.exists():
return load_json(meta_path)
raise FileNotFoundError(f"Metadata not found for sample: {sample_path}")
def _sample_name(sample_path: Path) -> str:
return sample_path.stem if sample_path.is_file() else sample_path.name
def _lang_device(sample_path: Path, metadata: dict[str, Any]) -> str:
return str(metadata.get("lang_device") or sample_path.parent.name)
def _resolve_referenced_image(
*,
metadata: dict[str, Any],
sample_path: Path,
dataset_root: Path,
data_type: str,
) -> Image.Image | None:
for key in ("image", "input_image", "initial_image", "reference_image"):
image_ref = metadata.get(key)
if not image_ref:
continue
candidate = dataset_root / TYPE_TO_FOLDER[data_type] / str(image_ref)
if candidate.exists():
image = Image.open(candidate)
image.load()
return image.convert("RGB")
if sample_path.is_dir():
local_image = find_first_image(sample_path)
if local_image:
image = Image.open(local_image)
image.load()
return image.convert("RGB")
return None
def _trajectory_steps(metadata: dict[str, Any]) -> list[dict[str, Any]]:
for key in ("trajectory", "steps", "frames"):
value = metadata.get(key)
if isinstance(value, list):
return [step for step in value if isinstance(step, dict)]
extracted: list[dict[str, Any]] = []
for index in range(1, 6):
value = metadata.get(f"step{index}") or metadata.get(str(index))
if isinstance(value, dict):
extracted.append(value)
return extracted
def _text_or_default(value: Any, default: str = "") -> str:
return str(value).strip() if value is not None else default
def _type1_prompt(metadata: dict[str, Any]) -> str:
caption = _text_or_default(metadata.get("caption") or metadata.get("instruction"), "Transform the reference GUI.")
return (
"Using the reference GUI screenshot, generate the next GUI state after the requested interaction.\n\n"
f"Requested change:\n{caption}\n\n"
"Requirements:\n"
"- Preserve layout, visual identity, and unrelated regions.\n"
"- Only apply the requested state change.\n"
"- Keep all text and controls readable.\n"
)
def _type2_prompt(goal: str, step_num: int) -> str:
return (
"Generate the next GUI state for a multi-step task.\n\n"
f"Overall goal: {goal}\n"
f"Current progress step: {step_num}/5\n\n"
"Requirements:\n"
"- The change should be incremental and plausible.\n"
"- Preserve layout and visual identity.\n"
"- Make text/buttons readable.\n"
)
def _type34_initial_prompt(metadata: dict[str, Any], first_step: dict[str, Any]) -> str:
app_name = _text_or_default(metadata.get("app_name"), "App")
final_goal = _text_or_default(metadata.get("final_goal") or metadata.get("instruction"), "Complete the task.")
visual_description = _text_or_default(
metadata.get("visual_description") or first_step.get("visual_description") or first_step.get("description"),
"A clean product-quality app home screen.",
)
return (
"Generate the first GUI frame for a task trajectory.\n\n"
f"App name: {app_name}\n"
f"Final goal: {final_goal}\n"
f"Visual description:\n{visual_description}\n\n"
"Requirements:\n"
"- Generate a production-looking UI screenshot only.\n"
"- Keep the layout coherent and readable.\n"
)
def _type34_next_prompt(step_num: int, step_info: dict[str, Any]) -> str:
action = _text_or_default(step_info.get("action") or step_info.get("instruction"), "Continue the task.")
visual_description = _text_or_default(
step_info.get("visual_description") or step_info.get("description"),
"Reflect the expected next GUI state.",
)
return (
"Using the previous frame as reference, generate the next GUI frame.\n\n"
f"Step {step_num} action: {action}\n"
f"Expected visual state:\n{visual_description}\n\n"
"Requirements:\n"
"- Only change UI regions affected by this action.\n"
"- Preserve persistent bars, layout, and style.\n"
"- Keep text and icons readable.\n"
)
def _type5_prompt(metadata: dict[str, Any]) -> str:
grounding = metadata.get("grounding") or {}
explanation = _text_or_default(
metadata.get("grounding_explanation") or grounding.get("effect") or grounding.get("description"),
"Predict the immediate GUI reaction to the indicated target.",
)
return (
"Using the reference GUI screenshot, predict the immediate GUI state after the grounded interaction.\n\n"
f"Expected effect: {explanation}\n"
f"Grounding metadata: {json.dumps(grounding, ensure_ascii=False)}\n\n"
"Requirements:\n"
"- Apply only the interaction-triggered change.\n"
"- Preserve unrelated regions.\n"
"- Keep the UI realistic and readable.\n"
)
def _make_storyboard_image(
frames: list[Image.Image],
*,
columns: int = 3,
background_color: tuple[int, int, int] = (255, 255, 255),
) -> Image.Image:
if not frames:
raise ValueError("Expected at least one frame to build a storyboard image.")
normalized = [frame.convert("RGB") for frame in frames]
frame_width = max(frame.width for frame in normalized)
frame_height = max(frame.height for frame in normalized)
rows = (len(normalized) + columns - 1) // columns
storyboard = Image.new("RGB", (frame_width * columns, frame_height * rows), color=background_color)
for index, frame in enumerate(normalized):
x_offset = (index % columns) * frame_width
y_offset = (index // columns) * frame_height
if frame.size != (frame_width, frame_height):
frame = frame.resize((frame_width, frame_height))
storyboard.paste(frame, (x_offset, y_offset))
return storyboard
def _trajectory_judge_payload(frames: list[Image.Image]) -> tuple[str, list[Image.Image]]:
storyboard = _make_storyboard_image(frames, columns=3)
prompt_suffix = (
"The attached image is a storyboard containing six frames arranged left-to-right, "
"top-to-bottom as frame0, frame1, frame2, frame3, frame4, frame5."
)
return prompt_suffix, [storyboard]
class LocalJudgeClient:
def __init__(self, base_url: str, api_key: str, model: str, timeout: int = 600):
self.base_url = base_url.rstrip("/")
self.api_key = api_key
self.model = model
self.timeout = timeout
def _build_scoring_prompt(self, task_prompt: str) -> str:
return (
"You are an expert evaluator for GUI image editing and GUI trajectory generation.\n"
"Evaluate whether the generated image(s) satisfy the task.\n\n"
"Score these five dimensions from 0 to 5:\n"
"- goal: whether the user goal is completed correctly\n"
"- logic: whether the transition/state change is logically correct\n"
"- cons: whether unrelated regions remain consistent\n"
"- ui: whether the UI layout/components remain realistic and coherent\n"
"- qual: whether the images are visually clear and artifact-free\n\n"
"Return JSON only. Do not add any prose outside JSON.\n"
"Use exactly this schema:\n"
"{\n"
' "goal": 0,\n'
' "logic": 0,\n'
' "cons": 0,\n'
' "ui": 0,\n'
' "qual": 0,\n'
' "reasoning": "short explanation"\n'
"}\n\n"
"Scoring task:\n"
f"{task_prompt}"
)
def _request_text(self, prompt: str, images: list[Image.Image]) -> str:
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
for image in images:
content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
response = requests.post(
build_openai_url(self.base_url, "/chat/completions"),
json={
"model": self.model,
"messages": [{"role": "user", "content": content}],
"temperature": 0,
},
headers={
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
},
timeout=self.timeout,
)
response.raise_for_status()
message_content = response.json()["choices"][0]["message"]["content"]
if isinstance(message_content, list):
return "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
return str(message_content)
def evaluate(self, *, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
primary_prompt = self._build_scoring_prompt(prompt)
text = self._request_text(primary_prompt, images)
try:
return extract_json_object(text)
except ValueError:
retry_prompt = (
self._build_scoring_prompt(prompt) + "\n\nYour previous response was not valid JSON. "
"Return only the JSON object with integer scores."
)
retry_text = self._request_text(retry_prompt, images)
try:
return extract_json_object(retry_text)
except ValueError:
return {
"goal": 0,
"logic": 0,
"cons": 0,
"ui": 0,
"qual": 0,
"reasoning": retry_text.strip() or text.strip() or "Judge response was not valid JSON.",
}
class GEBenchRunner:
def __init__(
self,
*,
dataset_root: Path,
output_root: Path,
base_url: str,
model: str,
api_key: str = "EMPTY",
width: int = 768,
height: int = 576,
num_inference_steps: int = 8,
output_compression: int | None = 98,
guidance_scale: float | None = None,
seed: int | None = 42,
):
self.dataset_root = dataset_root
self.output_root = output_root
self.model = model
self.width = width
self.height = height
self.num_inference_steps = num_inference_steps
self.output_compression = output_compression
self.guidance_scale = guidance_scale
self.seed = seed
self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
def generate(
self,
*,
data_type: str,
workers: int = 1,
max_samples: int | None = None,
samples_per_type: int | None = None,
) -> list[dict[str, Any]]:
sample_specs = select_balanced_gebench_samples(
{data_type: _iter_sample_specs(self.dataset_root, data_type)},
samples_per_type=samples_per_type,
)[data_type]
if max_samples is not None:
sample_specs = sample_specs[:max_samples]
results: list[dict[str, Any]] = []
if workers <= 1:
for sample_spec in sample_specs:
result = self._generate_one(data_type, sample_spec)
if result:
results.append(result)
return results
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = [executor.submit(self._generate_one, data_type, sample_spec) for sample_spec in sample_specs]
for future in as_completed(futures):
result = future.result()
if result:
results.append(result)
return results
def _generate_one(self, data_type: str, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
sample_path = sample_spec.sample_path
metadata = sample_spec.metadata
lang_device = sample_spec.lang_device
sample_name = sample_spec.sample_name
output_dir = ensure_dir(self.output_root / TYPE_TO_FOLDER[data_type] / lang_device / sample_name)
if data_type == "type1":
output_path = output_dir / "generated.png"
if output_path.exists():
return {
"data_type": data_type,
"sample_name": f"{lang_device}/{sample_name}",
"output_path": str(output_path),
}
source = _resolve_referenced_image(
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
)
if source is None:
return None
generated = self.client.generate_image_edit(
model=self.model,
prompt=_type1_prompt(metadata),
images=source,
width=self.width,
height=self.height,
num_inference_steps=self.num_inference_steps,
output_compression=self.output_compression,
guidance_scale=self.guidance_scale,
seed=self.seed,
)
save_image(output_path, generated)
return {
"data_type": data_type,
"sample_name": f"{lang_device}/{sample_name}",
"output_path": str(output_path),
}
if data_type == "type2":
goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
source = _resolve_referenced_image(
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
)
if source is None:
return None
frame0_path = output_dir / "frame0.png"
if not frame0_path.exists():
save_image(frame0_path, source)
previous = source
for step_num in range(1, 6):
frame_path = output_dir / f"frame{step_num}.png"
if frame_path.exists():
previous = Image.open(frame_path).convert("RGB")
continue
generated = self.client.generate_image_edit(
model=self.model,
prompt=_type2_prompt(goal, step_num),
images=previous,
width=self.width,
height=self.height,
num_inference_steps=self.num_inference_steps,
output_compression=self.output_compression,
guidance_scale=self.guidance_scale,
seed=self.seed,
)
save_image(frame_path, generated)
previous = generated
output_path = output_dir / "frame5.png"
return {
"data_type": data_type,
"sample_name": f"{lang_device}/{sample_name}",
"output_path": str(output_path),
}
if data_type in {"type3", "type4"}:
steps = _trajectory_steps(metadata)
frame0_path = output_dir / "frame0.png"
if frame0_path.exists():
previous = Image.open(frame0_path).convert("RGB")
else:
previous = self.client.generate_text_to_image(
model=self.model,
prompt=_type34_initial_prompt(metadata, steps[0] if steps else {}),
width=self.width,
height=self.height,
num_inference_steps=self.num_inference_steps,
output_compression=self.output_compression,
guidance_scale=self.guidance_scale,
seed=self.seed,
)
save_image(frame0_path, previous)
for step_num in range(1, 6):
frame_path = output_dir / f"frame{step_num}.png"
if frame_path.exists():
previous = Image.open(frame_path).convert("RGB")
continue
step_info = steps[step_num - 1] if step_num - 1 < len(steps) else {}
generated = self.client.generate_image_edit(
model=self.model,
prompt=_type34_next_prompt(step_num, step_info),
images=previous,
width=self.width,
height=self.height,
num_inference_steps=self.num_inference_steps,
output_compression=self.output_compression,
guidance_scale=self.guidance_scale,
seed=self.seed,
)
save_image(frame_path, generated)
previous = generated
output_path = output_dir / "frame5.png"
return {
"data_type": data_type,
"sample_name": f"{lang_device}/{sample_name}",
"output_path": str(output_path),
}
if data_type == "type5":
output_path = output_dir / "generated.png"
if output_path.exists():
return {
"data_type": data_type,
"sample_name": f"{lang_device}/{sample_name}",
"output_path": str(output_path),
}
source = _resolve_referenced_image(
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
)
if source is None:
return None
generated = self.client.generate_image_edit(
model=self.model,
prompt=_type5_prompt(metadata),
images=source,
width=self.width,
height=self.height,
num_inference_steps=self.num_inference_steps,
output_compression=self.output_compression,
guidance_scale=self.guidance_scale,
seed=self.seed,
)
save_image(output_path, generated)
return {
"data_type": data_type,
"sample_name": f"{lang_device}/{sample_name}",
"output_path": str(output_path),
}
raise ValueError(f"Unsupported data type: {data_type}")
class GEBenchEvaluator:
def __init__(self, *, dataset_root: Path, output_root: Path, judge: LocalJudgeClient):
self.dataset_root = dataset_root
self.output_root = output_root
self.judge = judge
def evaluate(
self,
*,
data_type: str,
workers: int = 1,
max_samples: int | None = None,
samples_per_type: int | None = None,
) -> dict[str, Any]:
output_type_dir = self.output_root / TYPE_TO_FOLDER[data_type]
sample_specs_by_name = {
(spec.lang_device, spec.sample_name): spec for spec in _iter_sample_specs(self.dataset_root, data_type)
}
if not output_type_dir.exists():
payload = {"data_type": data_type, "results": [], "summary": summarize_gebench_results([])}
write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
return payload
sample_dirs = [
sample_dir
for lang_dir in sorted(path for path in output_type_dir.iterdir() if path.is_dir())
for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir())
if (lang_dir.name, sample_dir.name) in sample_specs_by_name
]
sample_dirs = select_balanced_gebench_samples(
{data_type: sample_dirs},
samples_per_type=samples_per_type,
)[data_type]
if max_samples is not None:
sample_dirs = sample_dirs[:max_samples]
results: list[dict[str, Any]] = []
if workers <= 1:
for sample_dir in sample_dirs:
result = self._evaluate_one(
data_type,
sample_dir,
sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
)
if result:
results.append(result)
else:
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = [
executor.submit(
self._evaluate_one,
data_type,
sample_dir,
sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
)
for sample_dir in sample_dirs
]
for future in as_completed(futures):
result = future.result()
if result:
results.append(result)
payload = {"data_type": data_type, "results": results, "summary": summarize_gebench_results(results)}
write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
return payload
def _evaluate_one(self, data_type: str, sample_dir: Path, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
lang_device = sample_dir.parent.name
sample_name = sample_dir.name
dataset_sample = sample_spec.sample_path
metadata = sample_spec.metadata
if data_type == "type1":
source = _resolve_referenced_image(
metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
)
generated_path = find_first_image(sample_dir)
if source is None or generated_path is None:
return None
generated = Image.open(generated_path).convert("RGB")
raw_scores = self.judge.evaluate(prompt=_type1_prompt(metadata), images=[source, generated])
elif data_type == "type2":
frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
prompt_suffix, judge_images = _trajectory_judge_payload(frames)
raw_scores = self.judge.evaluate(
prompt=f"Evaluate a six-frame GUI trajectory.\nTask: {goal}\n{prompt_suffix}",
images=judge_images,
)
elif data_type in {"type3", "type4"}:
frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
instruction = _text_or_default(metadata.get("instruction") or metadata.get("caption"), "Complete the task.")
prompt_suffix, judge_images = _trajectory_judge_payload(frames)
raw_scores = self.judge.evaluate(
prompt=f"Evaluate a six-frame GUI trajectory.\nInstruction: {instruction}\n{prompt_suffix}",
images=judge_images,
)
elif data_type == "type5":
source = _resolve_referenced_image(
metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
)
generated_path = find_first_image(sample_dir)
if source is None or generated_path is None:
return None
generated = Image.open(generated_path).convert("RGB")
raw_scores = self.judge.evaluate(prompt=_type5_prompt(metadata), images=[source, generated])
else:
raise ValueError(f"Unsupported data type: {data_type}")
scores = _normalize_scores(raw_scores)
return {
"sample_name": f"{lang_device}/{sample_name}",
"data_type": data_type,
"scores": scores,
"overall": _compute_overall(scores),
"raw_scores": raw_scores,
}
def _data_types_arg(value: str) -> list[str]:
return list(TYPE_TO_FOLDER.keys()) if value == "all" else [value]
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Run local GEBench generation and scoring against vLLM-Omni.")
subparsers = parser.add_subparsers(dest="command", required=True)
generate = subparsers.add_parser("generate")
generate.add_argument("--dataset-root", type=Path, required=True)
generate.add_argument("--output-root", type=Path, required=True)
generate.add_argument("--base-url", type=str, required=True)
generate.add_argument("--model", type=str, required=True)
generate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
generate.add_argument("--api-key", type=str, default="EMPTY")
generate.add_argument("--width", type=int, default=768)
generate.add_argument("--height", type=int, default=576)
generate.add_argument("--num-inference-steps", type=int, default=8)
generate.add_argument("--output-compression", type=int, default=98)
generate.add_argument("--guidance-scale", type=float, default=None)
generate.add_argument("--seed", type=int, default=42)
generate.add_argument("--workers", type=int, default=1)
generate.add_argument("--max-samples", type=int, default=None)
generate.add_argument("--samples-per-type", type=int, default=None)
evaluate = subparsers.add_parser("evaluate")
evaluate.add_argument("--dataset-root", type=Path, required=True)
evaluate.add_argument("--output-root", type=Path, required=True)
evaluate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
evaluate.add_argument("--judge-base-url", type=str, required=True)
evaluate.add_argument("--judge-model", type=str, required=True)
evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
evaluate.add_argument("--workers", type=int, default=1)
evaluate.add_argument("--max-samples", type=int, default=None)
evaluate.add_argument("--samples-per-type", type=int, default=None)
summarize = subparsers.add_parser("summarize")
summarize.add_argument("--output-root", type=Path, required=True)
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
if args.command == "generate":
runner = GEBenchRunner(
dataset_root=args.dataset_root,
output_root=args.output_root,
base_url=args.base_url,
model=args.model,
api_key=args.api_key,
width=args.width,
height=args.height,
num_inference_steps=args.num_inference_steps,
output_compression=args.output_compression,
guidance_scale=args.guidance_scale,
seed=args.seed,
)
records: list[dict[str, Any]] = []
for data_type in _data_types_arg(args.data_type):
records.extend(
runner.generate(
data_type=data_type,
workers=args.workers,
max_samples=args.max_samples,
samples_per_type=args.samples_per_type,
)
)
payload = {"records": records, "summary": summarize_generated_records(records)}
write_json(args.output_root / "generation_manifest.json", payload)
return 0
if args.command == "evaluate":
judge = LocalJudgeClient(
base_url=args.judge_base_url,
api_key=args.judge_api_key,
model=args.judge_model,
)
evaluator = GEBenchEvaluator(dataset_root=args.dataset_root, output_root=args.output_root, judge=judge)
combined_results: list[dict[str, Any]] = []
for data_type in _data_types_arg(args.data_type):
payload = evaluator.evaluate(
data_type=data_type,
workers=args.workers,
max_samples=args.max_samples,
samples_per_type=args.samples_per_type,
)
combined_results.extend(payload["results"])
_write_json_with_timestamp(
args.output_root / "evaluations" / "summary.json",
{"summary": summarize_gebench_results(combined_results)},
)
return 0
if args.command == "summarize":
generation_summary = collect_gebench_generation_summary(args.output_root)
evaluation_dir = args.output_root / "evaluations"
result_records: list[dict[str, Any]] = []
if evaluation_dir.exists():
for file_path in sorted(evaluation_dir.glob("type*.json")):
payload = load_json(file_path)
result_records.extend(payload.get("results", []))
payload: dict[str, Any] = {"generation": generation_summary}
if result_records:
payload["evaluation"] = summarize_gebench_results(result_records)
_write_json_with_timestamp(args.output_root / "summary.json", payload)
print(json.dumps(payload, indent=2, ensure_ascii=False))
return 0
parser.error(f"Unknown command: {args.command}")
return 1
@@ -1,13 +0,0 @@
# ruff: noqa: E402, I001
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[3]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from benchmarks.accuracy.text_to_image.gbench import main
if __name__ == "__main__":
raise SystemExit(main())
+3 -2
View File
@@ -109,6 +109,8 @@ async def async_request_chat_completions(
resp_json = await response.json()
output.response_body = resp_json
output.success = True
if "peak_memory_mb" in resp_json:
output.peak_memory_mb = resp_json["peak_memory_mb"]
try:
choices = resp_json.get("choices", [])
if choices and isinstance(choices, list):
@@ -118,8 +120,7 @@ async def async_request_chat_completions(
if content and isinstance(content, list) and len(content) > 0:
first_item = content[0]
if isinstance(first_item, dict):
output.stage_durations = first_item.get("stage_durations") or {}
output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
output.stage_durations = first_item.get("stage_durations")
except (IndexError, TypeError, AttributeError):
pass
else:
@@ -1,4 +1,4 @@
# adapted from fastvideo
# adapted from sglang and fastvideo
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -289,9 +289,6 @@ class VBenchDataset(BaseDataset):
def _resize_data(self, data: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Resize data to match num_prompts."""
if not data:
raise ValueError("No benchmark data available. Install Pillow or provide --dataset-path.")
if not self.args.num_prompts:
return data
@@ -515,14 +512,6 @@ class TraceDataset(BaseDataset):
timestamp = self._coerce_float(row.get("timestamp"))
slo_ms = self._coerce_float(row.get("slo_ms"))
image_paths = row.get("image_paths")
if not image_paths:
single = row.get("image_path")
image_paths = [single] if single else None
if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i"]:
raise ValueError(
f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row."
)
override_w = self.args.width
override_h = self.args.height
@@ -725,19 +714,19 @@ async def iter_requests(
requests_list: list[RequestFuncInput],
request_rate: float,
) -> AsyncGenerator[RequestFuncInput, None]:
"""Yield requests using a Poisson process if request_rate is set.
"""Yield requests using a fixed interval if request_rate is set.
- If request_rate is inf, all requests are yielded immediately (no sleep).
- Otherwise, inter-arrival times follow an exponential distribution.
- Otherwise, requests are emitted at a fixed cadence of 1 / request_rate seconds.
"""
if request_rate != float("inf"):
if request_rate <= 0:
raise ValueError(f"request_rate must be positive or inf, got {request_rate}.")
interval_s = 1.0 / float(request_rate)
for i, req in enumerate(requests_list):
if request_rate != float("inf") and i > 0:
interval_s = random.expovariate(request_rate)
await asyncio.sleep(interval_s)
yield req
@@ -756,15 +745,6 @@ def calculate_metrics(
latencies = [o.latency for o in success_outputs]
peak_memories = [o.peak_memory_mb for o in success_outputs if o.peak_memory_mb > 0]
# Aggregate per-stage durations across all successful requests that reported them.
stage_duration_lists: dict[str, list[float]] = {}
for o in success_outputs:
for stage, duration in (o.stage_durations or {}).items():
stage_duration_lists.setdefault(stage, []).append(duration)
stage_durations_mean = {s: float(np.mean(v)) for s, v in stage_duration_lists.items()}
stage_durations_p50 = {s: float(np.percentile(v, 50)) for s, v in stage_duration_lists.items()}
stage_durations_p99 = {s: float(np.percentile(v, 99)) for s, v in stage_duration_lists.items()}
metrics = {
"duration": total_duration,
"completed_requests": num_success,
@@ -778,9 +758,6 @@ def calculate_metrics(
"peak_memory_mb_max": max(peak_memories) if peak_memories else 0,
"peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0,
"peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0,
"stage_durations_mean": stage_durations_mean,
"stage_durations_p50": stage_durations_p50,
"stage_durations_p99": stage_durations_p99,
}
if slo_enabled:
@@ -904,8 +881,6 @@ async def benchmark(args):
warm_req,
num_inference_steps=args.warmup_num_inference_steps,
)
if args.task == "t2v":
warm_req = replace(warm_req, num_frames=1)
warm_out = await limited_request_func(warm_req, session, None)
warmup_pairs.append((warm_req, warm_out))
@@ -978,12 +953,6 @@ async def benchmark(args):
print("{:<40} {:<15.2f}".format("Peak Memory Mean (MB):", metrics["peak_memory_mb_mean"]))
print("{:<40} {:<15.2f}".format("Peak Memory Median (MB):", metrics["peak_memory_mb_median"]))
if metrics["stage_durations_mean"]:
print(f"{'-' * 50}")
print("Stage Durations Mean (s):")
for stage, val in metrics["stage_durations_mean"].items():
print("{:<40} {:<15.4f}".format(f" {stage}:", val))
print("\n" + "=" * 60)
if args.output_file:
@@ -1,460 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Benchmark quantization quality loss for diffusion models (image & video).
Generates outputs with BF16 (baseline) and a quantized config using the same
seed, then computes LPIPS perceptual distance between them. Results are printed
as a Markdown table ready to paste into a PR description.
Requirements:
pip install lpips Pillow numpy
Image example (text-to-image):
python benchmarks/diffusion/quantization_quality.py \
--model Tongyi-MAI/Z-Image-Turbo \
--task t2i \
--quantization fp8 \
--prompts \
"an aerial view of a coral reef with crystal clear turquoise water" \
"a campfire in a dark forest with sparks rising into a starry sky" \
"a gourmet dessert plate with chocolate mousse and gold leaf" \
--height 1024 --width 1024 \
--num-inference-steps 50 --seed 42
Video example (text-to-video):
python benchmarks/diffusion/quantization_quality.py \
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--task t2v \
--quantization fp8 \
--prompts \
"A serene lakeside sunrise with mist over the water" \
"A cat walking across a wooden bridge in autumn" \
--height 720 --width 1280 \
--num-frames 81 --num-inference-steps 40 --seed 42
Multiple quantization methods:
python benchmarks/diffusion/quantization_quality.py \
--model Tongyi-MAI/Z-Image-Turbo \
--task t2i \
--quantization fp8 int8 bitsandbytes \
--prompts "a cup of coffee on the table" \
--height 1024 --width 1024 \
--num-inference-steps 50 --seed 42
Output directory structure (--output-dir, default: ./quant_bench_output):
quant_bench_output/
baseline/ # BF16 outputs
<method>/ # Quantized outputs per method
results.md # Markdown table
"""
import argparse
import gc
import time
from pathlib import Path
import numpy as np
import torch
def compute_lpips_images(
baseline_images: list,
quantized_images: list,
net: str = "alex",
) -> list[float]:
"""Compute LPIPS between paired lists of PIL images."""
import lpips
from torchvision import transforms
loss_fn = lpips.LPIPS(net=net).eval()
if torch.cuda.is_available():
loss_fn = loss_fn.cuda()
transform = transforms.Compose(
[
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
]
)
scores = []
for img_bl, img_qt in zip(baseline_images, quantized_images):
t_bl = transform(img_bl.convert("RGB")).unsqueeze(0)
t_qt = transform(img_qt.convert("RGB")).unsqueeze(0)
if torch.cuda.is_available():
t_bl, t_qt = t_bl.cuda(), t_qt.cuda()
with torch.no_grad():
score = loss_fn(t_bl, t_qt).item()
scores.append(score)
return scores
def compute_lpips_video(
baseline_frames: np.ndarray,
quantized_frames: np.ndarray,
net: str = "alex",
) -> float:
"""Compute mean per-frame LPIPS for a video pair.
Args:
baseline_frames: (F, H, W, C) float array in [0, 1].
quantized_frames: same shape.
Returns:
Mean LPIPS across all frames.
"""
import lpips
loss_fn = lpips.LPIPS(net=net).eval()
if torch.cuda.is_available():
loss_fn = loss_fn.cuda()
num_frames = min(len(baseline_frames), len(quantized_frames))
scores = []
for i in range(num_frames):
# Convert (H, W, C) float [0,1] -> (1, C, H, W) float [-1, 1]
f_bl = torch.from_numpy(baseline_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
f_qt = torch.from_numpy(quantized_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
if torch.cuda.is_available():
f_bl, f_qt = f_bl.cuda(), f_qt.cuda()
with torch.no_grad():
score = loss_fn(f_bl, f_qt).item()
scores.append(score)
return float(np.mean(scores))
def _build_omni_kwargs(args, quantization=None):
"""Build kwargs dict for Omni() constructor."""
from vllm_omni.diffusion.data import DiffusionParallelConfig
parallel_config = DiffusionParallelConfig(
ulysses_degree=args.ulysses_degree,
ring_degree=args.ring_degree,
tensor_parallel_size=args.tensor_parallel_size,
)
kwargs = {
"model": args.model,
"parallel_config": parallel_config,
"enforce_eager": args.enforce_eager,
}
if quantization:
kwargs["quantization_config"] = quantization
return kwargs
def _generate_image(omni, args, prompt, seed):
"""Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
torch.cuda.reset_peak_memory_stats()
start = time.perf_counter()
outputs = omni.generate(
{"prompt": prompt},
OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
generator=generator,
num_inference_steps=args.num_inference_steps,
),
)
elapsed = time.perf_counter() - start
peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
first = outputs[0]
req_out = first.request_output[0] if hasattr(first, "request_output") else first
img = req_out.images[0]
return img, elapsed, peak_mem
def _generate_video(omni, args, prompt, seed):
"""Generate a video and return (np.ndarray [F,H,W,C], time_seconds, memory_gib)."""
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
torch.cuda.reset_peak_memory_stats()
start = time.perf_counter()
outputs = omni.generate(
{"prompt": prompt, "negative_prompt": ""},
OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
generator=generator,
guidance_scale=args.guidance_scale,
num_inference_steps=args.num_inference_steps,
num_frames=args.num_frames,
),
)
elapsed = time.perf_counter() - start
peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
first = outputs[0]
if hasattr(first, "request_output") and isinstance(first.request_output, list):
inner = first.request_output[0]
if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
frames = inner.images[0] if inner.images else None
else:
frames = inner
elif hasattr(first, "images") and first.images:
frames = first.images
else:
raise ValueError("Could not extract video frames from output.")
if isinstance(frames, torch.Tensor):
video = frames.detach().cpu()
if video.dim() == 5:
video = video[0].permute(1, 2, 3, 0) if video.shape[1] in (3, 4) else video[0]
elif video.dim() == 4 and video.shape[0] in (3, 4):
video = video.permute(1, 2, 3, 0)
if video.is_floating_point():
video = video.clamp(-1, 1) * 0.5 + 0.5
frames_array = video.float().numpy()
else:
frames_array = np.asarray(frames)
if frames_array.ndim == 5:
frames_array = frames_array[0]
return frames_array, elapsed, peak_mem
def _unload_omni(omni):
"""Delete Omni instance and free GPU memory."""
del omni
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.synchronize()
def run_benchmark(args):
from vllm_omni.entrypoints.omni import Omni
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
is_video = args.task == "t2v"
prompts = args.prompts
seed = args.seed
# Determine configs to benchmark
configs = [] # list of (label, quantization_method)
for method in args.quantization:
configs.append((method, method))
# --- Baseline run ---
print("\n" + "=" * 60)
print("Running BF16 baseline...")
print("=" * 60)
bl_kwargs = _build_omni_kwargs(args, quantization=None)
omni_bl = Omni(**bl_kwargs)
baseline_outputs = {} # prompt -> (output, time, mem)
for prompt in prompts:
print(f" Generating: {prompt[:60]}...")
if is_video:
out, t, mem = _generate_video(omni_bl, args, prompt, seed)
else:
out, t, mem = _generate_image(omni_bl, args, prompt, seed)
baseline_outputs[prompt] = (out, t, mem)
bl_avg_time = np.mean([v[1] for v in baseline_outputs.values()])
bl_mem = baseline_outputs[prompts[0]][2] # use first prompt's memory
_unload_omni(omni_bl)
# Save baseline outputs
bl_dir = output_dir / "baseline"
bl_dir.mkdir(parents=True, exist_ok=True)
for i, prompt in enumerate(prompts):
out = baseline_outputs[prompt][0]
if is_video:
try:
from diffusers.utils import export_to_video
frames_list = list(out) if isinstance(out, np.ndarray) and out.ndim == 4 else out
export_to_video(frames_list, str(bl_dir / f"prompt_{i}.mp4"), fps=args.fps)
except ImportError:
np.save(bl_dir / f"prompt_{i}.npy", out)
else:
out.save(bl_dir / f"prompt_{i}.png")
# --- Quantized runs ---
all_results = [] # list of dicts
for config_label, quant_method in configs:
print(f"\n{'=' * 60}")
print(f"Running: {config_label}...")
print("=" * 60)
qt_kwargs = _build_omni_kwargs(args, quantization=quant_method)
omni_qt = Omni(**qt_kwargs)
qt_outputs = {}
for prompt in prompts:
print(f" Generating: {prompt[:60]}...")
if is_video:
out, t, mem = _generate_video(omni_qt, args, prompt, seed)
else:
out, t, mem = _generate_image(omni_qt, args, prompt, seed)
qt_outputs[prompt] = (out, t, mem)
qt_avg_time = np.mean([v[1] for v in qt_outputs.values()])
qt_mem = qt_outputs[prompts[0]][2]
_unload_omni(omni_qt)
# Save quantized outputs
qt_dir = output_dir / config_label.replace(" ", "_")
qt_dir.mkdir(parents=True, exist_ok=True)
# Compute LPIPS per prompt
per_prompt = []
for i, prompt in enumerate(prompts):
bl_out = baseline_outputs[prompt][0]
qt_out = qt_outputs[prompt][0]
if is_video:
lpips_score = compute_lpips_video(bl_out, qt_out, net=args.lpips_net)
try:
from diffusers.utils import export_to_video
frames_list = list(qt_out) if isinstance(qt_out, np.ndarray) and qt_out.ndim == 4 else qt_out
export_to_video(frames_list, str(qt_dir / f"prompt_{i}.mp4"), fps=args.fps)
except ImportError:
np.save(qt_dir / f"prompt_{i}.npy", qt_out)
else:
lpips_score = compute_lpips_images([bl_out], [qt_out], net=args.lpips_net)[0]
qt_out.save(qt_dir / f"prompt_{i}.png")
per_prompt.append({"prompt": prompt, "lpips": lpips_score})
mean_lpips = np.mean([p["lpips"] for p in per_prompt])
speedup = bl_avg_time / qt_avg_time if qt_avg_time > 0 else float("inf")
mem_reduction = (bl_mem - qt_mem) / bl_mem * 100
all_results.append(
{
"config": config_label,
"avg_time": qt_avg_time,
"speedup": speedup,
"memory_gib": qt_mem,
"mem_reduction_pct": mem_reduction,
"mean_lpips": mean_lpips,
"per_prompt": per_prompt,
}
)
# --- Print results ---
print("\n\n")
print("=" * 80)
print("RESULTS")
print("=" * 80)
# Summary table
lines = []
lines.append(f"## Quantization Quality Benchmark — {args.model.split('/')[-1]}")
lines.append(
f"Setup: {args.height}x{args.width}, {args.num_inference_steps} steps, "
f"seed={args.seed}, LPIPS ({args.lpips_net})"
)
if is_video:
lines.append(f"Video: {args.num_frames} frames")
lines.append("")
lines.append("### Summary")
lines.append("")
lines.append("| Config | Avg Time | Speedup | Memory (GiB) | Mem Reduction | Mean LPIPS |")
lines.append("|--------|----------|---------|--------------|---------------|------------|")
lines.append(f"| BF16 baseline | {bl_avg_time:.2f}s | 1.00x | {bl_mem:.2f} | — | (ref) |")
for r in all_results:
lines.append(
f"| {r['config']} | {r['avg_time']:.2f}s | {r['speedup']:.2f}x "
f"| {r['memory_gib']:.2f} | {r['mem_reduction_pct']:.0f}% "
f"| {r['mean_lpips']:.4f} |"
)
lines.append("")
lines.append("> LPIPS < 0.01 = imperceptible, > 0.1 = clearly noticeable.")
lines.append("")
# Per-prompt table
if len(prompts) > 1:
lines.append("### Per-Prompt LPIPS")
lines.append("")
header = "| Prompt |"
sep = "|--------|"
for r in all_results:
header += f" {r['config']} |"
sep += "--------|"
lines.append(header)
lines.append(sep)
for i, prompt in enumerate(prompts):
short = prompt[:50] + "..." if len(prompt) > 50 else prompt
row = f"| {short} |"
for r in all_results:
row += f" {r['per_prompt'][i]['lpips']:.4f} |"
lines.append(row)
lines.append("")
md = "\n".join(lines)
print(md)
# Save markdown
results_path = output_dir / "results.md"
results_path.write_text(md, encoding="utf-8")
print(f"\nResults saved to {results_path}")
print(f"Baseline outputs in {bl_dir}")
for r in all_results:
qt_dir = output_dir / r["config"].replace(" ", "_")
print(f"Quantized outputs in {qt_dir}")
def parse_args():
parser = argparse.ArgumentParser(
description="Benchmark quantization quality loss for diffusion models.",
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument("--model", required=True, help="Model name or local path.")
parser.add_argument(
"--task",
default="t2i",
choices=["t2i", "t2v"],
help="Task type: t2i (text-to-image) or t2v (text-to-video).",
)
parser.add_argument(
"--quantization",
nargs="+",
required=True,
help="One or more quantization methods to benchmark (e.g. fp8 int8 bitsandbytes).",
)
parser.add_argument(
"--prompts",
nargs="+",
default=["a cup of coffee on the table"],
help="One or more prompts to generate.",
)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--height", type=int, default=1024)
parser.add_argument("--width", type=int, default=1024)
parser.add_argument("--num-inference-steps", type=int, default=50)
parser.add_argument("--num-frames", type=int, default=81, help="Number of video frames (t2v only).")
parser.add_argument("--fps", type=int, default=24, help="Video FPS for saving (t2v only).")
parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (used for video).")
parser.add_argument("--output-dir", type=str, default="./quant_bench_output", help="Directory to save outputs.")
parser.add_argument(
"--lpips-net",
type=str,
default="alex",
choices=["alex", "vgg", "squeeze"],
help="LPIPS backbone network.",
)
parser.add_argument("--ulysses-degree", type=int, default=1)
parser.add_argument("--ring-degree", type=int, default=1)
parser.add_argument("--tensor-parallel-size", type=int, default=1)
parser.add_argument("--enforce-eager", action="store_true")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
run_benchmark(args)
+23 -4
View File
@@ -1,5 +1,5 @@
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
ARG VLLM_BASE_TAG=v0.18.0
ARG VLLM_BASE_TAG=v0.17.0
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
ARG APP_DIR=/workspace/vllm-omni
WORKDIR ${APP_DIR}
@@ -11,10 +11,29 @@ RUN apt-get update && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
# Install vllm-omni into the same uv-managed Python environment used by the base image.
# Use bash -c so that $(python3 -c ...) is expanded inside the container.
RUN uv pip install --system --no-cache-dir ".[dev]"
RUN uv pip uninstall --system -y vllm || true
# Install vLLM from precompiled wheel at the selected commit.
# Must use direct URL because the wheel has a PEP 440 local version identifier
# (e.g. +g0a0a1a198) which pip/uv refuse to install from a PEP 503 package index.
ENV VLLM_PRECOMPILED_WHEEL_COMMIT=89138b21cc246ae944c741d5c399c148e2b770ab
RUN VLLM_WHEEL_URL=$(python3 -c "import urllib.request,re; \
html=urllib.request.urlopen('https://wheels.vllm.ai/${VLLM_PRECOMPILED_WHEEL_COMMIT}/vllm/').read().decode(); \
m=re.search(r'>(\S+x86_64\.whl)<',html); \
print('https://wheels.vllm.ai/${VLLM_PRECOMPILED_WHEEL_COMMIT}/'+m.group(1).replace('+','%2B'))") && \
echo "Installing vLLM from: ${VLLM_WHEEL_URL}" && \
uv pip install --system --force-reinstall "${VLLM_WHEEL_URL}"
RUN uv pip install --system ".[dev]"
RUN uv pip install --system --upgrade \
"flashinfer-cubin==0.6.6" \
"nvidia-cublas-cu12==12.9.1.4" \
"numpy==2.2.6"
RUN uv pip install --system --upgrade \
"flashinfer-jit-cache==0.6.6" \
--index-url https://flashinfer.ai/whl/cu129
RUN ln -sf /usr/bin/python3 /usr/bin/python
ENTRYPOINT []
+1 -14
View File
@@ -2,22 +2,9 @@ ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.17.0rc1
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
WORKDIR /vllm-workspace/vllm
RUN git fetch origin --tags && git checkout v0.18.0
WORKDIR /vllm-workspace/vllm-ascend
RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
# Install vllm-ascend
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
source /usr/local/Ascend/nnal/atb/set_env.sh && \
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
python3 -m pip cache purge
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
COPY . .
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+1 -14
View File
@@ -2,22 +2,9 @@ ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
WORKDIR /vllm-workspace/vllm
RUN git fetch origin --tags && git checkout v0.18.0
WORKDIR /vllm-workspace/vllm-ascend
RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
# Install vllm-ascend
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
source /usr/local/Ascend/nnal/atb/set_env.sh && \
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
python3 -m pip cache purge
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
COPY . .
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+2 -2
View File
@@ -1,4 +1,4 @@
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.18.0
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.17.0
FROM ${BASE_IMAGE} AS base
# Declare a variable to know if we want to use the nightly build or the stable build.
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
# we should swap over to use stable release ASAP.
# We should avoid relying on custom commits.
ARG USE_NIGHTLY_BUILD=0
ARG VLLM_VERSION_OR_COMMIT_HASH=89138b21cc246ae944c741d5c399c148e2b770ab
ARG VLLM_VERSION_OR_COMMIT_HASH=2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
ARG ARG_PYTORCH_ROCM_ARCH
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
+1 -1
View File
@@ -76,7 +76,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
ENV UV_INDEX_STRATEGY="unsafe-best-match"
ENV UV_LINK_MODE="copy"
ARG VLLM_VERSION=v0.18.0
ARG VLLM_VERSION=v0.17.0
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
WORKDIR /workspace/vllm
+6 -21
View File
@@ -6,7 +6,6 @@ nav:
- getting_started/installation/*
- Serving:
- OpenAI-Compatible API:
- Diffusion Chat API: serving/diffusion_chat_api.md
- Image Generation: serving/image_generation_api.md
- Image Edit: serving/image_edit_api.md
- Text to Speech: serving/speech_api.md
@@ -29,7 +28,6 @@ nav:
- Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
- Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
- Voxtral TTS Offline Inference: user_guide/examples/offline_inference/voxtral_tts.md
- X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
- Online Serving:
- BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
@@ -54,28 +52,20 @@ nav:
- Features:
- Sleep Mode: features/sleep_mode.md
- Diffusion Features:
- Overview: user_guide/diffusion_features.md
- Feature Compatibility: user_guide/feature_compatibility.md
- Cache Acceleration:
- TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
- Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
- Acceleration Overview: user_guide/diffusion_acceleration.md
- TeaCache: user_guide/diffusion/teacache.md
- Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
- Quantization:
- Overview: user_guide/diffusion/quantization/overview.md
- FP8: user_guide/diffusion/quantization/fp8.md
- Int8: user_guide/diffusion/quantization/int8.md
- GGUF: user_guide/diffusion/quantization/gguf.md
- Parallelism:
- Overview: user_guide/diffusion/parallelism/overview.md
- CFG Parallel: user_guide/diffusion/parallelism/cfg_parallel.md
- Expert Parallel: user_guide/diffusion/parallelism/expert_parallel.md
- Hybrid Sharded Data Parallel: user_guide/diffusion/parallelism/hsdp.md
- Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
- Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
- VAE Patch Parallel: user_guide/diffusion/parallelism/vae_patch_parallel.md
- Step Execution: user_guide/diffusion/step_execution.md
- Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
- CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
- LoRA: user_guide/diffusion/lora.md
- Hybrid Sharded Data Parallel: design/feature/hsdp.md
- Custom Pipeline: features/custom_pipeline.md
- Step Execution: user_guide/diffusion/step_execution.md
- ComfyUI: features/comfyui.md
- Developer Guide:
- General:
@@ -88,7 +78,6 @@ nav:
- contributing/model/adding_omni_model.md
- contributing/model/adding_tts_model.md
- contributing/model/adding_diffusion_model.md
- contributing/model/adding_quantization_model.md
- CI: contributing/ci
- Design Documents:
- design/index.md
@@ -100,18 +89,14 @@ nav:
- design/feature/cfg_parallel.md
- design/feature/sequence_parallel.md
- design/feature/tensor_parallel.md
- design/feature/vae_parallel.md
- design/feature/hsdp.md
- design/feature/cache_dit.md
- design/feature/teacache.md
- design/feature/async_chunk_design.md
- design/feature/vae_parallel.md
- design/feature/diffusion_step_execution.md
- Module Design:
- design/module/ar_module.md
- design/module/dit_module.md
- design/module/entrypoint_module.md
- design/module/async_omni_architecture.md
- Docs Guide: contributing/DOCS_GUIDE.md
- API Reference:
- api/README.md
-41
View File
@@ -6,7 +6,6 @@ Main entry points for vLLM-Omni inference and serving.
- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
- [vllm_omni.entrypoints.cfg_companion_tracker.CfgCompanionTracker][]
- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
- [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
@@ -14,13 +13,11 @@ Main entry points for vLLM-Omni inference and serving.
- [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
- [vllm_omni.entrypoints.omni.Omni][]
- [vllm_omni.entrypoints.omni_base.OmniBase][]
- [vllm_omni.entrypoints.pd_utils.PDDisaggregationMixin][]
## Inputs
Input data structures for multi-modal inputs.
- [vllm_omni.inputs.data.OmniCustomPrompt][]
- [vllm_omni.inputs.data.OmniDiffusionSamplingParams][]
- [vllm_omni.inputs.data.OmniEmbedsPrompt][]
- [vllm_omni.inputs.data.OmniTextPrompt][]
@@ -46,18 +43,8 @@ Engine classes for offline and online inference.
- [vllm_omni.engine.OmniEngineCoreRequest][]
- [vllm_omni.engine.PromptEmbedsPayload][]
- [vllm_omni.engine.arg_utils.OmniEngineArgs][]
- [vllm_omni.engine.async_omni_engine.AsyncOmniEngine][]
- [vllm_omni.engine.mm_outputs.MultimodalCompletionOutput][]
- [vllm_omni.engine.mm_outputs.MultimodalPayload][]
- [vllm_omni.engine.orchestrator.Orchestrator][]
- [vllm_omni.engine.orchestrator.OrchestratorRequestState][]
- [vllm_omni.engine.output_modality.OutputModality][]
- [vllm_omni.engine.output_modality.TensorAccumulationStrategy][]
- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
- [vllm_omni.engine.output_processor.OmniRequestState][]
- [vllm_omni.engine.stage_engine_core_client.StageEngineCoreClient][]
- [vllm_omni.engine.stage_init_utils.StageMetadata][]
- [vllm_omni.engine.stage_init_utils.StartedLlmStage][]
## Core
@@ -69,19 +56,6 @@ Core scheduling and caching components.
- [vllm_omni.core.sched.output.OmniCachedRequestData][]
- [vllm_omni.core.sched.output.OmniNewRequestData][]
- [vllm_omni.core.sched.output.OmniSchedulerOutput][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.BASECFM][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.CausalConditionalCFM][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.CausalMaskedDiffWithDiT][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.ConditionalCFM][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConv1d][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConv1dUpsample][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConvRNNF0Predictor][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalHiFTGenerator][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.HiFTGenerator][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SineGen][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SineGen2][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.Snake][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SourceModuleHnNSF][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.EuclideanCodebook][]
@@ -93,17 +67,9 @@ Core scheduling and caching components.
Configuration classes.
- [vllm_omni.config.model.OmniModelConfig][]
- [vllm_omni.config.stage_config.ModelPipeline][]
- [vllm_omni.config.stage_config.StageConfig][]
- [vllm_omni.config.stage_config.StageConfigFactory][]
- [vllm_omni.config.stage_config.StageType][]
- [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
- [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
- [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
- [vllm_omni.model_executor.models.cosyvoice3.config.CosyVoice3Config][]
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechConfig][]
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechFastARConfig][]
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechSlowARConfig][]
- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioConfig][]
- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioTokenizerConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSConfig][]
@@ -117,10 +83,6 @@ Configuration classes.
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
## Workers
@@ -131,8 +93,6 @@ Worker classes and model runners for distributed inference.
- [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
- [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
- [vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase][]
- [vllm_omni.diffusion.worker.utils.DiffusionRequestState][]
- [vllm_omni.diffusion.worker.utils.RunnerOutput][]
- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.ExecuteModelState][]
- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.NPUARModelRunner][]
- [vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker][]
@@ -150,6 +110,5 @@ Worker classes and model runners for distributed inference.
- [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
- [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
- [vllm_omni.worker.gpu_memory_utils.parse_cuda_visible_devices][]
- [vllm_omni.worker.gpu_model_runner.CUDAGraphWrapper][]
- [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
- [vllm_omni.worker.mixins.OmniWorkerMixin][]
Binary file not shown.

Before

Width:  |  Height:  |  Size: 99 KiB

After

Width:  |  Height:  |  Size: 97 KiB

+4 -1
View File
@@ -20,4 +20,7 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
## Optimization Features
- **[Diffusion Features Overview](../user_guide/diffusion_features.md)** - Complete overview of all diffusion model features and supported models
- **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
- **[Cache-DiT Configuration](../user_guide/diffusion/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
- **[Parallelism Configuration](../user_guide/diffusion/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models
- **[CPU Offloading](../user_guide/diffusion/cpu_offload_diffusion.md)** - Enable CPU offloading (model-level and layerwise) for for DiT models
-6
View File
@@ -1,6 +0,0 @@
nav:
- CI_5levels.md
- failures.md
- test_guide.md
- test_markers.md
- test_style.md
+82 -74
View File
@@ -86,7 +86,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
/tests/e2e/online_serving/test_{model_name}_expansion.py<br>
/tests/e2e/offline_inference/test_{model_name}_expansion.py<br>
<strong>Performance:</strong><br>
/tests/dfx/perf/tests/test.json<br>
/tests/e2e/perf/nightly.json<br>
<strong>Doc Test:</strong><br>
tests/example/online_serving/test_{model_name}.py<br>
tests/example/offline_inference/test_{model_name}.py
@@ -104,7 +104,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
<td> Depends on reality </td>
<td>
<strong>Stability:</strong><br>
/tests/dfx/stability/tests/test.json<br>
tests/e2e/stability/weekly.json<br>
<strong>Reliability:</strong><br>
tests/e2e/reliability/test_{model_name}.py
</td>
@@ -225,11 +225,9 @@ vllm_omni/ tests/
│
└── e2e (tests side only) → ├── e2e/
├── online_serving/ ✅ non-empty
│ ├── test_qwen2_5_omni.py
│ ├── test_async_omni.py
│ ├── test_qwen3_omni.py
│ ├── test_qwen3_omni_expansion.py
│ ├── test_mimo_audio.py
│ ├── test_image_gen_edit.py
│ ├── test_images_generations_lora.py
│ └── stage_configs/
@@ -531,13 +529,13 @@ L4 level testing is a comprehensive quality audit before a version release. It e
### 3.2 Testing Content and Scope
- ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling.
- ***Performance Testing***: Uses the `tests/dfx/perf/tests/test.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
- ***Performance Testing***: Uses the `/tests/e2e/perf/nightly.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
- ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description.
### 3.3 Test Directory and Execution Files
- ***Functional Testing***: Same directories as L3.
- ***Performance Test Configuration***: `tests/dfx/perf/tests/test.json`
- ***Performance Test Configuration***: `/tests/e2e/perf/nightly.json`
- ***Documentation Example Tests***:
- - `tests/example/online_serving/test_{model_name}.py`
- `tests/example/offline_inference/test_{model_name}.py`
@@ -547,52 +545,13 @@ L4 level testing is a comprehensive quality audit before a version release. It e
- ***Trigger Timing***: **`Nightly`**, automatically executed every night.
- ***Execution Environment***: ***GPU*** server clusters to meet the resource demands of performance testing.
- ***Script Example***:
??? example "Test Examples: Documentation Example Tests"
--8<-- "docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md"
??? example "Test Examples: Performance Tests"
--8<-- "docs/contributing/ci/test_examples/l4_performance_tests.inc.md"
??? example "Test Examples: Functionality Tests"
--8<-- "docs/contributing/ci/test_examples/l4_functionality_tests.inc.md"
- ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).
## Chapter 4: L5 Level Testing - Stability and Reliability Testing
### 4.1 Testing Purpose
L5 level testing focuses on the performance of model services under ***long-running*** and ***abnormal fault*** scenarios. It aims to uncover deep-seated issues that only manifest under sustained pressure or extreme conditions, such as memory leaks, resource contention, gradual performance degradation, and lack of fault tolerance mechanisms. This is the final, yet crucial, line of defense for ensuring service high availability and production environment robustness.
### 4.2 Testing Content and Scope
- ***Long-term Stability (Stability) Testing***: Uses the `tests/dfx/stability/tests/test.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
- ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
### 4.3 Test Directory and Execution Files
- ***Stability Test Configuration***: `tests/dfx/stability/tests/test.json`
- ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
### 4.4 Execution Method and Example
- ***Trigger Timing***: **`Weekly`** (weekly) or **`Days before Release`** (several days before a major release). Due to long execution times, the frequency is lower.
- ***Execution Environment***: ***GPU*** servers, requiring a stable and exclusive testing environment.
- ***Script Example***:
<details>
<summary> Test Examples</summary>
When you want to add L4-level performance test cases, you can refer to the following format for case addition in tests/perf/tests/test.json:
When you want to add L5-level stability test cases, you can refer to the following format for case addition in `tests/dfx/stability/tests/test.json`:
```json
```JSON
{
"test_name": "test_qwen3_omni_stability",
"test_name": "test_qwen3_omni",
"server_params": {
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"stage_config_name": "qwen3_omni.yaml"
@@ -600,15 +559,17 @@ When you want to add L5-level stability test cases, you can refer to the followi
"benchmark_params": [
{
"dataset_name": "random",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"duration_sec": 43200,
"request_rate": 0.5,
"num_prompts_per_batch": 20,
"num_prompts": [10, 20],
"request_rate": [0.5, 1],
"random_input_len": 2500,
"random_output_len": 900,
"ignore_eos": true,
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration"
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": 100000,
"mean_audio_ttfp_ms": 100000,
"mean_audio_rtf": 100000
}
}
]
}
@@ -618,20 +579,20 @@ When you want to add L5-level stability test cases, you can refer to the followi
***Overview***
| Field | Required | Description |
| ---------------- | -------- | --------------------------------------------------------------------------- |
| test_name | Yes | Unique identifier for the stability test case |
| server_params | Yes | Server-side configuration parameters (model, stage configuration, etc.) |
| benchmark_params | Yes | Stability benchmark running parameters (supports multiple configurations) |
| Field | Required | Description |
| ---------------- | -------- | --------------------------------------------------------------- |
| test_name | Yes | Unique identifier for the test case |
| server_params | Yes | Server-side configuration parameters |
| benchmark_params | Yes | Benchmark running parameters (supports multiple configurations) |
#### server_params Configuration
##### Basic Parameters
| Parameter | Required | Example | Description |
| ----------------- | -------- | ---------------------------------- | ----------------------------------- |
| model | Yes | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path |
| stage_config_name | Yes | "qwen3_omni.yaml" | Stage configuration file name |
| Parameter | Required | Example | Description |
| ----------------- | -------- | ---------------------------------- | ----------------------------- |
| model | Yes | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path |
| stage_config_name | Yes | "qwen3_omni.yaml" | Stage configuration file name |
##### Dynamic Configuration (update/delete)
@@ -643,23 +604,70 @@ Supports incremental modifications based on the basic configuration:
| delete | Delete specified configuration items |
***Example***:
You can refer to Test Examples in Chapter 3.4
```
"update": {
"async_chunk": true, // Enable asynchronous chunk processing
"stage_args": {
"0": {
"engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
}
}
},
"delete": {
"stage_args": {
"2": ["custom_process_input_func"] // Delete this configuration for stage 2
}
}
```
#### benchmark_params Configuration
For stability testing, the key parameters are:
You can add any benchmark running parameters you need here. For all optional parameters, refer to the [benchmark documentation](https://github.com/vllm-project/vllm-omni/blob/main/docs/cli/bench/serve.md). General modifications are as follows:
- **duration_sec**: Total duration (in seconds) during which benchmark traffic is sent. The stability benchmark will keep sending batches until this duration is reached.
- **request_rate** / **max_concurrency**: Exactly one of them must be specified. They control how the traffic is generated for each batch:
- `request_rate`: Number of requests per second. The benchmark will send `num_prompts_per_batch` requests at the given rate.
- `max_concurrency`: Maximum number of concurrent requests. When this is used, `request_rate` is set to `inf` internally.
- **num_prompts_per_batch**: Number of prompts sent in each batch. Multiple batches will be executed sequentially within `duration_sec`.
All other optional parameters follow the same rules as the in Chapter 3.4.
1. Change the ---xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
2. For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
3. Add the baseline parameter to specify the required validation values, ensuring the validation metric names match those in the result.json generated by the benchmark.
4. The qps and concurrency modes are mutually exclusive. For detailed explanations, see the table below:
| Parameter | Type | Required | Example/Values | Description |
| --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| num_prompts | int / array | Yes | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
| request_rate | int / array | No | 1, [1, 2, 3] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. |
| max_concurrency | int / array | No | 1, [1, 2, 3] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. |
</details>
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_{model_name}.py`
- - ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).
## Chapter 4: L5 Level Testing - Stability and Reliability Testing
### 4.1 Testing Purpose
L5 level testing focuses on the performance of model services under ***long-running*** and ***abnormal fault*** scenarios. It aims to uncover deep-seated issues that only manifest under sustained pressure or extreme conditions, such as memory leaks, resource contention, gradual performance degradation, and lack of fault tolerance mechanisms. This is the final, yet crucial, line of defense for ensuring service high availability and production environment robustness.
### 4.2 Testing Content and Scope
- ***Long-term Stability (Stability) Testing***: Uses the `tests/e2e/stability/weekly.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
- ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
### 4.3 Test Directory and Execution Files
- ***Stability Test Configuration***: `tests/e2e/stability/weekly.json`
- ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
### 4.4 Execution Method and Example
- ***Trigger Timing***: **`Weekly`** (weekly) or **`Days before Release`** (several days before a major release). Due to long execution times, the frequency is lower.
- ***Execution Environment***: ***GPU*** servers, requiring a stable and exclusive testing environment.
- ***Script Example***:
<details>
<summary> Test Examples</summary>
```python
# WIP
```
</details>
- - ***Stability***: (Execution would be driven by the configuration in `weekly.json`).
- ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py`
## Summary
+4 -100
View File
@@ -1,106 +1,10 @@
# CI Failures
## Overview of CI Checks
What should I do when a CI job fails on my PR, but I don't think my PR caused the failure?
When you open a PR against vLLM-Omni, several CI checks run automatically:
| Check | Platform | What it does |
| ----- | -------- | ------------ |
| **pre-commit** | GitHub Actions | Runs linting (Ruff), formatting, spell-checking (typos), and YAML validation. |
| **Build Wheel** | GitHub Actions | Builds Python wheels for Python 3.11 and 3.12 on Ubuntu. Skipped for docs-only or Markdown-only changes (controlled by `paths-ignore` in the workflow). |
| **DCO** | GitHub | Verifies every commit has a `Signed-off-by` line. |
| **docs/readthedocs.org:vllm-omni** | Read the Docs | Builds the MkDocs documentation site. |
| **buildkite/vllm-omni** | Buildkite | Runs GPU-based tests on NVIDIA CUDA hardware (L4, H100). |
| **buildkite/vllm-omni-amd** | Buildkite | Runs GPU-based tests on AMD ROCm hardware (MI325). |
| **buildkite/vllm-omni-intel** | Buildkite | Runs GPU-based tests on Intel XPU hardware (Intel Arc BMG). |
## Common Case 1:
## Step 1: Identify the Failing Check
🚨 Error: The command was interrupted by a signal: signal: terminated
Click the **Details** link next to the failing check on your PR to open the build log. The most common failures fall into these categories:
### pre-commit failures
These are typically formatting or linting issues introduced by your PR. Fix them locally:
```bash
uv pip install pre-commit
pre-commit run --all-files
```
Then commit the fixes and push.
### DCO failures
Every commit must include a `Signed-off-by` line. If you forgot, amend your commits:
```bash
git commit --amend -s
git push --force-with-lease
```
For multiple commits, use an interactive rebase to add the sign-off to each one.
### Read the Docs failures
The documentation build uses MkDocs with `fail_on_warning: true`, so even a minor warning (not just errors) will cause the build to fail. To reproduce locally:
```bash
uv pip install -e ".[docs]"
mkdocs build --strict
```
Common causes include broken cross-references, invalid admonition syntax, or missing files referenced by `--8<--` includes.
### Buildkite failures
Buildkite runs GPU tests in Docker containers. These are the most complex checks and can fail for reasons unrelated to your PR (infrastructure issues, flaky tests, etc.). See the sections below for how to investigate.
## Step 2: Check if the Failure Is a Known Issue
Before spending time debugging, check whether the failure already exists on the `main` branch:
1. **Look at the Buildkite build log** — the test name and error message are usually enough to identify the issue.
2. **Check recent CI runs on `main`** — if the same test is failing there, the failure is not caused by your PR.
3. **Search existing issues** — look for open issues in the [vllm-omni issue tracker](https://github.com/vllm-project/vllm-omni/issues) with the test name or error message.
If the failure is already tracked, leave a comment on your PR noting that the failure is pre-existing and link the issue.
## Step 3: Investigate the Failure
If the failure appears to be new, investigate whether your changes caused it.
### Reading Buildkite Logs
1. Click **Details** next to the Buildkite check on your PR.
2. Find the failing step in the pipeline (e.g., "Diffusion Model Test", "Simple Unit Test").
3. Expand the step to see the full test output with the traceback.
### Running Tests Locally
For instructions on running tests locally (including specific test files, functions, and markers), see the [Running Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_guide/#running-tests) section in the Test Guide.
## Step 4: Raise an Issue or Fix It
### If the failure is pre-existing (not caused by your PR)
1. **Raise a new issue** if one doesn't already exist, using the title format:
`[CI Failure]: [job-name] - [test-path]`
2. Include the error message, relevant log excerpts, and the commit hash where the failure occurs (e.g., "Still failing on main as of commit `abc1234`").
3. Leave a comment on your PR linking to the issue and noting that the failure is unrelated to your changes.
### If the failure is caused by your PR
1. Fix the issue in your branch and push the update.
2. If the fix is non-trivial, consider adding a test to prevent regression.
## Common Failure Patterns
| Symptom | Likely Cause | Fix |
| ------- | ------------ | --- |
| `ruff` or formatting errors | Code style violation | Run `pre-commit run --all-files` |
| `Signed-off-by` missing | DCO check | Amend commits with `git commit --amend -s` |
| MkDocs build warning | Broken docs reference | Run `mkdocs build --strict` locally |
| `OOM` or `CUDA out of memory` | Test exceeds GPU memory | Check if your changes increased memory usage; use `--vae_use_slicing` / `--vae_use_tiling` for diffusion tests |
| Import errors | Missing or changed dependency | Check `pyproject.toml` and make sure dependencies are correct |
| Timeout (step exceeded N minutes) | Test is too slow or hangs | Profile the test; check for infinite loops or deadlocks |
| `Agent lost` in Buildkite | Infrastructure issue (not your fault) | Re-trigger the build; comment on your PR |
Reason: The test is terminated due to exceed the time limits. Sometimes rebuild the test will pass.
@@ -1,49 +0,0 @@
**Preferred Test Strategy**
Use one of the following patterns depending on page type:
- **Dynamic code-block extraction (preferred for offline docs)**
- Extract Python/Bash code blocks from markdown AST analyzer, then execute them directly in tests.
- Benefit: test logic stays automatically aligned with docs.
- Basic idea: Use `ReadmeSnippet.extract_readme_snippets` to extract a list of code blocks as a global variable in file,
use this list as `pytest.mark.parametrize` parameters, and pass each snippet item to `example_runner.run` inside the parametrized test.
Additionally pass an `output_subfolder` argument for the 2nd-level output folder explained in **Output Directory Structure** below.
If any extra environment variable is need for a test (e.g., the example script reads it), `example_runner.run` also accepts a 3rd `env` parameter.
- See [tests/examples/offline_inference/test_text_to_image.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/examples/offline_inference/test_text_to_image.py) for reference implementation.
- **Explicit copied scripts (used by online docs for now until further update)**
- For online serving pages, it is acceptable to copy code from docs into dedicated test functions, because only client-side, request-sending scripts are tested.
- Benefit: dynamic extraction is overly complex: need to tell server-launch and client-request scripts.
- Requirement: copied test code must be kept in sync with doc updates.
**Test Case Naming Convention**
- Dynamic code extraction (auto-generated internally):
- `test_{single_function_name_matching_file_name}[h2_heading_00X]`
- Example: `test_text_to_image[basic_usage_001]`
- Explicit copied scripts:
- `test_{h2_heading_00X}[{dummy_param_id_for_omni_server}]`
- Example: `test_api_calls_001[omni_server0]`
**Runtime Configuration**
In the example code tests, do **not** reduce `num_inference_steps` just to speed up the tests unless there is a strong CI reliability reason to do otherwise.
**Skipping Rules**
You may skip examples falling in the following categories using `pytest.mark.skip` or `pytest.skip`:
- Gradio UI scripts
- Scenarios that significantly overlap with existing tests and add little new coverage.
**Output Directory Structure**
Use a three-layer output structure to store output artifacts:
1. Root output directory
- Auto-detected from `OUTPUT_DIR` env var or auto-generated under `/tmp`.
2. Doc-page directory
- Define and use a clear page-level folder name in each `test_*.py` yourself (abbreviations are acceptable, e.g., `example_offline_t2i`).
3. Test-case directory
- Must match the case identifier (e.g., `basic_usage_001`).
- Auto-generated for dynamic extracted tests.
@@ -1,46 +0,0 @@
**Scope**
For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
**Test Case Design**
For a *high priority* model (currently listed in [issue #1832](https://github.com/vllm-project/vllm-omni/issues/1832)), we use several test cases, each with multiple features turned on, so that each supported feature is tested in at least one test case. This is to relieve the GPU workload on the CI machine. The suggested test case combination is as follows:
- If the model can fit into 4 L4 GPU (with quantization and tensor parallel always on) (20GB RAM each)
- (1 GPU) TeaCache + Layerwise CPU offloading + GGUF
- (4 GPUs) CacheDiT + Ulysses=2 + TP=2 + VAE=2 + FP8
- (4 GPUs) CacheDiT + Ring=2 + HSDP=2 + VAE=2 + GGUF
- (4 GPUs) TeaCache + CFG=2 + TP=2 + VAE=2 + FP8
- Otherwise, consider 2 H100 GPU environment (80GB RAM each) with the following tests
- (1 GPU) TeaCache + Layerwise CPU offloading + GGUF
- (2 GPUs) CacheDiT + Ulysses=2 + FP8
- (2 GPUs) CacheDiT + Ring=2 + GGUF
- (2 GPUs) TeaCache + CFG=2 + FP8
- (2 GPUs) CacheDiT + TP=2 + VAE=2 + FP8
- (2 GPUs) CacheDiT + HSDP=2 + VAE=2 + GGUF
- If 2 H100 GPU cannot handle the model either (e.g., HunyuanImage 3.0)
- Still design tests and feature combinations that can best fit real-world scenario.
- Do not include it in CI (or exclude it from the CI's filtering criteria). Instead, relevant PR authors are suggested to run these tests locally.
- Fallback plan
- If the model does not support layerwise CPU offloading, replace the corresponding test case with module-wise offloading
- If the model only supports specific or no caching feature, use this option or remove this option in all test cases.
- If the model only supports specific or no quantization feature, use this option or remove this option in all test cases.
- If the model does not support certain other features, remove this option from that test case. If, consequently, the coverage of this modified test case completely overlaps with others, remove this test case.
For a *normal priority* model, further reduce the number of test cases.
- Only write one or two test cases for the most common feature combinations.
- The author can explore themselves to see which feature combination balances output quality and performance. Alternatively, the author can refer to any example code in the PR that adds the model, or the example code in the PR that adds a feature (if the code involves this model of interest).
Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`.
**Code Style**
- Validation: test that the multimodal output files of your model have the correct shapes. `OpenAIClientHandler.send_diffusion_request` should have taken care of this.
- Test marks: always add `advanced_model` and `diffusion`. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/).
- To maximize code reuse, you may refer to
- `tests/conftest.py` for `omni_server` (running server in subprocess) and `openai_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper.
- `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
- [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
- Doc: add a concise docstring for each test function.
- Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).
@@ -1,89 +0,0 @@
When you want to add L4-level ***performance test*** cases, you can refer to the following format for case addition in tests/dfx/perf/tests/test.json:
```JSON
{
"test_name": "test_qwen3_omni",
"server_params": {
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"stage_config_name": "qwen3_omni.yaml"
},
"benchmark_params": [
{
"dataset_name": "random",
"num_prompts": [10, 20],
"max_concurrency": [1, 4],
"random_input_len": 2500,
"random_output_len": 900,
"ignore_eos": true,
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": [500, 800],
"mean_audio_ttfp_ms": [2000, 3500],
"mean_audio_rtf": [0.25, 0.35]
}
}
]
}
```
**Parameter Explanation**
*Overview*
| Field | Required | Description |
| ---------------- | -------- | --------------------------------------------------------------- |
| test_name | Yes | Unique identifier for the test case |
| server_params | Yes | Server-side configuration parameters |
| benchmark_params | Yes | Benchmark running parameters (supports multiple configurations) |
**`server_params` Configuration**
*Basic Parameters*
| Parameter | Required | Example | Description |
| ----------------- | -------- | ---------------------------------- | ----------------------------- |
| model | Yes | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path |
| stage_config_name | Yes | "qwen3_omni.yaml" | Stage configuration file name |
*Dynamic Configuration (update/delete)*
Supports incremental modifications based on the basic configuration:
| Operation | Description |
| --------- | ------------------------------------ |
| update | Update or add configuration items |
| delete | Delete specified configuration items |
**Example**:
```
"update": {
"async_chunk": true, // Enable asynchronous chunk processing
"stage_args": {
"0": {
"engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
}
}
},
"delete": {
"stage_args": {
"2": ["custom_process_input_func"] // Delete this configuration for stage 2
}
}
```
**`benchmark_params` Configuration**
You can add any benchmark running parameters you need here. For all optional parameters, refer to the [benchmark documentation](https://github.com/vllm-project/vllm-omni/blob/main/docs/cli/bench/serve.md). General modifications are as follows:
1. Change the --xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
2. For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
3. Optionally add a `baseline` object (see **Baseline thresholds** below). If you omit `baseline` or leave it empty, the performance test still runs but does not assert metric thresholds from this field.
4. The qps and concurrency modes are recommended to be mutually exclusive. For detailed explanations, see the table below:
| Parameter | Type | Required | Example/Values | Description |
| --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| num_prompts | int / array | Yes | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
| request_rate | float / array | No | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. |
| max_concurrency | int / array | No | 1, [1, 2, 3] | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode). |
| baseline | object | No | see above | Optional per-metric thresholds; keys must match benchmark output fields. Scalar, list (per sweep step), or object (keyed by concurrency or QPS string).
-8
View File
@@ -130,11 +130,9 @@ vllm_omni/ tests/
│
└── e2e (tests side only) → ├── e2e/
├── online_serving/ ✅ non-empty
│ ├── test_qwen2_5_omni.py
│ ├── test_async_omni.py
│ ├── test_qwen3_omni.py
│ ├── test_qwen3_omni_expansion.py
│ ├── test_mimo_audio.py
│ ├── test_image_gen_edit.py
│ ├── test_images_generations_lora.py
│ └── stage_configs/
@@ -159,12 +157,6 @@ vllm_omni/ tests/
├── qwen3_omni_ci.yaml
├── bagel_*.yaml
└── npu/, rocm/, etc.
examples/ tests
│ └── examples
├── online_serving/ → ├── online_serving/
│ └── {doc_page_title}/README.md │ └── test_{doc_page_title}.py ⬜
└── offline_inference/ → └── offline_inference/
└── {doc_page_title}/README.md └── test_{doc_page_title}.py ⬜
```
@@ -653,7 +653,25 @@ For a fair comparison, keep the same **prompt**, **seed**, **resolution**, **num
To ensure project maintainability and sustainable development, please submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
For comprehensive testing guidelines and the definition of test levels (L1-L5), please refer to the [Multi-Level Automated Testing System Documentation](../ci/CI_5levels.md). You are at least required to add an L4 *functionality* test described in that document.
For comprehensive testing guidelines and the definition of test levels (L1-L5), please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
The following tests are required to add:
- L4 test of the model's full *functionality* (i.e., all the *diffusion features* that are supported by this model), including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
- Test cases: Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`. The following test cases shall cover all features:
- 1 GPU: TeaCache & GGUF (or fallback to FP8, or disable it) & Layer-wise CPU offloading (or fallback to Module-wise)
- 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Ulysses = 2
- 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & Ring = 2
- 2 GPUs: TeaCache & FP8 (or fallback to GGUF, or disable it) & CFG Parallel = 2
- 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Tensor Parallel = 2 & VAE Patch Parallel = 2
- 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & HSDP = 2 & VAE Patch Parallel = 2
- Validation: test that the multimodal output files of your model have the correct shapes.
- Test marks: always add `advanced_model` and `diffusion`. Add `parallel` and GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/)
- To maximize code reuse, you may refer to
- `tests/conftest.py` for `omni_server` and `openai_client` fixtures, `generate_synthetic_image` and `assert_XXX_valid` helper.
- `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
- [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
- Doc: add a concise dostring for each test function.
- Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).
---
@@ -1,478 +0,0 @@
# Support FP8 Quantization
This section describes how to add FP8 quantization support to a diffusion model. We use the Qwen-Image pipeline as the reference implementation.
---
## Table of Contents
- [Overview](#overview)
- [Step-by-Step Implementation](#step-by-step-implementation)
- [Testing](#testing)
- [Troubleshooting](#troubleshooting)
- [Reference Implementations](#reference-implementations)
- [Summary](#summary)
---
## Overview
### What is FP8 Quantization?
FP8 quantization reduces model memory footprint by storing weights in 8-bit floating point (`float8_e4m3fn`) instead of 16-bit (`bfloat16`/`float16`). This saves ~50% of weight memory.
A diffusion pipeline typically has three quantizable components:
| Component | Layer Types | Quantization Approach |
|-----------|-----------|----------------------|
| **DiT (transformer)** | `nn.Linear` | Replace with vLLM quantized linear layers |
| **Text encoder** | `nn.Linear` | FP8 weight storage via hooks |
| **VAE** | `nn.Conv2d`, `nn.Conv3d` | FP8 weight storage via hooks |
### Architecture
The quantization system has two mechanisms:
**1. vLLM Quantized Linear Layers (for DiT)**
vLLM provides drop-in replacements for `nn.Linear` that perform FP8 W8A8 computation on supported hardware. These are configured via `quant_config` passed to layer constructors.
Key classes:
| Class | Location | Purpose |
|-------|----------|---------|
| `build_quant_config` | `vllm_omni/quantization/__init__.py` | Build quantization config from string/dict/per-component spec |
| `ComponentQuantizationConfig` | `vllm_omni/quantization/component_config.py` | Per-component routing by layer prefix |
**2. FP8 Weight Storage via Hooks (for text encoder and VAE)**
For models loaded via `from_pretrained()` (outside vLLM's weight pipeline), we use forward hooks to store weights in FP8 and dequantize on-the-fly:
```
At rest: weight.data = fp8_tensor (half memory)
Pre-hook: weight.data = fp8_weight.to(bf16) * scale
Forward: Normal computation in BF16
Post-hook: weight.data = fp8_weight (back to FP8)
```
Key function:
| Function | Location | Purpose |
|----------|----------|---------|
| `apply_fp8_weight_storage` | `vllm_omni/diffusion/models/utils.py` | Apply FP8 hooks to Linear/Conv layers (introduced in [#1412](https://github.com/vllm-project/vllm-omni/pull/1412)) |
---
## Step-by-Step Implementation
### Step 1: Add `quant_config` to the Transformer
Pass `quant_config` to vLLM parallel linear layers in the transformer. This enables FP8 W8A8 for DiT linear layers.
#### 1.1: Accept `quant_config` in Constructor
```python
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
class YourTransformer2DModel(nn.Module):
def __init__(
self,
*,
od_config: OmniDiffusionConfig,
quant_config: QuantizationConfig | None = None,
# ... other params
):
self.quant_config = quant_config
# Pass to sub-modules
self.blocks = nn.ModuleList([
YourTransformerBlock(quant_config=quant_config, ...)
for _ in range(num_layers)
])
```
#### 1.2: Pass `quant_config` to vLLM Linear Layers
Replace `nn.Linear` with vLLM parallel linear layers that accept `quant_config`:
```python
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
QKVParallelLinear,
RowParallelLinear,
)
class YourAttentionBlock(nn.Module):
def __init__(self, hidden_size, num_heads, quant_config=None, prefix=""):
# QKV projection
self.to_qkv = QKVParallelLinear(
hidden_size=hidden_size,
head_size=hidden_size // num_heads,
total_num_heads=num_heads,
quant_config=quant_config,
prefix=f"{prefix}.to_qkv",
)
# Output projection
self.to_out = RowParallelLinear(
input_size=hidden_size,
output_size=hidden_size,
quant_config=quant_config,
prefix=f"{prefix}.to_out",
)
class YourFeedForward(nn.Module):
def __init__(self, hidden_size, intermediate_size, quant_config=None, prefix=""):
self.gate_up_proj = ColumnParallelLinear(
input_size=hidden_size,
output_size=intermediate_size * 2,
quant_config=quant_config,
prefix=f"{prefix}.gate_up_proj",
)
self.down_proj = RowParallelLinear(
input_size=intermediate_size,
output_size=hidden_size,
quant_config=quant_config,
prefix=f"{prefix}.down_proj",
)
```
!!! important
Always pass the `prefix` parameter. It is used by `ignored_layers` matching to skip sensitive layers.
#### 1.3: Identify Sensitive Layers
Some layers degrade output quality when quantized. Common sensitive layers:
| Layer | Why Sensitive | Typical Models | Quant Methods Affected |
|-------|-------------|---------------|----------------------|
| `img_mlp` | Processes denoising latents with shifting dynamic range | Qwen-Image | FP8, Int8 |
| `feed_forward` | FFN layers in DiT blocks, large dynamic range | Z-Image | Int8 |
| `proj_out` | Final output projection, small quantization errors amplified | Various | FP8, Int8 |
| `lm_head` | Output vocabulary projection, precision-critical | Omni (thinker) | NVFP4 |
| `mlp.gate` | MoE router gates, routing decisions are precision-critical | MoE models | NVFP4 |
**Quantization quality benchmarks** (LPIPS: lower = better, <0.01 imperceptible, >0.1 noticeable):
| Model | Method | All Layers | With `ignored_layers` | Recommendation |
|-------|--------|-----------|----------------------|----------------|
| Qwen-Image-2512 | Int8 | LPIPS 0.0197 | 0.0027 (skip `img_mlp`) | Skip `img_mlp` |
| Z-Image-Turbo | Int8 | LPIPS 0.1597 | 0.0290 (skip `feed_forward`) | Skip `feed_forward` |
| Z-Image-Turbo | FP8 | LPIPS ~0.005 | — | All layers OK |
To identify sensitive layers for a new model:
1. Run inference with all layers quantized
2. Compare output quality with BF16 baseline using LPIPS (see `benchmarks/diffusion/quantization_quality.py`)
3. Selectively disable quantization per layer (via `ignored_layers`) until quality matches
4. Document the recommended `ignored_layers` in the model's supported models table
### Step 2: Wire Up Quantization in the Pipeline
#### 2.1: Create `quant_config` and Pass to Transformer
In the pipeline's `__init__`, extract the vLLM quant config and pass it to the transformer:
```python
from vllm_omni.quantization import build_quant_config
class YourPipeline(nn.Module):
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__()
# ... load scheduler, text_encoder, vae ...
# Build quant config from OmniDiffusionConfig
quant_config = build_quant_config(od_config.quantization_config)
# Pass to transformer
self.transformer = YourTransformer2DModel(
od_config=od_config,
quant_config=quant_config,
# ... other kwargs
)
```
#### 2.2: Apply FP8 Weight Storage to Text Encoder and VAE
For models loaded via `from_pretrained()`, apply hook-based FP8 weight storage.
!!! note
The `apply_fp8_weight_storage` function is provided by `vllm_omni/diffusion/models/utils.py`, introduced in [PR #1412](https://github.com/vllm-project/vllm-omni/pull/1412).
```python
from vllm_omni.diffusion.models.utils import apply_fp8_weight_storage
class YourPipeline(nn.Module):
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__()
self.text_encoder = SomeTextEncoder.from_pretrained(...)
self.vae = SomeVAE.from_pretrained(...).to(self.device)
# Apply FP8 weight storage to text encoder and VAE
if (
od_config.quantization_config is not None
and getattr(od_config.quantization_config, "quant_method", None) == "fp8"
):
apply_fp8_weight_storage(self.vae)
apply_fp8_weight_storage(self.text_encoder)
# ... rest of init
```
`apply_fp8_weight_storage` walks all `nn.Linear`, `nn.Conv2d`, and `nn.Conv3d` modules in the model and applies FP8 hooks to each.
#### 2.3: Fix `load_weights` to Include Pre-loaded Components
Since text encoders and VAEs are loaded via `from_pretrained()` (not through vLLM's weight pipeline), their parameters won't appear in the `loaded_weights` set. Mark them explicitly:
```python
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self)
loaded_weights = loader.load_weights(weights)
# Mark pre-loaded component weights as loaded
loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
return loaded_weights
```
### Step 3: Add to Supported Models Documentation
Update the FP8 supported models table in `docs/user_guide/diffusion/quantization/fp8.md`:
```markdown
| Your-Model | `org/Your-Model` | ✅ | ✅ | ✅ | `sensitive_layer` |
```
---
## Testing
### Quick Verification
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# Test with FP8 enabled
omni = Omni(model="your-model", quantization="fp8")
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
### Quality Gate Test (LPIPS)
We provide a pytest-based quality gate at `tests/diffusion/quantization/test_quantization_quality.py`.
It generates outputs with both BF16 and your quantized method using the same seed, computes
[LPIPS](https://github.com/richzhang/PerceptualSimilarity) perceptual distance, and **fails if it
exceeds a threshold**. This is the recommended way to validate that a quantization method does not
unacceptably degrade output quality.
**Step 1: Add your test config** — edit `QUALITY_CONFIGS` in the test file:
```python
QualityTestConfig(
id="int8_z_image", # pytest ID
model="Tongyi-MAI/Z-Image-Turbo",
quantization="int8", # your quantization method
task="t2i", # "t2i" or "t2v"
prompt="a cup of coffee on a wooden table, morning light",
max_lpips=0.15, # fail threshold (see guide below)
num_inference_steps=20,
),
```
**Step 2: Run the test:**
```bash
# Install dependency
pip install lpips
# Run all quality tests
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m ""
# Run only your method
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "int8"
# Run a specific model
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image"
```
**Step 3: Read the output:**
```
Quantization Quality: int8_z_image
Model: Tongyi-MAI/Z-Image-Turbo
Method: int8
LPIPS: 0.0290 (threshold: 0.15)
BF16 memory: 19.15 GiB
Quant memory: 14.80 GiB (23% reduction)
Result: PASS
```
**LPIPS threshold guide:**
| LPIPS Range | Meaning | Recommendation |
|-------------|---------|----------------|
| < 0.01 | Imperceptible | Ideal — no visible difference |
| 0.01 – 0.05 | Minor differences | Good — acceptable for most use cases |
| 0.05 – 0.15 | Noticeable on close inspection | Acceptable with `ignored_layers` tuning |
| 0.15 – 0.30 | Clearly visible | Investigate — likely needs `ignored_layers` |
| > 0.30 | Significant degradation | Quantization method may not suit this model |
!!! tip
If your test fails, don't raise the threshold — instead, identify sensitive layers and add
them to `ignored_layers` (see next section). Then re-run the test with the lower LPIPS.
### Quality Benchmark (Detailed)
For a full multi-prompt benchmark with Markdown output (suitable for PR descriptions), use the
standalone benchmark script:
```bash
python benchmarks/diffusion/quantization_quality.py \
--model Tongyi-MAI/Z-Image-Turbo \
--quantization fp8 \
--prompts \
"a cup of coffee on the table" \
"an aerial view of a coral reef" \
--height 1024 --width 1024 \
--num-inference-steps 50 --seed 42
```
This generates BF16 and quantized outputs, computes per-prompt LPIPS, and prints a Markdown
summary table ready to paste into your PR.
### Identifying Sensitive Layers with `ignored_layers`
Not all layers tolerate quantization equally. **`ignored_layers`** lets you keep specific layers
in BF16 while quantizing everything else. This is critical for achieving good quality with
aggressive quantization methods (Int8, NVFP4).
**Why some layers are sensitive:**
| Layer Type | Why Sensitive | Typical Impact |
|------------|--------------|----------------|
| `img_mlp` | Processes denoising latents with shifting dynamic range across timesteps. No built-in normalization to absorb quantization error. | Color shifts, blurriness |
| `feed_forward` | FFN layers with large dynamic range in DiT blocks. | Artifacts, loss of detail |
| `proj_out` | Final output projection — small errors get amplified since there are no subsequent layers to correct them. | Overall quality loss |
| `lm_head` | Vocabulary projection in LLMs — precision-critical for token selection. | Garbage text output |
| `mlp.gate` | MoE router gates — routing decisions are binary and precision-critical. | Wrong expert selection, quality collapse |
**How to find sensitive layers for your model:**
1. Run the quality gate test with all layers quantized:
```bash
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "your_model"
```
2. If LPIPS exceeds your threshold, try skipping common sensitive layers one at a time:
```python
QualityTestConfig(
id="fp8_your_model_skip_mlp",
model="your-org/Your-Model",
quantization="fp8",
task="t2i",
prompt="a cup of coffee on a wooden table",
max_lpips=0.05,
),
```
And in the test, modify the Omni call:
```python
omni_qt = Omni(
model=config.model,
quantization_config={
"method": config.quantization,
"ignored_layers": ["img_mlp"],
},
)
```
3. Compare LPIPS with and without `ignored_layers` to isolate the problematic layers.
4. Document the recommended `ignored_layers` in the supported models table.
**Real-world examples:**
| Model | Method | All Layers (LPIPS) | With `ignored_layers` | Which Layers to Skip |
|-------|--------|-------------------|----------------------|---------------------|
| Qwen-Image-2512 | Int8 | 0.0197 | 0.0027 | `img_mlp` |
| Z-Image-Turbo | Int8 | 0.1597 | 0.0290 | `feed_forward` |
| Z-Image-Turbo | FP8 | ~0.005 | — | None needed |
### What to Verify (Checklist)
1. **Model loads without errors** — Check logs for quantization initialization messages
2. **Quality gate passes** — LPIPS within threshold vs BF16 baseline
3. **Memory reduction** — Expect ~30-50% savings on quantized components
4. **Sensitive layers identified** — Document which layers (if any) need `ignored_layers`
5. **`ignored_layers` pattern works** — Verify that the layer name prefix matches your transformer's module names
---
## Troubleshooting
### Issue: Output Quality Degradation
**Symptoms:** Generated images have visible artifacts, color shifts, or blurriness compared to BF16 baseline.
**Causes & Solutions:**
- **Sensitive layers being quantized:**
**Solution:** Identify and skip sensitive layers:
```python
omni = Omni(
model="your-model",
quantization_config={
"method": "fp8",
"ignored_layers": ["img_mlp"],
},
)
```
- **All layers sensitive (rare):**
**Solution:** The model may not be suitable for FP8 quantization. Consider weight-only quantization or other methods.
### Issue: Missing Weights Warning
**Symptoms:** Warnings about unloaded weights for `vae.*` or `text_encoder.*` parameters.
**Solution:** Ensure `load_weights` marks pre-loaded component weights. See [Step 2.3](#23-fix-load_weights-to-include-pre-loaded-components).
### Issue: CUDA Out of Memory
**Symptoms:** OOM errors during model loading despite FP8 being enabled.
**Causes & Solutions:**
- **Text encoder/VAE loaded in BF16 before FP8 conversion:**
FP8 weight storage is applied after `from_pretrained()`, so peak load memory is still BF16. This is a one-time cost during initialization.
---
## Reference Implementations
| Model | Pipeline File | Transformer File | Notes |
|-------|-------------|-----------------|-------|
| **Z-Image** | `models/z_image/pipeline_z_image.py` | `models/z_image/z_image_transformer.py` | DiT + text encoder FP8 |
| **Qwen-Image** | `models/qwen_image/pipeline_qwen_image.py` | `models/qwen_image/qwen_image_transformer.py` | DiT + text encoder + VAE FP8 |
All files are under `vllm_omni/diffusion/`.
---
## Summary
Adding FP8 quantization support to a new model:
1. ✅ **Add `quant_config` to transformer** — Accept and pass to vLLM parallel linear layers with proper `prefix`
2. ✅ **Identify sensitive layers** — Profile which layers need `ignored_layers` for quality
3. ✅ **Apply FP8 weight storage** — Call `apply_fp8_weight_storage()` on text encoder and VAE in pipeline `__init__`
4. ✅ **Fix `load_weights`** — Mark pre-loaded component weights as loaded
5. ✅ **Update docs** — Add model to the supported models table in `docs/user_guide/diffusion/quantization/fp8.md`
6. ✅ **Test** — Verify quality, memory savings, and sensitive layer behavior
+42 -110
View File
@@ -4,54 +4,48 @@
vLLM-Omni uses the PyTorch Profiler to analyze performance across both **multi-stage omni-modality models** and **diffusion models**.
### 1. Configure Profiling in the Stage YAML
### 1. Set the Output Directory
Before running any script, set this environment variable. The system detects this and automatically saves traces here.
Enable profiling by adding `profiler_config` under `engine_args` for the stage(s) you want to profile in your stage config YAML:
```yaml
stage_args:
- stage_id: 0
stage_type: llm
engine_args:
# ... other engine args ...
profiler_config:
profiler: torch
torch_profiler_dir: ./perf
```bash
export VLLM_TORCH_PROFILER_DIR=./profiles
```
| Field | Description |
|---|---|
| `profiler` | Profiler backend to use. Currently supports `torch`. |
| `torch_profiler_dir` | Directory where trace files are saved. Created automatically if it doesn't exist. |
> **Tip:** Only enable `profiler_config` on stages you actually need to profile. Stages without it will not start a profiler, keeping overhead minimal.
### 2. Profiling Omni-Modality Models
**Selective Stage Profiling**
It is best to limit profiling to one iteration to keep trace files manageable.
It is highly recommended to profile specific stages to prevent producing overly large trace files:
```python
# Profile all stages
omni_llm.start_profile()
# Only profile Stage 1
omni_llm.start_profile(stages=[1])
# Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
omni_llm.start_profile(stages=[0, 2])
```bash
export VLLM_PROFILER_MAX_ITERS=1
```
> **Important:** Always pass the same `stages` list to both `start_profile()` and `stop_profile()`. If you omit `stages` from `stop_profile()`, it defaults to stopping all stages — including ones that were never started — which will produce errors.
**Selective Stage Profiling**
The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
```python
# Profile all stages
omni.start_profile()
# Only profile Stage 1
omni.start_profile(stages=[1])
```
```python
# Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
omni.start_profile(stages=[0, 2])
```
**Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
```python
profiler_stages = [0] # Only profile the stages you need
from vllm_omni.entrypoints.omni import Omni
# 1. Start profiling
omni.start_profile(stages=profiler_stages)
omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
# 1. Start profiling if enabled
if profiler_enabled:
omni.start_profile(stages=[0])
# Initialize generator
omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
@@ -65,57 +59,42 @@ for stage_outputs in omni_generator:
# ... [Output processing logic for text/audio would go here] ...
# Update count to track when to stop profiling
processed_count += len(stage_outputs.request_output)
processed_count += 1
# 2. Check if all requests are done to stop the profiler safely
if profiler_enabled and processed_count >= total_requests:
print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
# Stop the profiler while workers are still active
# Pass the same stages list used in start_profile()
omni_llm.stop_profile(stages=profiler_stages)
omni.stop_profile()
# Wait for traces to flush to disk
print("[Info] Waiting 30s for workers to write trace files to disk...")
time.sleep(30)
print("[Info] Trace export wait time finished.")
omni_llm.close()
omni.close()
```
**CLI Usage** (using `end2end.py`):
```bash
# Profile only Stage 0 (Thinker)
python end2end.py --output-wav output_audio \
--query-type text --enable-profiler --profiler-stages 0
# Profile Stage 0 and Stage 2
python end2end.py --output-wav output_audio \
--query-type text --enable-profiler --profiler-stages 0 2
# Profile all stages (omit --profiler-stages)
python end2end.py --output-wav output_audio \
--query-type text --enable-profiler
```
**Examples**:
1. **Qwen2.5-Omni**: [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
2. **Qwen3-Omni**: [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
### 3. Profiling diffusion models
Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding. Standalone diffusion scripts use `--profiler-dir` to enable profiling.
Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding.
**CLI Usage:**
```bash
```python
python image_to_video.py \
--model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
--image qwen-bear.png \
--prompt "A cat playing with yarn, smooth motion" \
--profiler-dir \
\
# Minimize Spatial Dimensions (Optional but helpful):
# Drastically reduces memory usage so the profiler doesn't
@@ -145,9 +124,8 @@ python image_to_video.py \
--flow-shift 12.0 \
--fps 16 \
--output i2v_output.mp4
```
> **Note:** For diffusion stages within a multi-stage omni pipeline, use `profiler_config` in the stage YAML instead (see Section 1).
```
**Examples**:
@@ -155,62 +133,16 @@ python image_to_video.py \
2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**: [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
### 4. Profiling Online Serving
### 4. Analyzing Omni Traces
When `profiler_config` is set in the stage YAML, the server automatically exposes `/start_profile` and `/stop_profile` HTTP endpoints.
**1. Start the server** with a stage YAML that has `profiler_config` enabled:
```bash
vllm serve Qwen/Qwen2.5-Omni-7B \
--omni \
--stage-configs-path qwen2_5_omni.yaml \
--port 8091
```
Or for one stage diffusion models:
```bash
vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --port 8091 --profiler-config '{"profiler": "torch", "torch_profiler_dir": "./vllm_profile"}'
```
**2. Start profiling** by sending a POST request:
```bash
# Profile all stages that have profiler_config set
curl -X POST http://localhost:8091/start_profile
# Profile specific stages only
curl -X POST http://localhost:8091/start_profile \
-H "Content-Type: application/json" \
-d '{"stages": [0]}'
```
**3. Send your inference requests** as normal while the profiler is running.
**4. Stop profiling** and collect traces:
```bash
# Stop all stages
curl -X POST http://localhost:8091/stop_profile
# Stop specific stages (must match the stages you started)
curl -X POST http://localhost:8091/stop_profile \
-H "Content-Type: application/json" \
-d '{"stages": [0]}'
```
Trace files are written to the `torch_profiler_dir` specified in your stage YAML.
> **Important:** Always stop the same stages you started. Stopping a stage that was never started will produce errors.
### 5. Analyzing Traces
Output files are saved to the `torch_profiler_dir` specified in your stage YAML config.
Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
**Output**
**Chrome Trace** (`.json.gz`): Visual timeline of kernels and stages. Open in Perfetto UI.
**Chrome Trace** (```.json.gz```): Visual timeline of kernels and stages. Open in Perfetto UI.
**Viewing Tools:**
- [Perfetto](https://ui.perfetto.dev/) (recommended)
- `chrome://tracing` (Chrome only)
- [Perfetto](https://ui.perfetto.dev/)(recommended)
- ```chrome://tracing```(Chrome only)
**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation: [vLLM Profiling Guide](https://docs.vllm.ai/en/stable/contributing/profiling/)
+5 -5
View File
@@ -113,13 +113,13 @@ def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], No
Some models require custom handling:
- **Single or dual-transformer:** Models that may use one or two transformers (e.g., Wan2.2)
- **Dual-transformer:** Models with separate high-noise and low-noise transformers (e.g., Wan2.2)
- **Multi-block-list:** Models with multiple block lists in one transformer (e.g., LongCatImage with `transformer_blocks` + `single_transformer_blocks`)
- **Special forward patterns:** Models with non-standard block execution patterns
### Example 1: Single or Dual-Transformer Model (Wan2.2)
### Example 1: Dual-Transformer Model (Wan2.2)
Wan2.2 can use either a single transformer or two transformers (one for high-noise steps and one for low-noise steps). The implementation automatically detects the mode based on the presence of `transformer_2`.
Wan2.2 uses two transformers: one for high-noise steps and one for low-noise steps.
**Key difference:** Use `BlockAdapter` to wrap multiple transformers with separate configurations.
@@ -256,7 +256,7 @@ cache_config={
}
```
Check the [user guide for cache_dit](../../user_guide/diffusion/cache_acceleration/cache_dit.md) for more adjustable parameters.
Check the [user guide for cache_dit](../../user_guide/diffusion/cache_dit_acceleration.md) for more adjustable parameters.
---
@@ -267,7 +267,7 @@ Complete examples in the codebase:
| Model | Path | Pattern | Notes |
|-------|------|---------|-------|
| **Standard DiT** | [`cache_dit_backend.py::enable_cache_for_dit`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Default enabler | Single transformer, automatic |
| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Single or dual-transformer | Auto-detects mode based on transformer_2 presence |
| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Dual-transformer | Separate high/low noise transformers |
| **LongCat** | [`cache_dit_backend.py::enable_cache_for_longcat_image`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_longcat_image) | Multi-block-list | Two block lists in one transformer |
| **BAGEL** | [`cache_dit_backend.py::enable_cache_for_bagel`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_bagel) | Omni model | Complex architecture |
+10 -67
View File
@@ -34,7 +34,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
| Method | Purpose | Automatic Behavior |
|--------|---------|-------------------|
| [`predict_noise_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with CFG | Detects parallel mode, distributes computation, gathers results |
| [`scheduler_step_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Step scheduler | All ranks step locally (no broadcast needed) |
| [`scheduler_step_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Step scheduler with sync | Rank 0 steps, broadcasts latents to all ranks |
| [`combine_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine positive/negative | Applies CFG formula with optional normalization |
| [`predict_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Forward pass wrapper | Override for custom transformer calls |
| [`cfg_normalize_function()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Normalize CFG output | Override for custom normalization |
@@ -47,7 +47,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
- Rank 0 computes positive prompt prediction
- Rank 1 computes negative prompt prediction
- Results are gathered via `all_gather()`
- All ranks compute CFG combine locally (deterministic, identical results)
- Combined on rank 0 using CFG formula
- **Sequential mode** (when `cfg_world_size == 1`):
- Single rank computes both positive and negative predictions
@@ -55,7 +55,14 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
`scheduler_step_maybe_with_cfg()` ensures consistent latent states across all ranks:
- All ranks compute the scheduler step locally — no broadcast needed because `predict_noise_maybe_with_cfg` already ensures all ranks have identical noise predictions after `all_gather` + local combine.
- **CFG-Parallel mode**:
- Only rank 0 performs the scheduler step (applies noise prediction to update latents)
- Updated latents are broadcast to all other ranks via `broadcast()`
- All ranks maintain synchronized latent states for the next iteration
- **Sequential mode**:
- Single rank directly performs the scheduler step
- No synchronization needed
---
@@ -170,70 +177,6 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin):
# return noise_pred
```
### Override `combine_cfg_noise()` for Multi-Output Models
When `predict_noise()` returns a tuple (e.g., video + audio), the default `combine_cfg_noise()` applies CFG to every element. Override it to apply different logic per element — for example, CFG on video but positive-only on audio:
```python
class MyVideoAudioPipeline(nn.Module, CFGParallelMixin):
def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, scale, normalize):
(video_pos, audio_pos) = positive_noise_pred
(video_neg, audio_neg) = negative_noise_pred
video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize)
return (video_combined, audio_pos) # audio: positive only, no CFG
```
This also requires `predict_noise()` to return a tuple (see [Override predict_noise](#override-predict_noise-for-custom-transformer-calls) above).
### Implement a Composite Scheduler for Multi-Output Models
When each output has its own denoising schedule, implement a composite scheduler that dispatches to per-output schedulers. Assign it to `self.scheduler` so the default `scheduler_step()` works without override.
**Complete example (video + audio with separate schedulers and diffuse loop):**
```python
class VideoAudioScheduler:
"""Composite scheduler dispatching to video and audio schedulers."""
def __init__(self, video_scheduler, audio_scheduler):
self.video_scheduler = video_scheduler
self.audio_scheduler = audio_scheduler
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
video_out = self.video_scheduler.step(noise_pred[0], t[0], latents[0], return_dict=False, generator=generator)[0]
audio_out = self.audio_scheduler.step(noise_pred[1], t[1], latents[1], return_dict=False, generator=generator)[0]
return ((video_out, audio_out),)
class MyVideoAudioPipeline(nn.Module, CFGParallelMixin):
def __init__(self, ...):
self.scheduler = VideoAudioScheduler(video_sched, audio_sched)
def predict_noise(self, **kwargs):
video_pred, audio_pred = self.transformer(**kwargs)
return (video_pred, audio_pred)
def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, scale, normalize):
# ... (as above)
def diffuse(self, video_latents, audio_latents, timesteps_video, timesteps_audio, ...):
for t_v, t_a in zip(timesteps_video, timesteps_audio):
positive_kwargs = {...}
negative_kwargs = {...} if do_true_cfg else None
video_pred, audio_pred = self.predict_noise_maybe_with_cfg(
do_true_cfg=do_true_cfg, true_cfg_scale=self.guidance_scale,
positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs,
)
video_latents, audio_latents = self.scheduler_step_maybe_with_cfg(
(video_pred, audio_pred), (t_v, t_a),
(video_latents, audio_latents), do_true_cfg=do_true_cfg,
generator=generator,
)
return video_latents, audio_latents
```
> **Note:** If you use a non-deterministic scheduler, e.g., DDPM, please set `self.scheduler_step_maybe_with_cfg(..., generator=torch.Generator(device).manual_seed(seed))` explicitly to control the randomness of scheduler step among ranks.
---
## Testing
+2 -43
View File
@@ -7,7 +7,6 @@ This section describes how to add Sequence Parallel (SP) to a diffusion transfor
## Table of Contents
- [Overview](#overview)
- [UAA Mode (Experimental)](#uaa-mode-experimental)
- [Approach 1: Non-Intrusive `_sp_plan` (Recommended)](#approach-1-non-intrusive-_sp_plan-recommended)
- [Approach 2: Intrusive Modification (For Complex Cases)](#approach-2-intrusive-modification-for-complex-cases)
- [Testing](#testing)
@@ -47,41 +46,6 @@ from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
---
## UAA Mode (Experimental)
`ulysses_mode="advanced_uaa"` enables the experimental UAA ("Ulysses Anything Attention") feature, which lets Ulysses attention handle arbitrary sequence lengths and arbitrary attention head counts. The same idea is also supported by [Cache-DiT](https://cache-dit.readthedocs.io/en/latest/user_guide/CONTEXT_PARALLEL/#uaa-ulysses-anything-attention).
Use it when plain Ulysses-SP would otherwise fail because:
- the local sequence shards are not evenly divisible after split hooks, or
- the attention head count is not divisible by `ulysses_degree`.
### Design Summary
1. **Strict mode stays unchanged.**
`ulysses_mode="strict"` keeps the original fast path and still requires divisible sequence/head shapes.
2. **UAA uses variable all-to-all split sizes for sequence shards.**
Before the Ulysses Q/K/V exchange, each rank all-gathers its local sequence length and uses those lengths as `all_to_all_single(..., output_split_sizes=seq_lens)`. This lets Ulysses gather the full sequence even when each rank started with a different local shard length.
3. **UAA pads heads only inside the Ulysses exchange.**
If `head_cnt % ulysses_degree != 0`, UAA pads the head dimension up to the next multiple of `ulysses_degree`, performs the forward/reverse all-to-all, then slices the temporary head padding away after the reverse exchange. The same rule is applied to joint attention tensors.
4. **Hybrid Ulysses + Ring is still shape-constrained.**
Ring attention expects every rank in a ring group to exchange exactly the same post-Ulysses sequence shape. UAA therefore validates those shapes before entering the ring path and raises a clear error if ring peers disagree on `S_global`.
5. **Tiny scalar gathers stay out of TorchDynamo tracing.**
`_all_gather_int()` is marked with `@torch.compiler.disable` so the scalar `item()` conversions used by UAA metadata collection do not get traced into `torch.compile`.
### UAA vs `auto_pad`
- `auto_pad=True` pads sequence tokens in `_sp_plan` and requires attention backends that support `attention_mask`.
- `advanced_uaa` does not depend on mask-based token padding inside Ulysses attention. It is therefore a better fit for non-divisible head counts and uneven Ulysses shard sizes.
- `auto_pad=True` remains incompatible with Ring attention because the ring backend does not consume `attention_mask`.
- `advanced_uaa` is still experimental and hybrid mode remains limited by Ring's equal-shape requirement.
---
## Approach 1: Non-Intrusive `_sp_plan` (Recommended)
The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries.
@@ -400,14 +364,9 @@ _sp_plan = {
- **Sequence Length not divisible by sp_size:**
**Problem:** strict Ulysses sequence parallel requires divisible shapes. If the shard length is uneven, or if the model head count is not divisible by `ulysses_degree`, the strict path will raise an error.
**Problem:** `SequenceParallelInput(auto_pad=False)` - auto_pad should be True to enable automatic sequence padding.
**Solutions:**
1. Use `ulysses_mode="advanced_uaa"` for Ulysses-SP when you want the experimental uneven-shape path without relying on attention-mask padding.
2. If the model already uses `_sp_plan` token padding and the attention backend supports `attention_mask`, set `auto_pad=True` and add attention-mask plumbing.
> **Experimental Feature:** `ulysses_mode="advanced_uaa"` is experimental. It is intended to relax Ulysses divisibility constraints, but hybrid Ulysses + Ring still requires equal post-Ulysses sequence lengths inside each ring group.
**Solution:** In `SequenceParallelInput`, set `auto_pad=True` and add attention mask support.
> **Experimental Feature:** `auto_pad=True` is an experimental feature and may be changed in the future. We plan to improve this solution to involve minimal changes to model files. More details are [here](https://github.com/vllm-project/vllm-omni/issues/1324).
+1 -1
View File
@@ -369,7 +369,7 @@ images = omni.generate(
2. **Compare performance** - Measure speedup vs baseline (expect 1.5x-2.0x)
3. **Verify output quality** - Visually compare cached vs uncached outputs (should be nearly identical)
See more detailed examples in [user guide for teacache](../../user_guide/diffusion/cache_acceleration/teacache.md).
See more detailed examples in [user guide for teacache](../../user_guide/diffusion/teacache.md).
---
+47 -86
View File
@@ -127,117 +127,80 @@ def step(self, requests: list[OmniDiffusionRequest]):
## 2. Scheduler
**Location**: `vllm_omni/diffusion/sched/`
**Location**: `vllm_omni/diffusion/scheduler.py`
### Architecture
The scheduler is a **request-state scheduler**. It owns request lifecycle management and scheduling decisions, while execution stays in `DiffusionEngine` and the executor.
The `Scheduler` is implemented as a **Singleton** pattern to ensure a single coordination point across the system, i.e., only one scheduler instance exists for coordination.
### Key Components
#### 2.1 Scheduler Interface
#### 2.1 Message Queue System
```python
class SchedulerInterface(ABC):
def add_request(self, request: OmniDiffusionRequest) -> str: ...
def schedule(self) -> DiffusionSchedulerOutput: ...
def update_from_output(
self,
sched_output: DiffusionSchedulerOutput,
output: DiffusionOutput,
) -> set[str]: ...
class Scheduler:
def initialize(self, od_config: OmniDiffusionConfig):
# Broadcast queue: scheduler -> all workers
self.mq = MessageQueue(
n_reader=od_config.num_gpus,
n_local_reader=od_config.num_gpus,
local_reader_ranks=list(range(od_config.num_gpus)),
)
# Result queue: rank 0 worker -> scheduler
self.result_mq = None # Initialized later
```
**Responsibilities**:
**Communication Pattern**:
- **Lifecycle contract**: Defines how the engine adds requests, triggers one scheduling cycle, and feeds executor results back.
- **Broadcast Queue**: One-to-many communication (scheduler → all workers)
- **Stable boundary**: `DiffusionSchedulerOutput` is the only scheduling result consumed by `DiffusionEngine`.
- **Result Queue**: One-to-one communication (rank 0 → scheduler)
- **Pluggability**: Different scheduler policies can reuse the same engine integration path.
- **Shared Memory**: Uses `MessageQueue` (ZMQ-based) for efficient IPC
#### 2.2 Request State Model
#### 2.2 Request Distribution
```python
class DiffusionRequestStatus(enum.IntEnum):
WAITING = ...
RUNNING = ...
PREEMPTED = ...
FINISHED_COMPLETED = ...
FINISHED_ABORTED = ...
FINISHED_ERROR = ...
def add_req(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
# Broadcast request to all workers
self.mq.enqueue(requests)
@dataclass
class DiffusionRequestState:
sched_req_id: str
req: OmniDiffusionRequest
status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
# Wait for result from Rank 0
output = self.result_mq.dequeue()
return output
```
**Design Features**:
- **Scheduler-owned ID**: Each `OmniDiffusionRequest` is tracked by an internal `sched_req_id`, separated from public `request_id` values.
- **Broadcast Model**: All workers receive the same request (for tensor parallelism)
- **Explicit lifecycle**: Requests move through waiting, running, optional preemption, and terminal states.
- **Single Response**: Only rank 0 sends results back (avoids duplicate outputs)
- **Centralized error handling**: Completion, abort, and error states are all normalized in the scheduler layer.
- **Synchronous**: Blocks until result is received (can be made async)
#### 2.3 Shared Bookkeeping in `_BaseScheduler`
#### 2.3 Singleton Pattern
```python
class _BaseScheduler(SchedulerInterface):
def __init__(self) -> None:
self._request_states = {}
self._request_id_to_sched_req_id = {}
self._waiting = deque()
self._running = []
self._finished_req_ids = set()
self._max_batch_size = 1
class Scheduler:
_instance = None
def __new__(cls, *args, **kwargs):
if not cls._instance:
cls._instance = super().__new__(cls)
return cls._instance
# Global singleton instance
scheduler = Scheduler()
```
**Design Features**:
**Benefits**:
- **Common state storage**: Shared request maps and waiting/running sets live in the base class.
- **Single Point of Control**: Ensures consistent state
- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
- **Easy Access**: Global `scheduler` instance accessible everywhere
- **Current constraint**: `_max_batch_size` remains `1` because the current engine path is still synchronous request-mode execution.
#### 2.4 Current `RequestScheduler` Policy
```python
class RequestScheduler(_BaseScheduler):
def schedule(self) -> DiffusionSchedulerOutput:
# 1. keep existing RUNNING requests in the scheduling result
# 2. pull WAITING requests while capacity remains
# 3. move newly admitted requests into RUNNING
```
**Behavior**:
- **FIFO request scheduling**: Waiting requests are promoted in queue order.
- **Single-request admission**: The current policy only admits one active request at a time.
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished scheduler ids.
#### 2.5 Engine-Driven Execution Loop
```python
sched_req_id = scheduler.add_request(request)
while True:
sched_output = scheduler.schedule()
output = executor.add_req(req)
finished_req_ids = scheduler.update_from_output(sched_output, output)
```
**Design Decisions**:
- **Separation of concerns**: Scheduler manages state and policy; executor handles runtime execution.
- **No scheduler-owned IPC**: Scheduler no longer talks to workers directly.
- **Conservative concurrency**: The current request-mode implementation still allows only one active request at a time.
- **Resource Management**: Centralized queue management
---
@@ -917,9 +880,8 @@ def initialize_model_parallel(
└─> Model-specific transformations
3. Scheduling
└─> scheduler.add_request(request)
└─> scheduler.schedule()
└─> DiffusionEngine submits scheduled request to executor.add_req(req)
└─> scheduler.add_req(requests)
└─> Broadcast via MessageQueue to all workers
4. Worker Execution
└─> WorkerProc.worker_busy_loop()
@@ -933,9 +895,8 @@ def initialize_model_parallel(
└─> vae.decode()
5. Result Collection
└─> Executor returns DiffusionOutput
└─> scheduler.update_from_output(...)
└─> DiffusionEngine pops finished request state
└─> Rank 0 sends DiffusionOutput via result queue
└─> Scheduler receives and returns
6. Post-processing
└─> post_process_func(output)
+1 -1
View File
@@ -30,7 +30,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
### Pre-built wheels
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, and 0.18.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
=== "NVIDIA CUDA"
@@ -29,23 +29,18 @@ uv pip install vllm --torch-backend=auto
uv pip install vllm-omni
```
To run Gradio demos, also install the optional extras:
```bash
uv pip install 'vllm-omni[demo]'
```
# --8<-- [end:pre-built-wheels]
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library
If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
```bash
uv pip install vllm==0.18.0 --torch-backend=auto
uv pip install vllm==0.17.0 --torch-backend=auto
```
The 0.18.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
The release 0.17.0 of vLLM is based on PyTorch 2.10.0 which requires CUDA 12.9 environment.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -55,22 +50,17 @@ cd vllm-omni
uv pip install -e .
```
To run Gradio demos, install with optional extras:
```bash
uv pip install -e '.[demo]'
```
<details><summary>(Optional) Installation of vLLM from source</summary>
If you want to check, modify or debug with source code of vLLM, install the library from source with the following instructions:
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.18.0
git checkout v0.17.0
```
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
```bash
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.18.0/vllm-0.18.0+cu129-cp38-abi3-manylinux_2_35_x86_64.whl
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_x86_64.whl
```
Install vllm with command below (If you have no existing PyTorch).
```bash
@@ -101,7 +91,7 @@ docker run --runtime nvidia --gpus 2 \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
--ipc=host \
vllm/vllm-omni:v0.18.0 \
vllm/vllm-omni:v0.16.0 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
```
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
```
#### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library
If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
```bash
uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
```
The pre-built 0.18.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
The release 0.17.0 of vLLM requires ROCm 7.0 environment.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.18.0
git checkout v0.17.0
python3 -m pip install -r requirements/rocm.txt
python3 setup.py develop
```
@@ -130,7 +130,7 @@ docker run --rm \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
vllm/vllm-omni-rocm:v0.18.0 \
vllm/vllm-omni-rocm:v0.16.0 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
@@ -149,7 +149,7 @@ docker run --rm -it \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--entrypoint bash \
vllm/vllm-omni-rocm:v0.18.0
vllm/vllm-omni-rocm:v0.16.0
```
# --8<-- [end:pre-built-images]
@@ -10,10 +10,10 @@ The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built
```bash
# Update the vllm-ascend image
# Atlas A2:
# export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1
# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
# Atlas A3:
# export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1-a3
export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1
# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1-a3
export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
docker run --rm \
--name vllm-omni-npu \
--shm-size=1g \
@@ -34,19 +34,20 @@ docker run --rm \
-it $IMAGE bash
cd /vllm-workspace/vllm
git pull origin main
git fetch origin --tags
git checkout v0.18.0
git checkout v0.16.0
# Because vllm-ascend will release v0.18.0rc1 after vllm-omni 0.16.0,
# Because vllm-ascend will release v0.16.0rc1 after vllm-omni 0.16.0,
# we have to pin vllm-ascend at the current commit.
cd /vllm-workspace/vllm-ascend
git pull origin main
git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
pip install -v -e .
# Inside the container, install vLLM-Omni from source
cd /vllm-workspace
git clone -b v0.18.0 https://github.com/vllm-project/vllm-omni.git
git clone -b v0.16.0 https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
pip install -v -e . --no-build-isolation
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
@@ -69,19 +70,21 @@ You can also build vLLM-Omni from the latest main branch if you want to use the
```bash
# Pin vLLM version to 0.18.0
cd /vllm-workspace/vllm
git pull origin main
git fetch origin --tags
git checkout v0.18.0
VLLM_TARGET_DEVICE=empty pip install -v -e .
# Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
cd /vllm-workspace/vllm-ascend
git pull origin main
git fetch origin --tags
git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
git checkout 1e05c4908f31737bc4eef865a9f351d030a77c9d
pip install -v -e .
# Install vLLM-Omni from the latest main branch
git clone https://github.com/vllm-project/vllm-omni.git
cd /vllm-workspace/vllm-omni
git clone https://github.com/vllm-project/vllm-omni.git
pip install -v -e . --no-build-isolation
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
export VLLM_WORKER_MULTIPROC_METHOD=spawn
@@ -104,10 +107,10 @@ Here's an example deployment command that has been verified on 4 x NPUs:
```bash
# Atlas A2:
# export IMAGE=quay.io/ascend/vllm-omni:v0.18.0
# export IMAGE=quay.io/ascend/vllm-omni:v0.16.0
# Atlas A3:
# export IMAGE=quay.io/ascend/vllm-omni:v0.18.0-a3
export IMAGE=quay.io/ascend/vllm-omni:v0.18.0
# export IMAGE=quay.io/ascend/vllm-omni:v0.16.0-a3
export IMAGE=quay.io/ascend/vllm-omni:v0.16.0
docker run --rm \
--name vllm-omni-npu \
--shm-size=1g \
+2 -2
View File
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
source .venv/bin/activate
# On CUDA
uv pip install vllm==0.18.0 --torch-backend=auto
uv pip install vllm==0.17.0 --torch-backend=auto
# On ROCm
uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
uv pip install vllm==0.17.0 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
-2
View File
@@ -139,8 +139,6 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]:
# Skip excluded modules (avoid importing vllm during docs build)
excluded_prefixes = [
"vllm_omni.diffusion.models.qwen_image",
"vllm_omni.diffusion.quantization",
"vllm_omni.quantization",
"vllm_omni.entrypoints.async_diffusion",
"vllm_omni.entrypoints.openai",
"vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts",
+69 -43
View File
@@ -6,7 +6,7 @@ vLLM-Omni supports unified multimodal comprehension and generation models across
If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm_omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.
## List of Supported Models
## List of Supported Models for Nvidia GPU / AMD GPU
<style>
th {
@@ -15,46 +15,72 @@ th {
}
</style>
| Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU |
|--------------|--------|-------------------|------------|---------|-----|-----------|
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | |
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ |
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
| `MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | |
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
| Architecture | Models | Example HF Models |
|--------------|--------|-------------------|
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` |
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` |
|`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` |
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` |
|`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
|`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
|`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
|`CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` |
|`MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` |
|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
|`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
|`OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` |
|`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` |
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
|`GlmImageForConditionalGeneration` | GLM-Image | `zai-org/GLM-Image` |
|`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` |
|`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` |
|`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
|`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` |
|`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` |
|`VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` |
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
## List of Supported Models for NPU
<style>
th {
white-space: nowrap;
min-width: 0 !important;
}
</style>
| Architecture | Models | Example HF Models |
|--------------|--------|-------------------|
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` |
|`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
|`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
|`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
-78
View File
@@ -1,78 +0,0 @@
# Diffusion Chat Completions API
vLLM-Omni supports generating and editing images via the `/v1/chat/completions`
endpoint using diffusion models. This page explains how to pass generation
parameters (such as `num_inference_steps`, `height`, `width`) to diffusion
models through this endpoint.
!!! tip
For dedicated endpoints that accept generation parameters as top-level
fields, see [Image Generation API](image_generation_api.md) and
[Image Edit API](image_edit_api.md).
## Passing Generation Parameters
The `/v1/chat/completions` endpoint follows the OpenAI Chat API schema, which
does not natively include diffusion-specific fields like `num_inference_steps`
or `height`. How you pass these extra fields depends on your client.
### curl / Python `requests`
Wrap generation parameters inside an `"extra_body"` key in the JSON body:
```bash
curl -s http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"messages": [
{"role": "user", "content": "A beautiful landscape painting"}
],
"extra_body": {
"num_inference_steps": 50,
"seed": 42
}
}'
```
### OpenAI Python SDK
Use the `extra_body` **keyword argument**. The SDK automatically merges these
fields into the top-level request body:
```python
response = client.chat.completions.create(
model="Qwen/Qwen-Image",
messages=[{"role": "user", "content": "A beautiful landscape painting"}],
extra_body={
"num_inference_steps": 50,
"seed": 42,
},
)
```
!!! note "SDK `extra_body` vs. JSON `extra_body`"
These two `extra_body` usages look similar but work differently under the
hood. The SDK flattens the dict into the top-level request JSON, while the
curl/requests approach sends it as a nested `"extra_body"` key. Both are
handled correctly by the server.
!!! note "About the `ignored fields` warning"
You may see a log message like:
```
WARNING: The following fields were present in the request but ignored: {'height', 'width', ...}
```
This is **harmless**. It is emitted by vLLM's request validation layer
because these fields are not part of the standard OpenAI
`ChatCompletionRequest` schema. The fields are still stored internally
and correctly forwarded to the diffusion pipeline.
## Model-Specific Examples
For complete examples with full request/response details, see the model-specific
guides:
- [Text-to-Image (Qwen-Image)](../user_guide/examples/online_serving/text_to_image.md)
- [Image-to-Image (Qwen-Image-Edit, Qwen-Image-Layered)](../user_guide/examples/online_serving/image_to_image.md)
- [GLM-Image](../user_guide/examples/online_serving/glm_image.md)
+4 -174
View File
@@ -160,7 +160,6 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
| `audio_sample` | file | Yes | Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4) |
| `consent` | string | Yes | Consent recording ID |
| `name` | string | Yes | Name for the new voice |
| `ref_text` | string | No | Transcript of the audio. When provided, enables in-context voice cloning (higher quality). Without it, only the speaker embedding is extracted. |
**Response Example:**
@@ -183,8 +182,7 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
curl -X POST http://localhost:8091/v1/audio/voices \
-F "audio_sample=@/path/to/voice_sample.wav" \
-F "consent=user_consent_id" \
-F "name=custom_voice_1" \
-F "ref_text=The exact transcript of the audio sample."
-F "name=custom_voice_1"
```
## Streaming Text Input (WebSocket)
@@ -319,9 +317,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
}' --output cloned.wav
```
### Upload Voice
Upload voice (speaker embedding only):
upload voice
```bash
curl -X POST http://localhost:8091/v1/audio/voices \
-F "audio_sample=@/path/to/voice_sample.wav" \
@@ -329,183 +325,17 @@ curl -X POST http://localhost:8091/v1/audio/voices \
-F "name=custom_voice_1"
```
Upload voice with transcript (in-context cloning, higher quality):
```bash
curl -X POST http://localhost:8091/v1/audio/voices \
-F "audio_sample=@/path/to/voice_sample.wav" \
-F "consent=user_consent_id" \
-F "name=custom_voice_2" \
-F "ref_text=The exact transcript of the audio sample."
```
### Use Uploaded Voice
use upload voice
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, this is a cloned voice",
"task_type": "Base",
"voice": "custom_voice_1"
}' --output cloned.wav
```
## Batch Speech Generation
The batch endpoint synthesizes multiple texts in a single request, returning all results as JSON with base64-encoded audio.
### Endpoint
```
POST /v1/audio/speech/batch
Content-Type: application/json
```
### Request Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `items` | array | **required** | List of items to synthesize (1–32) |
| `model` | string | server's model | Model to use |
| `voice` | string | null | Default voice for all items |
| `response_format` | string | "wav" | Default audio format for all items |
| `speed` | float | 1.0 | Default playback speed (0.25–4.0) |
| `task_type` | string | null | Default TTS task type |
| `language` | string | null | Default language |
| `instructions` | string | null | Default voice style instructions |
| `ref_audio` | string | null | Default reference audio (Base task) |
| `ref_text` | string | null | Default reference transcript (Base task) |
| `max_new_tokens` | integer | null | Default max tokens |
Each item in the `items` array requires only `input` (the text). All other fields are optional and override the batch-level defaults when set:
| Field | Type | Description |
|-------|------|-------------|
| `input` | string | **required** — text to synthesize |
| `voice` | string | Override voice for this item |
| `response_format` | string | Override format for this item |
| `speed` | float | Override speed for this item |
| `task_type` | string | Override task type |
| `language` | string | Override language |
| `instructions` | string | Override instructions |
| `ref_audio` | string | Override reference audio |
| `ref_text` | string | Override reference transcript |
| `max_new_tokens` | integer | Override max tokens |
### Response Format
```json
{
"id": "speech-batch-abc123",
"results": [
{
"index": 0,
"status": "success",
"audio_data": "<base64-encoded audio>",
"media_type": "audio/wav"
},
{
"index": 1,
"status": "error",
"error": "Input text cannot be empty"
}
],
"total": 2,
"succeeded": 1,
"failed": 1
}
```
### Examples
**Basic batch with shared defaults:**
```bash
curl -X POST http://localhost:8091/v1/audio/speech/batch \
-H "Content-Type: application/json" \
-d '{
"items": [
{"input": "Hello, how are you?"},
{"input": "Goodbye, see you later!"}
],
"voice": "vivian",
"language": "English"
}'
```
**Per-item overrides (different voices and formats):**
```bash
curl -X POST http://localhost:8091/v1/audio/speech/batch \
-H "Content-Type: application/json" \
-d '{
"items": [
{"input": "Hello!", "voice": "vivian", "response_format": "mp3"},
{"input": "你好!", "voice": "ryan", "language": "Chinese"}
],
"response_format": "wav"
}'
```
**Voice cloning with shared reference audio (Base task):**
```bash
curl -X POST http://localhost:8091/v1/audio/speech/batch \
-H "Content-Type: application/json" \
-d '{
"items": [
{"input": "First sentence in the cloned voice."},
{"input": "Second sentence in the cloned voice."}
],
"task_type": "Base",
"ref_audio": "https://example.com/reference.wav",
"ref_text": "Transcript of the reference audio"
}'
```
Setting `ref_audio` at the batch level applies it to all items, avoiding the need to repeat it per item.
**Decoding the response in Python:**
```python
import base64
import httpx
response = httpx.post(
"http://localhost:8091/v1/audio/speech/batch",
json={
"items": [
{"input": "First sentence."},
{"input": "Second sentence."},
],
"voice": "vivian",
},
timeout=300.0,
)
for result in response.json()["results"]:
if result["status"] == "success":
audio_bytes = base64.b64decode(result["audio_data"])
with open(f"output_{result['index']}.wav", "wb") as f:
f.write(audio_bytes)
```
### Configuration
| Parameter | Source | Default | Description |
|-----------|--------|---------|-------------|
| `tts_batch_max_items` | engine kwarg | 32 | Maximum number of items per batch request |
All items are fanned out to `generate()` concurrently. The engine's stage worker automatically batches them up to the configured `max_batch_size` and queues the rest — no client-side throttling needed.
For best throughput, use a batch-optimized stage config with `max_batch_size > 1`:
```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml \
--omni --port 8091 --trust-remote-code --enforce-eager
```
The default `qwen3_tts.yaml` uses `max_batch_size: 1` (single request). The `qwen3_tts_batch.yaml` config sets `max_batch_size: 4` for ~4x throughput.
## Supported Models
### Qwen3-TTS
@@ -1,285 +0,0 @@
# Cache-DiT Guide
## Table of Content
- [Overview](#overview)
- [Quick Start](#quick-start)
- [Example Script](#example-script)
- [Acceleration Methods](#acceleration-methods)
- [Configuration Parameters](#configuration-parameters)
- [Best Practices](#best-practices)
- [Troubleshooting](#troubleshooting)
- [Summary](#summary)
- [Additional Resources](#additional-resources)
---
## Overview
Cache-DiT accelerates diffusion transformer models through intelligent caching mechanisms, providing significant speedup with minimal quality loss. It supports multiple acceleration techniques that can be combined for optimal performance:
- **DBCache**: Dual Block Cache for reducing redundant computations
- **TaylorSeer**: Taylor expansion-based forecasting for faster inference
- **SCM**: Step Computation Masking for selective step computation
See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
---
## Quick Start
### Basic Usage
Enable cache-dit acceleration by simply setting `cache_backend="cache_dit"`:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="Qwen/Qwen-Image",
cache_backend="cache_dit", # Enable Cache-DiT with defaults
)
outputs = omni.generate(
"a beautiful landscape",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
**Note**: When `cache_config` is not provided, Cache-DiT uses optimized default values. See the [Configuration Parameters](#configuration-parameters) section for details.
### Custom Configuration
To customize cache-dit settings, provide a `cache_config` dictionary, for example:
```python
omni = Omni(
model="Qwen/Qwen-Image",
cache_backend="cache_dit",
cache_config={
"Fn_compute_blocks": 1,
"Bn_compute_blocks": 0,
"max_warmup_steps": 4,
"residual_diff_threshold": 0.12,
},
)
```
---
## Example Script
### Offline Inference
Use the example script under `examples/offline_inference/text_to_image`:
```bash
cd examples/offline_inference/text_to_image
python text_to_image.py \
--model Qwen/Qwen-Image \
--prompt "a cup of coffee on the table" \
--cache-backend cache_dit \
--num-inference-steps 50
```
See the [text_to_image.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/text_to_image/text_to_image.py) for detailed configuration options.
The script uses cache-dit acceleration with a hybrid configuration combining DBCache, SCM, and TaylorSeer:
```python
omni = Omni(
model="Qwen/Qwen-Image",
cache_backend="cache_dit",
cache_config={
# Scheme: Hybrid DBCache + SCM + TaylorSeer
"Fn_compute_blocks": 1, # Optimized for single-transformer models
"Bn_compute_blocks": 0, # Number of backward compute blocks
"max_warmup_steps": 4, # Maximum warmup steps (works for few-step models)
"residual_diff_threshold": 0.24, # Higher threshold for more aggressive caching
"max_continuous_cached_steps": 3, # Limit to prevent precision degradation
# TaylorSeer parameters [cache-dit only]
"enable_taylorseer": False, # Disabled by default (not suitable for few-step models)
"taylorseer_order": 1, # TaylorSeer polynomial order
# SCM (Step Computation Masking) parameters [cache-dit only]
"scm_steps_mask_policy": None, # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
"scm_steps_policy": "dynamic", # SCM steps policy: "dynamic" or "static"
}
)
```
You can customize the configuration by modifying the `cache_config` dictionary to use only specific methods (e.g., DBCache only, DBCache + SCM, etc.) based on your quality and speed requirements.
For image-to-image tasks, use the example script under `examples/offline_inference/image_to_image`:
```bash
cd examples/offline_inference/image_to_image
python image_edit.py \
--model Qwen/Qwen-Image-Edit \
--prompt "make the sky more colorful" \
--image path/to/input/image.jpg \
--cache-backend cache_dit \
--num-inference-steps 50 \
--cache-dit-max-continuous-cached-steps 3 \
--cache-dit-residual-diff-threshold 0.24 \
--cache-dit-enable-taylorseer
```
See the [image_edit.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py) for detailed configuration options.
### Online Serving
```bash
# Default configuration (recommended)
vllm serve Qwen/Qwen-Image --omni --port 8091 --cache-backend cache_dit
# Custom configuration
vllm serve Qwen/Qwen-Image --omni --port 8091 \
--cache-backend cache_dit \
--cache-config '{"Fn_compute_blocks": 1, "residual_diff_threshold": 0.12}'
```
---
## Acceleration Methods
For comprehensive illustration, please view Cache-DiT [User Guide](https://cache-dit.readthedocs.io/en/latest/user_guide/OVERVIEWS/).
### 1. DBCache (Dual Block Cache)
DBCache intelligently caches intermediate transformer block outputs when the residual differences between consecutive steps are small, reducing redundant computations without sacrificing quality.
**Example Configuration**:
```python
cache_config={
"Fn_compute_blocks": 8, # Use first 8 blocks for difference computation
"Bn_compute_blocks": 0, # No additional fusion blocks
"max_warmup_steps": 8, # Cache after 8 warmup steps
"residual_diff_threshold": 0.12, # Lower threshold for faster inference
"max_cached_steps": -1, # No limit on cached steps
}
```
**Performance Tips**:
- Default `Fn_compute_blocks=1` works well for most cases. Some models (e.g., [FLUX.2-klein](https://github.com/wtomin/vllm-omni/blob/main/vllm_omni/diffusion/cache/cache_dit_backend.py#L363)) use a larger value for `Fn_compute_blocks` for a balanced performance.
- Increase `residual_diff_threshold` (e.g., 0.12-0.15) for faster inference with slight quality trade-off, or decrease from default 0.24 for higher quality.
- Default `max_warmup_steps=4` is optimized for few-step models. Increase to 6-8 for more steps if needed.
### 2. TaylorSeer
TaylorSeer uses Taylor expansion to forecast future hidden states, allowing the model to skip some computation steps while maintaining quality.
**Example Configuration**:
```python
cache_config={
"enable_taylorseer": True,
"taylorseer_order": 1, # First-order Taylor expansion
}
```
**Performance Tips**:
- TaylorSeer is **not suitable for few-step distilled models**.
- Use `taylorseer_order=1` for most cases (good balance of speed and quality).
- Combine with DBCache for maximum acceleration.
- Higher orders (2-3) may improve quality but reduce speed gains.
### 3. SCM (Step Computation Masking)
SCM allows you to specify which steps must be computed and which can use cached results, similar to LeMiCa/EasyCache style acceleration.
`scm_steps_mask_policy` options (number of compute steps out of 28):
| Policy | Compute Steps | Speed | Quality |
|--------|--------------|-------|---------|
| `None` (default) | All | Baseline | Best |
| `"slow"` | 18 / 28 | Moderate | High |
| `"medium"` | 15 / 28 | Balanced | Good |
| `"fast"` | 11 / 28 | Fast | Moderate |
| `"ultra"` | 8 / 28 | Fastest | Lower |
**Example Configuration**:
```python
cache_config={
"scm_steps_mask_policy": "medium", # Balanced speed/quality
"scm_steps_policy": "dynamic", # Use dynamic cache
}
```
**Performance Tips**:
- SCM is disabled by default. Enable it by setting a policy value if you need additional acceleration.
- Start with `"medium"` policy and adjust based on quality requirements.
- Use `"fast"` or `"ultra"` for maximum speed when quality can be slightly compromised.
- `"dynamic"` policy generally provides better quality than `"static"`.
- SCM mask is automatically regenerated when `num_inference_steps` changes during inference.
---
## Configuration Parameters
In `cache_config` passed to `Omni` constructor, it accepts the arguments of `DBCacheConfig` ([Cache-DiT API Reference](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/)). Key parameters are listed below:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `Fn_compute_blocks` | int | 1 | First n blocks for difference computation (optimized for single-transformer models) |
| `Bn_compute_blocks` | int | 0 | Last n blocks for fusion |
| `max_warmup_steps` | int | 4 | Steps before caching starts (optimized for few-step distilled models) |
| `max_cached_steps` | int | -1 | Max cached steps (-1 = unlimited) |
| `max_continuous_cached_steps` | int | 3 | Max consecutive cached steps (prevents precision degradation) |
| `residual_diff_threshold` | float | 0.24 | Residual difference threshold (higher for more aggressive caching) |
| `num_inference_steps` | int \| None | None | Initial inference steps for SCM mask generation (optional, auto-refreshed during inference) |
| `enable_taylorseer` | bool | False | Enable TaylorSeer acceleration (not suitable for few-step distilled models) |
| `taylorseer_order` | int | 1 | Taylor expansion order |
| `scm_steps_mask_policy` | str \| None | None | SCM mask policy (None, "slow", "medium", "fast", "ultra") |
| `scm_steps_policy` | str | "dynamic" | SCM computation policy ("dynamic" or "static") |
---
## Best Practices
### When to Use
**Good for:**
- Production deployments requiring fast inference
- Diffusion transformer models (DiT architecture)
- Scenarios where 1.5x-3x speedup is valuable
**Not for:**
- Non-DiT architectures (use model-specific acceleration instead)
- Models already using few-step distillation (< 10 steps)
---
## Troubleshooting
### Common Issue 1: Quality Degradation
**Symptoms**: Generated images have visible artifacts or lower quality
**Solution**:
```python
# Reduce aggressiveness - use more conservative settings
cache_config={
"residual_diff_threshold": 0.20, # Lower threshold (closer to default 0.24)
"Fn_compute_blocks": 8, # Use more blocks for better decisions
"max_warmup_steps": 6, # Longer warmup
"scm_steps_mask_policy": "slow", # More compute steps
}
```
---
## Summary
Using Cache-DiT acceleration:
1. ✅ **Enable Cache-DiT** - Set `cache_backend="cache_dit"` to get 1.5x-3x speedup with optimized defaults
2. ✅ **(Optional) Customize** - Adjust `cache_config` parameters for specific speed/quality trade-offs
@@ -1,194 +0,0 @@
# TeaCache Guide
## Table of Content
- [Overview](#overview)
- [Quick Start](#quick-start)
- [Example Script](#example-script)
- [Configuration Parameters](#configuration-parameters)
- [Best Practices](#best-practices)
- [Troubleshooting](#troubleshooting)
- [Summary](#summary)
---
## Overview
TeaCache accelerates diffusion model inference by caching transformer computations when consecutive timesteps are similar, providing **1.5x-2.0x speedup** with minimal quality loss. It dynamically decides whether to reuse cached outputs based on input similarity, making it ideal for production deployments where inference speed matters without sacrificing generation quality.
See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
---
## Quick Start
### Basic Usage
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="Qwen/Qwen-Image",
cache_backend="tea_cache",
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
### Custom Configuration
```python
omni = Omni(
model="Qwen/Qwen-Image",
cache_backend="tea_cache",
cache_config={
"rel_l1_thresh": 0.2, # Controls speed/quality tradeoff
},
)
```
### Using Environment Variable
You can also enable TeaCache via environment variable:
```bash
export DIFFUSION_CACHE_BACKEND=tea_cache
```
Then initialize without explicitly setting `cache_backend`:
```python
from vllm_omni import Omni
omni = Omni(
model="Qwen/Qwen-Image",
cache_config={"rel_l1_thresh": 0.2}
)
```
---
## Example Script
### Offline Inference
Use python script under `examples/offline_inference/text_to_image/` or `examples/offline_inference/image_to_image/` with CLI:
```bash
# Text-to-image example
python examples/offline_inference/text_to_image/text_to_image.py \
--model Qwen/Qwen-Image \
--cache-backend tea_cache
# Image-to-image example
python examples/offline_inference/image_to_image/image_edit.py \
--model Qwen/Qwen-Image-Edit \
--image input.png \
--prompt "Edit description" \
--cache-backend tea_cache \
--tea-cache-rel-l1-thresh 0.25
```
See the [text_to_image.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/text_to_image/text_to_image.py) or [image_edit.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py) for detailed configuration options.
### Online Serving
```bash
# Default configuration
vllm serve Qwen/Qwen-Image --omni --port 8091 --cache-backend tea_cache
# Custom configuration
vllm serve Qwen/Qwen-Image --omni --port 8091 \
--cache-backend tea_cache \
--cache-config '{"rel_l1_thresh": 0.2}'
```
---
## Configuration Parameters
In `OmniDiffusionConfig`
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `rel_l1_thresh` | float | `0.2` | Similarity threshold for cache reuse. Lower values prioritize quality (less caching), higher values prioritize speed (more caching). Suggested range: 0.1-0.8 |
| `coefficients` | list[float] \| None | `None` | Polynomial coefficients for rescaling L1 distance. Must contain exactly 5 elements if provided. If `None`, uses model-specific defaults based on transformer type. |
Users can find the default model coefficients in [`vllm_omni/diffusion/cache/teacache/config.py`](https://github.com/vllm-project/vllm-omni/blob/main/vllm_omni/diffusion/cache/teacache/config.py), for example:
```python
_MODEL_COEFFICIENTS = {
# Qwen-Image transformer coefficients from ComfyUI-TeaCache
# Tuned specifically for Qwen's dual-stream transformer architecture
# Used for all Qwen-Image Family pipelines, in general
"QwenImageTransformer2DModel": [
-4.50000000e02,
2.80000000e02,
-4.50000000e01,
3.20000000e00,
-2.00000000e-02,
],
...
}
```
---
## Best Practices
### When to Use
**Good for:**
- Production deployments requiring faster inference, tolerant of minimal quality loss
- Scenarios where 1.5-2x speedup is valuable
- Useful for single-card acceleration
**Not for:**
- Maximum quality requirements where no degradation is acceptable
- Very short inference runs (< 20 steps) where caching overhead may outweigh benefits
---
## Troubleshooting
### Common Issue 1: Quality Degradation
**Symptoms**: Generated images show artifacts, reduced detail, or inconsistent quality compared to non-cached results
**Solution**:
```python
# Lower the threshold for more conservative caching
cache_config={"rel_l1_thresh": 0.1}
```
### Common Issue 2: Limited Speedup
**Symptoms**: Actual speedup is less than expected (< 1.3x)
**Solutions**:
1. Increase the threshold to enable more aggressive caching:
```python
cache_config={"rel_l1_thresh": 0.8}
```
2. Ensure you're using sufficient inference steps (35+ recommended)
3. Check that your model architecture is supported (see Supported Models section)
---
## Summary
1. ✅ **Enable TeaCache** - Set `cache_backend="tea_cache"` to get 1.5x-2.0x speedup with optimized defaults
2. ✅ **(Optional) Customize** - Adjust thresholds and polynomial coefficients for specific speed/quality trade-offs
@@ -0,0 +1,228 @@
# Cache-DiT Acceleration Guide
This guide explains how to use cache-dit acceleration in vLLM-Omni to speed up diffusion model inference.
## Overview
Cache-dit is a library that accelerates diffusion transformer models through intelligent caching mechanisms. It supports multiple acceleration techniques that can be combined for optimal performance:
- **DBCache**: Dual Block Cache for reducing redundant computations
- **TaylorSeer**: Taylor expansion-based forecasting for faster inference
- **SCM**: Step Computation Masking for selective step computation
## Quick Start
### Basic Usage
Enable cache-dit acceleration by simply setting `cache_backend="cache_dit"`. Cache-dit will use its recommended default parameters:
```python
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# Simplest way: just enable cache-dit with default parameters
omni = Omni(
model="Qwen/Qwen-Image",
cache_backend="cache_dit",
)
images = omni.generate(
"a beautiful landscape",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
**Default Parameters**: When `cache_config` is not provided, cache-dit uses optimized default values. See the [Configuration Reference](#configuration-reference) section for a complete list of all parameters and their default values.
### Custom Configuration
To customize cache-dit settings, provide a `cache_config` dictionary, for example:
```python
omni = Omni(
model="Qwen/Qwen-Image",
cache_backend="cache_dit",
cache_config={
"Fn_compute_blocks": 1,
"Bn_compute_blocks": 0,
"max_warmup_steps": 4,
"residual_diff_threshold": 0.12,
},
)
```
## Online Serving (OpenAI-Compatible)
Enable Cache-DiT for online serving by passing `--cache-backend cache_dit` when starting the server:
```bash
# Use Cache-DiT default (recommended) parameters
vllm serve Qwen/Qwen-Image --omni --port 8091 --cache-backend cache_dit
```
To customize Cache-DiT settings for online serving, pass a JSON string via `--cache-config`:
```bash
vllm serve Qwen/Qwen-Image --omni --port 8091 \
--cache-backend cache_dit \
--cache-config '{"Fn_compute_blocks": 1, "Bn_compute_blocks": 0, "max_warmup_steps": 4, "residual_diff_threshold": 0.12}'
```
## Acceleration Methods
For comprehensive illustration, please view cache-dit [User_Guide](https://cache-dit.readthedocs.io/en/latest/user_guide/OVERVIEWS/)
### 1. DBCache (Dual Block Cache)
DBCache intelligently caches intermediate transformer block outputs when the residual differences between consecutive steps are small, reducing redundant computations without sacrificing quality.
**Key Parameters**:
- `Fn_compute_blocks` (int, default: 1): Number of **first n** transformer blocks used to compute stable feature differences. Higher values provide more accurate caching decisions but increase computation.
- `Bn_compute_blocks` (int, default: 0): Number of **last n** transformer blocks used for additional fusion. These blocks act as an auto-scaler for approximate hidden states.
- `max_warmup_steps` (int, default: 4): Number of initial steps where caching is disabled to ensure the model learns sufficient features before caching begins. Optimized for few-step distilled models.
- `residual_diff_threshold` (float, default: 0.24): Threshold for residual difference. Higher values lead to faster performance but may reduce precision. Default uses a relatively higher threshold for more aggressive caching.
- `max_cached_steps` (int, default: -1): Maximum number of cached steps. Set to -1 for unlimited caching.
- `max_continuous_cached_steps` (int, default: 3): Maximum number of consecutive cached steps. Limits consecutive caching to prevent precision degradation.
**Example Configuration**:
```python
cache_config={
"Fn_compute_blocks": 8, # Use first 8 blocks for difference computation
"Bn_compute_blocks": 0, # No additional fusion blocks
"max_warmup_steps": 8, # Cache after 8 warmup steps
"residual_diff_threshold": 0.12, # Higher threshold for faster inference
"max_cached_steps": -1, # No limit on cached steps
}
```
**Performance Tips**:
- Default `Fn_compute_blocks=1` works well for most cases. Increase to 8-12 for larger models or when more accuracy is needed
- Increase `residual_diff_threshold` (e.g., 0.12-0.15) for faster inference with slight quality trade-off, or decrease from default 0.24 for higher quality
- Default `max_warmup_steps=4` is optimized for few-step models. Increase to 6-8 for more steps if needed
### 2. TaylorSeer
TaylorSeer uses Taylor expansion to forecast future hidden states, allowing the model to skip some computation steps while maintaining quality.
**Key Parameters**:
- `enable_taylorseer` (bool, default: False): Enable TaylorSeer acceleration
- `taylorseer_order` (int, default: 1): Order of Taylor expansion. Higher orders provide better accuracy but require more computation.
**Example Configuration**:
```python
cache_config={
"enable_taylorseer": True,
"taylorseer_order": 1, # First-order Taylor expansion
}
```
**Performance Tips**:
- Use `taylorseer_order=1` for most cases (good balance of speed and quality)
- Combine with DBCache for maximum acceleration
- Higher orders (2-3) may improve quality but reduce speed gains
### 3. SCM (Step Computation Masking)
SCM allows you to specify which steps must be computed and which can use cached results, similar to LeMiCa/EasyCache style acceleration.
**Key Parameters**:
- `scm_steps_mask_policy` (str | None, default: None): Predefined mask policy. Options:
- `None`: SCM disabled (default)
- `"slow"`: More compute steps, higher quality (18 compute steps out of 28)
- `"medium"`: Balanced (15 compute steps out of 28)
- `"fast"`: More cache steps, faster inference (11 compute steps out of 28)
- `"ultra"`: Maximum speed (8 compute steps out of 28)
- `scm_steps_policy` (str, default: "dynamic"): Policy for cached steps:
- `"dynamic"`: Use dynamic cache for masked steps (recommended)
- `"static"`: Use static cache for masked steps
**Example Configuration**:
```python
cache_config={
"scm_steps_mask_policy": "medium", # Balanced speed/quality
"scm_steps_policy": "dynamic", # Use dynamic cache
}
```
**Performance Tips**:
- SCM is disabled by default (`scm_steps_mask_policy=None`). Enable it by setting a policy value if you need additional acceleration
- Start with `"medium"` policy and adjust based on quality requirements
- Use `"fast"` or `"ultra"` for maximum speed when quality can be slightly compromised
- `"dynamic"` policy generally provides better quality than `"static"`
- SCM mask is automatically regenerated when `num_inference_steps` changes during inference
## Configuration Reference
### DiffusionCacheConfig Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `Fn_compute_blocks` | int | 1 | First n blocks for difference computation (optimized for single-transformer models) |
| `Bn_compute_blocks` | int | 0 | Last n blocks for fusion |
| `max_warmup_steps` | int | 4 | Steps before caching starts (optimized for few-step distilled models) |
| `max_cached_steps` | int | -1 | Max cached steps (-1 = unlimited) |
| `max_continuous_cached_steps` | int | 3 | Max consecutive cached steps (prevents precision degradation) |
| `residual_diff_threshold` | float | 0.24 | Residual difference threshold (higher for more aggressive caching) |
| `num_inference_steps` | int \| None | None | Initial inference steps for SCM mask generation (optional, auto-refreshed during inference) |
| `enable_taylorseer` | bool | False | Enable TaylorSeer acceleration (not suitable for few-step distilled models) |
| `taylorseer_order` | int | 1 | Taylor expansion order |
| `scm_steps_mask_policy` | str \| None | None | SCM mask policy (None, "slow", "medium", "fast", "ultra") |
| `scm_steps_policy` | str | "dynamic" | SCM computation policy ("dynamic" or "static") |
## Example: Accelerate Text-to-Image Generation with CacheDiT
See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example with cache-dit acceleration.
```bash
# Enable cache-dit with hybrid acceleration
cd examples/offline_inference/text_to_image
python text_to_image.py \
--model Qwen/Qwen-Image \
--prompt "a cup of coffee on the table" \
--cache-backend cache_dit \
--num-inference-steps 50
```
The script uses cache-dit acceleration with a hybrid configuration combining DBCache, SCM, and TaylorSeer:
```python
omni = Omni(
model="Qwen/Qwen-Image",
cache_backend="cache_dit",
cache_config={
# Scheme: Hybrid DBCache + SCM + TaylorSeer
# DBCache
"Fn_compute_blocks": 8,
"Bn_compute_blocks": 0,
"max_warmup_steps": 4,
"residual_diff_threshold": 0.12,
# TaylorSeer
"enable_taylorseer": True,
"taylorseer_order": 1,
# SCM
"scm_steps_mask_policy": "fast", # Set to None to disable SCM
"scm_steps_policy": "dynamic",
},
)
```
You can customize the configuration by modifying the `cache_config` dictionary to use only specific methods (e.g., DBCache only, DBCache + SCM, etc.) based on your quality and speed requirements.
To test another model, you can modify `--model` with the target model identifier like `Tongyi-MAI/Z-Image-Turbo` and update `cache_config` according the model architecture (e.g., number of transformer blocks).
## Additional Resources
- [Cache-DiT User Guide](https://cache-dit.readthedocs.io/en/latest/user_guide/OVERVIEWS/)
- [Cache-DiT Benchmark](https://cache-dit.readthedocs.io/en/latest/benchmark/HYBRID_CACHE/)
- [DBCache Technical Details](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/)
@@ -1,169 +0,0 @@
# CFG-Parallel Guide
## Table of Content
- [Overview](#overview)
- [Quick Start](#quick-start)
- [Example Script](#example-script)
- [Configuration Parameters](#configuration-parameters)
- [Best Practices](#best-practices)
- [Troubleshooting](#troubleshooting)
- [Summary](#summary)
---
## Overview
CFG-Parallel accelerates diffusion models by distributing positive and negative classifier-free guidance (CFG) passes across different GPUs, providing ~1.8x speedup when CFG is enabled. It's ideal for image editing tasks that require guidance scales greater than 1.0.
See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
---
## Quick Start
### Basic Usage
Simplest working example:
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from PIL import Image
omni = Omni(
model="Qwen/Qwen-Image-Edit",
parallel_config=DiffusionParallelConfig(cfg_parallel_size=2), # Enable CFG-Parallel
)
input_image = Image.open("input.png").convert("RGB")
outputs = omni.generate(
{
"prompt": "turn this cat to a dog",
"negative_prompt": "low quality, blurry",
"multi_modal_data": {"image": input_image},
},
OmniDiffusionSamplingParams(
true_cfg_scale=4.0,
num_inference_steps=50,
),
)
```
---
## Example Script
### Offline Inference
Use python script under `examples/offline_inference/image_to_image/image_edit.py`:
```bash
cd examples/offline_inference/image_to_image/
python image_edit.py \
--model "Qwen/Qwen-Image-Edit" \
--image "input.png" \
--prompt "turn this cat to a dog" \
--negative-prompt "low quality, blurry" \
--cfg-scale 4.0 \
--output "edited_image.png" \
--cfg-parallel-size 2
```
### Online Serving
Enable CFG-Parallel in online serving:
```bash
# Default configuration
vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2
```
---
## Configuration Parameters
In `DiffusionParallelConfig`
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `cfg_parallel_size` | int | 1 | Number of GPUs for CFG parallelism. Set to 2 to enable CFG-Parallel (rank 0 for positive, rank 1 for negative branch) |
!!! info
Most models support `cfg_parallel_size=2` (positive branch on rank 0, negative branch on rank 1). **Bagel** is an exception: it supports `cfg_parallel_size=3`, which adds a third branch on rank 2 for full three-way CFG parallelism.
---
## Best Practices
### When to Use
**Good for:**
- Tasks requiring classifier-free guidance
- Multi-GPU setups (at least 2 GPUs available)
- Combining with other parallelism methods (sequence/tensor parallel)
**Not for:**
- Single GPU setups
- Models that don't support CFG-Parallel (check [supported models](../../diffusion_features.md#supported-models))
- Workloads without negative prompts or classifier-free guidance
- Very short inference runs (< 10 steps) where parallelism overhead may outweigh benefits
### Expected Performance
| Configuration | Speedup | Quality | Use Case |
|--------------|---------|---------|----------|
| CFG-Parallel (2 GPUs) | 1.5~1.8x | No degradation | Large model, VRAM limited |
---
## Troubleshooting
### Common Issue 1: No Speedup with CFG-Parallel
**Symptoms**: CFG-Parallel enabled but no performance improvement
**Solutions**:
1. **Ensure CFG scale is set correctly:**
```python
# Bad: No CFG effect
sampling_params = OmniDiffusionSamplingParams(num_inference_steps=50)
# Good: CFG-Parallel will work
sampling_params = OmniDiffusionSamplingParams(
num_inference_steps=50,
true_cfg_scale=4.0 # Must be > 1.0
)
```
2. **Add negative prompt:**
```python
outputs = omni.generate(
{
"prompt": "beautiful landscape",
"negative_prompt": "low quality, blurry", # Required for best results
"multi_modal_data": {"image": input_image}
},
sampling_params
)
```
3. **Check model support:**
- Verify your model in [supported models](../../diffusion_features.md#supported-models)
- Some models don't support CFG-Parallel
---
## Summary
1. ✅ **Enable CFG-Parallel** - Set `cfg_parallel_size=2` in `DiffusionParallelConfig` to get speedup when using CFG
2. ✅ **Set CFG Scale** - Ensure `true_cfg_scale > 1.0` in `OmniDiffusionSamplingParams` for CFG-Parallel to take effect
3. ✅ **Check Model Support** - Verify your model supports CFG-Parallel in [supported models](../../diffusion_features.md#supported-models)
@@ -1,87 +0,0 @@
# Expert Parallelism Guide
## Table of Content
- [Overview](#overview)
- [Quick Start](#quick-start)
- [Configuration Parameters](#configuration-parameters)
- [Best Practices](#best-practices)
- [Summary](#summary)
---
## Overview
Unlike Tensor Parallelism which shards every layer's weights, Expert Parallelism (EP) only shards the MoE expert MLP blocks. This significantly reduces the memory footprint of MoE models (e.g., HunyuanImage3.0) while maintaining constant dense-equivalent compute efficiency.
During the forward pass, a gating mechanism routes tokens to their designated experts, requiring all-to-all communication to dispatch tokens to the correct ranks and combine results.
See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
!!! note "EP Size Constraint"
The effective EP size equals `tp × sp × cfg × dp`. At least one of TP/SP/CFG/DP must be set when EP is enabled.
---
## Quick Start
### Basic Usage
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="tencent/HunyuanImage-3.0",
parallel_config=DiffusionParallelConfig(
tensor_parallel_size=8,
enable_expert_parallel=True,
),
)
outputs = omni.generate(
"A brown and white dog is running on the grass",
OmniDiffusionSamplingParams(
num_inference_steps=50,
width=1024,
height=1024,
),
)
```
---
## Configuration Parameters
In `DiffusionParallelConfig`:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `enable_expert_parallel` | bool | False | Enable Expert Parallelism for MoE models |
EP size is derived automatically as `tp × sp × cfg × dp` — configure at least one of those to set the EP degree.
---
## Best Practices
### When to Use
**Good for:**
- MoE models (e.g., HunyuanImage3.0) with numbers of experts
- Memory-constrained multi-GPU setups where only expert blocks need sharding
**Not for:**
- Dense models (no MoE layers) — EP has no effect
- Single GPU setups
---
## Summary
1. ✅ **Enable EP** - Set `enable_expert_parallel=True` in `DiffusionParallelConfig` for MoE models
2. ✅ **Set parallelism degree** - At least one of `tensor_parallel_size` / `ulysses_degree` / `cfg_parallel_size` must be > 1 to define the EP size
@@ -1,149 +0,0 @@
# HSDP Guide
## Table of Content
- [Overview](#overview)
- [Quick Start](#quick-start)
- [Example Script](#example-script)
- [Configuration Parameters](#configuration-parameters)
- [Best Practices](#best-practices)
- [Summary](#summary)
---
## Overview
HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
Unlike Tensor Parallelism which splits computation, HSDP uses PyTorch's FSDP2 to shard and redistribute weights at runtime. Each GPU only holds a fraction of the model weights, and weights are gathered on-demand during forward passes.
See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
**Operating Modes:**
- **Standalone Mode**: HSDP alone without other parallelism. Must specify `hsdp_shard_size` explicitly.
- **Combined Mode**: HSDP overlays on top of other parallelism (Ulysses-SP, CFG-Parallel). HSDP dimensions must match world_size.
---
## Quick Start
### Basic Usage
Simplest working example (standalone HSDP, shard across 4 GPUs):
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
parallel_config=DiffusionParallelConfig(
use_hsdp=True,
hsdp_shard_size=4, # Shard across 4 GPUs
),
)
outputs = omni.generate(
"A cat playing piano",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
### Combined with Sequence Parallel
```python
omni = Omni(
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
parallel_config=DiffusionParallelConfig(
ulysses_degree=4, # Sequence parallel
use_hsdp=True, # HSDP overlays on SP
),
)
```
---
## Example Script
### Offline Inference
Use Python script under `examples/offline_inference/image_to_video/`:
```bash
# Standalone HSDP: shard across 4 GPUs
python examples/offline_inference/image_to_video/image_to_video.py \
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--use-hsdp \
--hsdp-shard-size 4
# Combined HSDP + Sequence Parallel
python examples/offline_inference/image_to_video/image_to_video.py \
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--ulysses-degree 4 \
--use-hsdp
```
### Online Serving
**Standalone HSDP** (shard model across 4 GPUs):
```bash
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 \
--use-hsdp --hsdp-shard-size 4
```
**Combined with Sequence Parallel**:
```bash
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 \
--use-hsdp --usp 4
```
---
## Configuration Parameters
In `DiffusionParallelConfig`:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `use_hsdp` | bool | False | Enable HSDP |
| `hsdp_shard_size` | int | -1 | Number of GPUs to shard weights across. `-1` = auto (requires other parallelism > 1) |
| `hsdp_replicate_size` | int | 1 | Number of replica groups. Each group holds a full sharded copy |
**Constraints:**
- `hsdp_replicate_size × hsdp_shard_size == world_size`
- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
---
## Best Practices
### When to Use
**Good for:**
- Very large models (e.g., Wan2.2 14B)
- Multi-GPU setups where memory reduction is the primary goal
- Combining with Sequence Parallelism for large video models
**Not for:**
- Models that fit comfortably in single-GPU memory
- Use cases requiring Tensor Parallelism (HSDP and TP are mutually exclusive)
### Adding HSDP Support to New Models
For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../../design/feature/hsdp.md).
---
## Summary
1. ✅ **Enable HSDP** - Set `use_hsdp=True` and `hsdp_shard_size` to reduce per-GPU memory for large models
2. ✅ **Combine with SP** - Use together with `ulysses_degree` for video models requiring both memory reduction and sequence parallelism
3. ⚠️ **Incompatible with TP** - `tensor_parallel_size` must be 1 when HSDP is enabled
@@ -1,16 +0,0 @@
# Parallelism Acceleration Guide
This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion model inference and reducing per-device memory requirements.
## Supported Methods
| Method | Description |
|--------|-------------|
| **[Tensor Parallelism](tensor_parallel.md)** | Shards DiT weights across GPUs to reduce per-GPU memory |
| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation |
| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs |
| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) |
See [Supported Models](../../diffusion_features.md#supported-models) for per-model compatibility.
@@ -1,233 +0,0 @@
# Sequence Parallelism Guide
## Table of Content
- [Overview](#overview)
- [Quick Start](#quick-start)
- [Example Script](#example-script)
- [Best Practices](#best-practices)
- [Troubleshooting](#troubleshooting)
- [Summary](#summary)
---
## Overview
Sequence parallelism splits the input along the sequence dimension across multiple GPUs, allowing each device to process only a portion of the sequence. vLLM-Omni provides 1.5x-3.6x speedup for large images and videos using DeepSpeed Ulysses, Ring-Attention, or hybrid approaches. Use sequence parallelism when generating high-resolution images/videos that don't fit on a single GPU or require faster inference.
See supported models list in [Diffusion Features - Supported Models](../../diffusion_features.md#supported-models).
**Supported Methods:**
- **DeepSpeed Ulysses Sequence Parallel (Ulysses-SP)** ([paper](https://arxiv.org/pdf/2309.14509)): Uses all-to-all communication for subset of attention heads per device
- **Ring-Attention** ([paper](https://arxiv.org/abs/2310.01889)): Uses ring-based P2P communication with sharded sequence dimension throughout
- **Hybrid Ulysses + Ring**: Combines both for larger scale parallelism (`ulysses_degree × ring_degree`)
---
## Quick Start
### Basic Usage - Ulysses-SP
Simplest working example with Ulysses Sequence Parallel:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="Qwen/Qwen-Image",
parallel_config=DiffusionParallelConfig(ulysses_degree=2) # Enable Ulysses-SP
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50, width=1024, height=1024),
)
```
!!! note "Experimental UAA mode"
`ulysses_mode="advanced_uaa"` is an experimental extension to Ulysses-SP. It lets Ulysses attention handle arbitrary sequence lengths and arbitrary attention head counts without relying on `attention_mask`-based token padding.
In hybrid Ulysses + Ring mode, Ring still requires every rank in the same ring group to observe the same post-Ulysses sequence length. If that condition is not met, vLLM-Omni raises a validation error instead of entering the ring kernel with inconsistent shapes.
To enable the experimental UAA mode, use a model/configuration that requires it. For example, `Tongyi-MAI/Z-Image-Turbo` has 30 attention heads, so `ulysses_degree=4` requires UAA because 30 is not divisible by 4:
```python
omni = Omni(
model="Tongyi-MAI/Z-Image-Turbo",
parallel_config=DiffusionParallelConfig(
ulysses_degree=4,
ulysses_mode="advanced_uaa",
),
)
```
### Alternative Methods
**Ring-Attention** (better for very long sequences):
```python
omni = Omni(
model="Qwen/Qwen-Image",
parallel_config=DiffusionParallelConfig(ring_degree=2) # Enable Ring-Attention
)
```
**Hybrid Ulysses + Ring** (for larger scale):
```python
omni = Omni(
model="Qwen/Qwen-Image",
parallel_config=DiffusionParallelConfig(ulysses_degree=2, ring_degree=2) # 4 GPUs total
)
```
---
## Example Script
### Offline Inference
Use Python script under `examples/offline_inference/text_to_image/text_to_image.py`:
**Ulysses-SP:**
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model Qwen/Qwen-Image \
--prompt "A cat sitting on a windowsill" \
--ulysses-degree 2 \
--width 1024 --height 1024
```
**Ring-Attention:**
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model Qwen/Qwen-Image \
--prompt "A cat sitting on a windowsill" \
--ring-degree 2 \
--width 1024 --height 1024
```
**Hybrid Ulysses + Ring:**
```bash
# Hybrid: 2 Ulysses × 2 Ring = 4 GPUs total
python examples/offline_inference/text_to_image/text_to_image.py \
--model Qwen/Qwen-Image \
--prompt "A cat sitting on a windowsill" \
--ulysses-degree 2 --ring-degree 2 \
--width 1024 --height 1024
```
### Online Serving
**Ulysses-SP:**
```bash
# Text-to-image (requires >= 2 GPUs)
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
```
**Ulysses-SP with UAA mode** (for models with non-divisible head counts):
```bash
vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 --usp 4 --ulysses-mode advanced_uaa
```
**Ring-Attention:**
```bash
# Text-to-image (requires >= 2 GPUs)
vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2
```
**Hybrid Ulysses + Ring:**
```bash
# Text-to-image (requires >= 4 GPUs)
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
```
---
## Configuration Parameters
In `DiffusionParallelConfig`:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `ulysses_degree` | int | 1 | Number of GPUs for Ulysses-SP. Uses all-to-all communication. |
| `ring_degree` | int | 1 | Number of GPUs for Ring-Attention. Uses P2P ring communication. |
| `ulysses_mode` | str | `"default"` | Ulysses attention mode. Set to `"advanced_uaa"` to handle arbitrary sequence lengths and head counts without padding. |
**Notes:**
- Total sequence parallel size equals to `ulysses_degree × ring_degree`
- Degrees must evenly divide the sequence length for optimal performance (or use `ulysses_mode="advanced_uaa"` for Ulysses-SP)
## Best Practices
### When to Use
**Good for:**
- Large images (1024x1024 or higher) or videos
- Fast inter-GPU communication, larger bandwidth (e.g., NVLink)
**Not for:**
- Small images (<1024px) - overhead exceeds benefit, use single GPU with cache instead
---
## Troubleshooting
### Common Issue 1: Performance Not Scaling
**Symptoms**: Adding GPUs doesn't improve speed proportionally, or higher parallelism degree is slower
**Diagnosis:**
```bash
# Check GPU topology
nvidia-smi topo -m
```
**Solutions:**
1. Check inter-GPU communication - NVLink is better than PCIe
2. Reduce parallelism degree if over-parallelized:
```python
# If 4 GPUs is slower than 2
parallel_config=DiffusionParallelConfig(ulysses_degree=2)
```
3. Try to switch between Ring-Attention and Ulysses-SP
- Ring-Attention has advantages, like communication-computation overlap, but the block-wise loop overhead is relatively higher, especially for short sequences
- Ulysses-SP: can benefit from larger bandwidth (such as NVLink), with two major constraints, the sequence length should be divisible by usp size, and the number of heads should be divisible by usp size (or use `ulysses_mode="advanced_uaa"`)
### Common Issue 2: Out of Memory (OOM)
**Symptoms**: CUDA OOM errors or process crashes with memory errors
**Solutions:**
1. Increase parallelism degree to split sequence more:
```python
parallel_config=DiffusionParallelConfig(ulysses_degree=4) # From 2
```
2. Combine with other parallelism method, e.g., tensor parallel, and memory optimization methods, e.g., cpu offloading.
## Summary
1. ✅ **Enable Sequence Parallelism** - Set `ulysses_degree` or `ring_degree` for long sequence generation
2. ✅ **UAA mode** - Use `ulysses_mode="advanced_uaa"` when head count is not divisible by `ulysses_degree`
3. ✅ **Troubleshooting** - Check GPU topology with `nvidia-smi topo -m`, reduce degree if performance doesn't scale
@@ -1,151 +0,0 @@
# Tensor Parallelism Guide
## Table of Content
- [Overview](#overview)
- [Quick Start](#quick-start)
- [Example Script](#example-script)
- [Configuration Parameters](#configuration-parameters)
- [Best Practices](#best-practices)
- [Troubleshooting](#troubleshooting)
- [Summary](#summary)
---
## Overview
Tensor Parallelism (TP) shards some model weights across multiple GPUs, usually the Linear layers. This enables running large models that don't fit on a single GPU. It's essential for memory-constrained setups or very large models.
See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
!!! note "TP Limitations for Diffusion Models"
We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
- Good news: The text_encoder typically has minimal impact on overall inference performance.
- Bad news: When TP is enabled, every TP process retains a full copy of the text_encoder weights, leading to significant GPU memory waste.
We are actively refactoring this design to address this. For details and progress, please refer to [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
---
## Quick Start
### Basic Usage
Simplest working example:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="Tongyi-MAI/Z-Image-Turbo",
parallel_config=DiffusionParallelConfig(tensor_parallel_size=2), # Enable TP
)
outputs = omni.generate(
"a cat reading a book",
OmniDiffusionSamplingParams(num_inference_steps=9),
)
```
---
## Example Script
### Offline Inference
Use Python script under `examples/offline_inference`, and enable TP:
```bash
# Text-to-Image with Qwen-Image
python examples/offline_inference/text_to_image/text_to_image.py \
--model Qwen/Qwen-Image \
--tensor-parallel-size 2
# Image Editing with Qwen-Image-Edit
python examples/offline_inference/image_to_image/image_edit.py \
--model Qwen/Qwen-Image-Edit \
--image input.png \
--prompt "Edit description" \
--tensor-parallel-size 2
```
### Online Serving
You can enable tensor parallelism in online serving via `--tensor-parallel-size`:
```bash
# Text-to-Image with Qwen-Image on 2 GPUs
vllm serve Qwen/Qwen-Image --omni --port 8091 \
--tensor-parallel-size 2
# Text-to-Image with Z-Image (TP=2 only)
vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 \
--tensor-parallel-size 2
```
---
## Configuration Parameters
In `DiffusionParallelConfig`:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `tensor_parallel_size` | int | 1 | Number of GPUs to shard model weights across. Must divide number of heads. |
---
## Best Practices
### When to Use
**Good for:**
- Large models that don't fit on a single GPU, especially for models with large DiT blocks (transformer layers)
- Memory-constrained environments
**Not for:**
- When maximum throughput is needed and memory is sufficient
- Models with incompatible dimensions (e.g., Z-Image `num_heads=30`, which now supports `tensor_parallel_size=2`)
## Troubleshooting
### Common Issue 1: Out of Memory (OOM)
**Symptoms**: CUDA OOM errors during model loading or inference, process crashes with memory errors
**Solution**:
```python
# Step 1: Enable TP with smallest degree
parallel_config=DiffusionParallelConfig(tensor_parallel_size=2)
# Step 2: If still OOM, increase TP degree
parallel_config=DiffusionParallelConfig(tensor_parallel_size=4)
```
### Common Issue 2: Divisibility Error
**Symptoms**: Error like "Model dimension X not divisible by tensor_parallel_size Y"
**Solutions**:
1. Check model-specific constraints (e.g., Z-Image only supports TP=2)
2. Use a smaller TP size that divides model dimensions
3. Consult [Supported Models](../../diffusion_features.md#supported-models) for compatible TP sizes
---
## Summary
1. ✅ **Enable TP** - Set `--tensor-parallel-size` to reduce per-GPU memory
2. ✅ **Increase TP size** - Only increase if OOM persists
3. ⚠️ **Text encoder not sharded** - Known limitation
@@ -1,200 +0,0 @@
# VAE Patch Parallelism Guide
## Table of Content
- [Overview](#overview)
- [Quick Start](#quick-start)
- [Example Script](#example-script)
- [Configuration Parameters](#configuration-parameters)
- [Best Practices](#best-practices)
- [Troubleshooting](#troubleshooting)
- [Summary](#summary)
---
## Overview
VAE Patch Parallelism distributes the VAE (Variational AutoEncoder) decode/encode computation across multiple GPUs by splitting the latent space into spatial tiles or patches. Each GPU processes a subset of tiles in parallel, significantly reducing peak memory consumption during the VAE decode stage while maintaining output quality.
This is particularly useful for:
- **High-resolution image generation** where VAE decode can become a memory bottleneck
- **Memory-constrained environments** where the VAE decode activation peak exceeds available VRAM
- **Multi-GPU setups** where you want to leverage distributed resources for the VAE stage
See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
VAE Patch Parallelism uses two strategies based on image size:
| Strategy | Use Case | How It Works | Overlap Handling | Output Quality |
|----------|----------|--------------|------------------|----------------|
| **Tiled Decode** | Large images (triggers VAE tiling) | Distributes existing VAE tiling computation across ranks. Each rank decodes a subset of overlapping tiles. | Uses VAE's native `blend_v` and `blend_h` functions to seamlessly merge overlapping regions | Bit-identical (same logic as single-GPU tiling) |
| **Patch Decode** | Small images (no VAE tiling) | Splits latent into spatial patches with halos. Each rank decodes one patch with boundary context. | Halo regions provide edge context; core regions are directly stitched without blending | Near-identical (diff < 0.5%, visually imperceptible) |
VAE Patch Parallelism **reuses the DiT process group** (`dit_group`) and does not initialize a separate ProcessGroup. This means:
- **Shared ranks**: VAE patch parallelism uses the same GPU ranks as DiT parallelism (Tensor Parallel, Sequence Parallel, etc.)
- **Combined usage**: VAE patch parallelism is typically used together with other parallelism methods
- **Configuration alignment**: The `vae_patch_parallel_size` should be no greater than the size of your DiT process group
---
## Quick Start
### Basic Usage
Simplest working example:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
# TP=2 for DiT, VAE patch parallel also uses these 2 GPUs
omni = Omni(
model="Tongyi-MAI/Z-Image-Turbo",
parallel_config=DiffusionParallelConfig(
tensor_parallel_size=2, # Enable tensor parallelism for DiT
vae_patch_parallel_size=2, # Enable VAE patch parallelism
),
vae_use_tiling=True, # Required for VAE patch parallelism
)
outputs = omni.generate(
"a futuristic city at sunset, high resolution, 8k",
OmniDiffusionSamplingParams(
num_inference_steps=9,
height=1152, # High resolution benefits from VAE patch parallel
width=1152,
),
)
```
---
## Example Script
### Offline Inference
Use Python script under `examples/offline_inference/text_to_image/`:
```bash
# Text-to-Image with Z-Image
python examples/offline_inference/text_to_image/text_to_image.py \
--model Tongyi-MAI/Z-Image-Turbo \
--prompt "a futuristic city at sunset" \
--height 1152 \
--width 1152 \
--tensor-parallel-size 2 \
--vae-patch-parallel-size 2 \
--vae-use-tiling
```
### Online Serving
You can enable VAE patch parallelism in online serving via `--vae-patch-parallel-size`:
```bash
# Text-to-Image with Z-Image, TP=2 + VAE patch parallel=2
vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 \
--tensor-parallel-size 2 \
--vae-patch-parallel-size 2 \
--vae-use-tiling
```
---
## Configuration Parameters
In `DiffusionParallelConfig`:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. |
Additional requirements:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `vae_use_tiling` | bool | False | Must be set to `True` when using VAE patch parallelism. |
!!! note "Automatic VAE Tiling"
When `vae_patch_parallel_size > 1` and the model has a distributed VAE (`DistributedVaeMixin`), the system automatically sets `vae_use_tiling=True` if not already enabled.
---
## Best Practices
### When to Use
**Good for:**
- High-resolution image generation and long video generation
- Memory-constrained setups where VAE decode causes OOM
- Multi-GPU environments
**Not for:**
- Low-resolution images/videos where VAE decode is not a bottleneck
- Single GPU setups should use vae tiling decode, but not parallel vae tiling decode
- Models that do not support vae patch parallel
---
## Troubleshooting
### Common Issue 1: Model Not Support VAE Patch Parallel
**Symptoms**:
```
WARNING: vae_patch_parallel_size=2 is set but VAE patch parallelism is NOT enabled for xxxPipeline; ignoring.
```
**Root Cause**: VAE Patch Parallelism requires the model's VAE to implement `DistributedVaeMixin`. At startup, `vllm_omni/diffusion/registry.py` checks whether the instantiated pipeline has a `.vae` attribute that is an instance of `DistributedVaeMixin`. If it does not, the setting is silently ignored:
```python
vae_pp_size = od_config.parallel_config.vae_patch_parallel_size
is_distributed_vae = hasattr(model, "vae") and isinstance(model.vae, DistributedVaeMixin)
if vae_pp_size > 1 and not is_distributed_vae:
logger.warning(
"vae_patch_parallel_size=%d is set but VAE patch parallelism is NOT enabled for %s; ignoring.",
vae_pp_size,
od_config.model_class_name,
)
```
**Solutions**:
1. **Use a supported model** (recommended): check [Supported Models](../../diffusion_features.md#supported-models) for the VAE-Patch-Parallel column.
2. To add support for a new model, implement `DistributedVaeMixin` on its VAE class (contributions are welcome).
### Common Issue 2: `vae_patch_parallel_size` Exceeds DiT Process Group Size
**Symptoms**: Shows warning message, and vae patch parallel size is resized to DiT process group size
**Root Cause**: VAE Patch Parallelism reuses the DiT process group.
**Recommendation**: Always set `vae_patch_parallel_size` to be no greater than your DiT process group size.
Note that the size of DiT process group size equals to:
```text
dit_parallel_size = data_parallel_size
× cfg_parallel_size
× sequence_parallel_size
× pipeline_parallel_size
× tensor_parallel_size
```
_sequence_parallel_size = ulysses_degree × ring_degree_
---
## Summary
1. ✅ **Enable VAE Patch Parallelism** - Set `vae_patch_parallel_size`, `vae_use_tiling=True` in `DiffusionParallelConfig` to reduce VAE decode peak memory
2. ✅ **Use Long Sequence** - VAE patch parallelism benefits are most apparent at long sequence decoding
3. ✅ **Combine with other parallelism methods** - Suggest to use together with Tensor Parallel or CFG-Parallel for maximum memory savings
@@ -0,0 +1,448 @@
# Parallelism Acceleration Guide
This guide includes how to use parallelism methods in vLLM-Omni to speed up diffusion model inference as well as reduce the memory requirement on each device.
## Overview
The following parallelism methods are currently supported in vLLM-Omni:
1. DeepSpeed Ulysses Sequence Parallel (DeepSpeed Ulysses-SP) ([arxiv paper](https://arxiv.org/pdf/2309.14509)): Ulysses-SP splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
2. [Ring-Attention](#ring-attention) - splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results, keeping the sequence dimension sharded
3. Classifier-Free-Guidance Parallel (CFG-Parallel): CFG-Parallel runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step.
4. [Tensor Parallelism](#tensor-parallelism): Tensor parallelism shards model weights across devices. This can reduce per-GPU memory usage. Note that for diffusion models we currently shard the majority of layers within the DiT.
5. [VAE Patch Parallelism](#vae-patch-parallelism): VAE patch parallelism shards VAE decode spatially across ranks. This can reduce the peak memory of VAE decode and (depending on resolution and communication overhead) speed up VAE decode.
6. [HSDP](#hsdp): Hybrid Sharded Data Parallel shards model weights across GPUs using PyTorch FSDP2. This reduces per-GPU memory usage, enabling inference of large models on GPUs with limited memory.
7. [Expert Parallel](#expert-parallelism): Expert Parallelism shards the Experts of a Mixture-of-Experts (MoE) layer across multiple devices. During the forward, a gating mechanism routes tokens to their designated experts, necessitating cross-cards communication(all-to-all) to dispatch tokens to the correct ranks and combine the results. This parallelism allows for massive scaling of model parameters without a proportional increase in the computational load per device.
The following table shows which models are currently supported by parallelism method:
### ImageGen
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel | HSDP |
|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|:----:|
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ❌ |
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ❌ |
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | N/A | ❌ |
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | N/A | ❌ |
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ |
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ |
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ |
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | N/A | ❌ |
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ | N/A | ❌ |
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ✅ |
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | N/A | ✅ |
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | ✅ |
| **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ |
| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ✅ | ❌ | ❌ | N/A | ❌ |
!!! note "TP Limitations for Diffusion Models"
We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
- Good news: The text_encoder typically has minimal impact on overall inference performance.
- Bad news: When TP is enabled, every TP process retains a full copy of the text_encoder weights, leading to significant GPU memory waste.
We are actively refactoring this design to address this. For details and progress, please refer to [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
!!! note "Why Z-Image is TP=2 only"
Z-Image Turbo is currently limited to `tensor_parallel_size` of **1 or 2** due to model shape divisibility constraints.
For example, the model has `n_heads=30` and a final projection out dimension of `64`, so valid TP sizes must divide both 30 and 64; the only common divisors are **1 and 2**.
### VideoGen
| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | VAE-Patch-Parallel |
|-------|------------------|:----------:|:--------------:|:---------------:|:----:| :----:|
| **Wan2.1** | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
| **Wan2.1** | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
| **LTX-2** | `Lightricks/LTX-2` | ✅ | ✅ | ✅ | ❌ | ❌ |
### Tensor Parallelism
Tensor parallelism splits model parameters across GPUs. In vLLM-Omni, tensor parallelism is configured via `DiffusionParallelConfig.tensor_parallel_size`.
#### Offline Inference
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="Tongyi-MAI/Z-Image-Turbo",
parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
)
outputs = omni.generate(
"a cat reading a book",
OmniDiffusionSamplingParams(
num_inference_steps=9,
width=512,
height=512,
),
)
```
### VAE Patch Parallelism
VAE patch parallelism distributes the VAE decode workload across multiple ranks by splitting the latent spatially. It is configured via `DiffusionParallelConfig.vae_patch_parallel_size` and can be combined with other parallelism methods (e.g., TP).
!!! note "Enablement and feature gate"
- VAE patch parallelism is currently **enabled only for validated pipelines** (check [ImageGen](#imagegen) and [VideoGen](#videogen) for more information).
- If `vae_patch_parallel_size > 1` is set for a validated pipeline, vLLM-Omni will automatically enable `vae_use_tiling` as a safety gate. (We use `vae_use_tiling` because it indicates the VAE supports diffusers tiling parameters like `tile_latent_min_size` and `tile_overlap_factor`.)
#### Offline Inference
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="Tongyi-MAI/Z-Image-Turbo",
parallel_config=DiffusionParallelConfig(
tensor_parallel_size=2,
vae_patch_parallel_size=2,
),
vae_use_tiling=True,
)
outputs = omni.generate(
prompt="a cat reading a book",
num_inference_steps=9,
width=1024,
height=1024,
)
```
#### How it works (method selection)
VAE patch parallelism automatically selects between two internal decode methods based on whether diffusers tiling would kick in:
- `_distributed_tiled_decode`: Used when the latent spatial size exceeds `vae.tile_latent_min_size` (i.e., diffusers tiled decode). Each rank decodes a subset of tiles; rank0 gathers and runs the same overlap+blend+stitch logic as diffusers. This matches the single-rank diffusers tiled output.
- `_distributed_patch_decode`: Used when diffusers tiling would not kick in. Each rank decodes a grid patch expanded with a latent-space halo; then rank0 gathers the cropped core patches and stitches them into the full image. This path has no blending and can introduce small numerical differences compared to the non-parallel decode.
### Sequence Parallelism
#### Ulysses-SP
##### Offline Inference
An example of offline inference script using [Ulysses-SP](https://arxiv.org/pdf/2309.14509) is shown below:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ulysses_degree = 2
omni = Omni(
model="Qwen/Qwen-Image",
parallel_config=DiffusionParallelConfig(ulysses_degree=2)
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)
```
See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
##### Online Serving
You can enable Ulysses-SP in online serving for diffusion models via `--usp`:
```bash
# Text-to-image (requires >= 2 GPUs)
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
```
##### Benchmarks
!!! note "Benchmark Disclaimer"
These benchmarks are provided for **general reference only**. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:
- Specific model and use case
- Hardware configuration
- Careful parameter tuning
- Different inference settings (e.g., number of steps, image resolution)
To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** model generating images (**2048x2048** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA H800 GPUs. `sdpa` is the attention backends.
| Configuration | Ulysses degree |Generation Time | Speedup |
|---------------|----------------|---------|---------|
| **Baseline (diffusers)** | - | 112.5s | 1.0x |
| Ulysses-SP | 2 | 65.2s | 1.73x |
| Ulysses-SP | 4 | 39.6s | 2.84x |
| Ulysses-SP | 8 | 30.8s | 3.65x |
#### Ring-Attention
Ring-Attention ([arxiv paper](https://arxiv.org/abs/2310.01889)) splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results. Unlike Ulysses-SP which uses all-to-all communication, Ring-Attention keeps the sequence dimension sharded throughout the computation and circulates Key/Value blocks through a ring topology.
##### Offline Inference
An example of offline inference script using Ring-Attention is shown below:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ring_degree = 2
omni = Omni(
model="Qwen/Qwen-Image",
parallel_config=DiffusionParallelConfig(ring_degree=2)
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)
```
See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
##### Online Serving
You can enable Ring-Attention in online serving for diffusion models via `--ring`:
```bash
# Text-to-image (requires >= 2 GPUs)
vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2
```
##### Benchmarks
!!! note "Benchmark Disclaimer"
These benchmarks are provided for **general reference only**. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:
- Specific model and use case
- Hardware configuration
- Careful parameter tuning
- Different inference settings (e.g., number of steps, image resolution)
To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** model generating images (**1024x1024** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA A100 GPUs. `flash_attn` is the attention backends.
| Configuration | Ring degree |Generation Time | Speedup |
|---------------|----------------|---------|---------|
| **Baseline (diffusers)** | - | 45.2s | 1.0x |
| Ring-Attention | 2 | 29.9s | 1.51x |
| Ring-Attention | 4 | 23.3s | 1.94x |
#### Hybrid Ulysses + Ring
You can combine both Ulysses-SP and Ring-Attention for larger scale parallelism. The total sequence parallel size equals `ulysses_degree × ring_degree`.
##### Offline Inference
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
# Hybrid: 2 Ulysses × 2 Ring = 4 GPUs total
omni = Omni(
model="Qwen/Qwen-Image",
parallel_config=DiffusionParallelConfig(ulysses_degree=2, ring_degree=2)
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)
```
##### Online Serving
```bash
# Text-to-image (requires >= 4 GPUs)
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
```
##### Benchmarks
!!! note "Benchmark Disclaimer"
These benchmarks are provided for **general reference only**. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:
- Specific model and use case
- Hardware configuration
- Careful parameter tuning
- Different inference settings (e.g., number of steps, image resolution)
To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** model generating images (**1024x1024** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA A100 GPUs. `flash_attn` is the attention backends.
| Configuration | Ulysses degree | Ring degree | Generation Time | Speedup |
|---------------|----------------|-------------|-----------------|---------|
| **Baseline (diffusers)** | - | - | 45.2s | 1.0x |
| Hybrid Ulysses + Ring | 2 | 2 | 24.3s | 1.87x |
### CFG-Parallel
#### Offline Inference
CFG-Parallel is enabled through `DiffusionParallelConfig(cfg_parallel_size=2)`, which runs one rank for the positive branch and one rank for the negative branch.
An example of offline inference using CFG-Parallel (image-to-image) is shown below:
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
image_path = "path_to_image.png"
omni = Omni(
model="Qwen/Qwen-Image-Edit",
parallel_config=DiffusionParallelConfig(cfg_parallel_size=2),
)
input_image = Image.open(image_path).convert("RGB")
outputs = omni.generate(
{
"prompt": "turn this cat to a dog",
"negative_prompt": "low quality, blurry",
"multi_modal_data": {"image": input_image},
},
OmniDiffusionSamplingParams(
true_cfg_scale=4.0,
num_inference_steps=50,
),
)
```
Notes:
- CFG-Parallel is only effective when a `negative_prompt` is provided AND a guidance scale (or `cfg_scale`) is greater than 1.
See `examples/offline_inference/image_to_image/image_edit.py` for a complete working example.
```bash
cd examples/offline_inference/image_to_image/
python image_edit.py \
--model "Qwen/Qwen-Image-Edit" \
--image "qwen_image_output.png" \
--prompt "turn this cat to a dog" \
--negative-prompt "low quality, blurry" \
--cfg-scale 4.0 \
--output "edited_image.png" \
--cfg-parallel-size 2
```
#### Online Serving
You can enable CFG-Parallel in online serving for diffusion models via `--cfg-parallel-size`:
```bash
vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2
```
### HSDP
HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
Unlike Tensor Parallelism which splits computation, HSDP uses PyTorch's FSDP2 to shard and redistribute weights at runtime. Each GPU only holds a fraction of the model weights, and weights are gathered on-demand during forward passes.
#### Configuration
HSDP is configured via `DiffusionParallelConfig`:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `use_hsdp` | bool | False | Enable HSDP |
| `hsdp_shard_size` | int | -1 | Number of GPUs to shard weights across. -1 = auto (requires other parallelism > 1) |
| `hsdp_replicate_size` | int | 1 | Number of replica groups. Each group holds a full sharded copy |
**Constraints:**
- `hsdp_replicate_size × hsdp_shard_size == world_size`
- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
#### Operating Modes
HSDP can work in two modes:
- **Standalone Mode**: HSDP alone without other parallelism. Must specify `hsdp_shard_size` explicitly.
- **Combined Mode**: HSDP overlays on top of other parallelism (Ulysses Sequence Parallel, CFG Parallel). HSDP dimensions must match world_size.
#### Offline Inference
**Standalone HSDP** (shard across 4 GPUs, no other parallelism):
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
parallel_config=DiffusionParallelConfig(
use_hsdp=True,
hsdp_shard_size=4, # Shard across 4 GPUs
),
)
outputs = omni.generate(
"A cat playing piano",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
**Combined HSDP + Sequence Parallel**:
```python
omni = Omni(
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
parallel_config=DiffusionParallelConfig(
ulysses_degree=4, # Sequence parallel
use_hsdp=True, # HSDP overlays on SP
),
)
```
#### Online Serving
**Standalone HSDP** (shard model across 4 GPUs):
```bash
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --hsdp-shard-size 4
```
**Combined with Sequence Parallel**:
```bash
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --usp 4
```
#### Adding HSDP Support to New Models
For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../design/feature/hsdp.md).
### Expert Parallelism
Unlike Tensor Parallelism which shards every layer's weights, EP only shards the MoE expert MLP blocks. This significantly reduces the memory footprint of MoE models (e.g., HunyuanImage3.0) while maintaining constant dense-equivalent compute efficiency. Expert Parallelism is enabled via `DiffusionParallelConfig.enable_expert_parallel`. And `self.ep = tp * sp * cfg * dp` for now, so at least one of TP/SP/CFG/DP should set when EP enabled.
#### Offline Inference
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="tencent/HunyuanImage-3.0",
parallel_config=DiffusionParallelConfig(tensor_parallel_size=8, enable_expert_parallel=True),
)
outputs = omni.generate(
"A brown and white dog is running on the grass",
OmniDiffusionSamplingParams(
num_inference_steps=50,
width=1024,
height=1024,
),
)
```
@@ -63,9 +63,6 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Skip sensitive layers | `img_mlp` |
| Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
| HunyuanImage-3 | `tencent/HunyuanImage3` | All layers | None |
| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None |
| Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | All layers | None |
## Combining with Other Features
+4 -34
View File
@@ -6,13 +6,13 @@
3. Keep user-facing knobs minimal and consistent across offline and online flows.
## Scope
1. Models: Z-Image, Flux2-klein, and Qwen-Image family.
1. Models: Z-Image, and Flux2-klein.
2. Components: diffusion transformer weights, loader paths, and quantization configs.
3. Modes: native GGUF (transformer-only weights).
## Architecture Overview
1. `OmniDiffusionConfig` accepts `quantization` or `quantization_config`.
2. Diffusion quantization wrapper (`DiffusionGGUFConfig`) produces vLLM `QuantizationConfig` objects for linear layers.
2. Diffusion quantization wrapper (`DiffusionGgufConfig`) produces vLLM `QuantizationConfig` objects for linear layers.
3. `DiffusersPipelineLoader` branches on quantization method and loads either HF weights or GGUF weights for the transformer.
4. GGUF transformer loading is routed through model-specific adapters (e.g., Flux2Klein).
5. vLLM GGUF path uses `GGUFConfig` and `GGUFLinearMethod` for matmul.
@@ -124,21 +124,19 @@ x @ weight.T
1. `DiffusersPipelineLoader.load_model` detects `quantization_config.method == "gguf"`.
2. `gguf_model` is resolved as one of: local file, `repo/file.gguf`, or `repo:quant_type`.
3. GGUF weights are routed through adapters in `vllm_omni/diffusion/model_loader/gguf_adapters/`.
4. Name mapping is applied per-architecture (Z-Image, Flux2Klein, Qwen-Image).
4. Name mapping is applied per-architecture (Z-Image, Flux2Klein).
5. GGUF weights are loaded into transformer modules, remaining non-transformer weights come from the HF checkpoint.
## GGUF Adapter Design
1. `GGUFAdapter` is an abstract base class for model-specific adapters.
2. `Flux2KleinGGUFAdapter` implements Flux2-Klein remapping + qkv split + adaLN swap.
3. `ZImageGGUFAdapter` implements Z-Image qkv + ffn shard handling and linear qweight routing.
4. `QwenImageGGUFAdapter` implements Qwen-Image fused attention projection remapping.
5. `get_gguf_adapter(...)` strictly selects by model class/config; unsupported models raise an error (no fallback adapter).
4. `get_gguf_adapter(...)` strictly selects by model class/config; unsupported models raise an error (no fallback adapter).
Adapter paths:
- Base: `vllm_omni/diffusion/model_loader/gguf_adapters/base.py`
- Z-Image: `vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py`
- Flux2-Klein: `vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py`
- Qwen-Image: `vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py`
## User Usage (Offline)
@@ -197,31 +195,3 @@ curl -X POST http://localhost:8000/v1/images/generations \
"num_inference_steps": 4
}'
```
## Qwen-Image Example
### Offline
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model Qwen/Qwen-Image \
--gguf-model QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf \
--quantization gguf \
--prompt "a red paper kite hanging from a pine tree in a winter courtyard" \
--height 1024 \
--width 1024 \
--seed 42 \
--num_inference_steps 20 \
--output outputs/qwen_image_q4km.png
```
### Online
```bash
vllm serve Qwen/Qwen-Image \
--omni \
--port 8000 \
--quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
```
Notes:
1. The GGUF repo provides transformer weights only; the base model repo still provides tokenizer, text encoder, scheduler, and VAE.
2. The same adapter path covers `QwenImagePipeline`, `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, and `QwenImageLayeredPipeline`.
@@ -1,86 +1,22 @@
# Quantization
# Quantization for Diffusion Transformers
vLLM-Omni provides a unified quantization framework that supports both diffusion models (DiT, text encoders, VAEs) and multi-stage omni models (e.g., Qwen3-Omni thinker). Quantization reduces memory usage and accelerates inference.
vLLM-Omni supports quantization of DiT linear layers to reduce memory usage and accelerate inference.
## Supported Methods
### Diffusion Models
| Method | Guide | Description | Tested Models | Min GPU |
|--------|-------|-------------|---------------|---------|
| FP8 | [FP8](fp8.md) | FP8 W8A8, dynamic or static | Z-Image, Qwen-Image, Flux, Bagel | SM 89 (Ada) |
| Int8 | [Int8](int8.md) | Int8 W8A8 | Z-Image, Qwen-Image | SM 89 (Ada) / Ascend NPU |
| GGUF | [GGUF](gguf.md) | GGUF format, dequant+GEMM for N-D tensors | Z-Image, Flux | SM 60 |
### Multi-stage Omni Models (Pre-quantized Checkpoints)
| Method | Description | Tested Models | Min GPU |
|--------|-------------|---------------|---------|
| ModelOpt FP8 | Pre-quantized FP8 via NVIDIA ModelOpt | Qwen3-Omni (thinker) | SM 89 (Ada/Hopper) |
| ModelOpt NVFP4 | Pre-quantized NVFP4 via NVIDIA ModelOpt | Qwen3-Omni (experimental, quality issues) | SM 100 (Blackwell) |
!!! note
AWQ, GPTQ, and BitsAndBytes are supported by vLLM's upstream quantization registry but have **not been tested** with vLLM-Omni pipelines. They may work via `build_quant_config()` but are not validated.
### Pre-quantized LLM Checkpoints (Multi-stage Models)
For multi-stage models like Qwen3-Omni, the unified quantization framework auto-detects
pre-quantized checkpoints via `quantization_config` in the HF config. Quantization is
automatically scoped to the thinker's `language_model` — audio encoder, vision encoder,
talker, and code2wav remain in BF16.
| Format | `quant_algo` | Hardware | Example | Status |
|--------|-------------|----------|---------|--------|
| ModelOpt FP8 | `FP8` | Ada/Hopper (SM 89+) | `asdazd/Qwen3-Omni-30B-A3B-Instruct_modelopt_FP8` | Tested — 47% memory reduction, comparable throughput |
| ModelOpt NVFP4 | `NVFP4` | Blackwell (SM 100+) | `shunyang90/Qwen3-Omni-30B-A3B-Instruct-NVFP4` | Experimental — loads and runs but output quality is unacceptable |
**Tested FP8 results (Qwen3-Omni, 1×H200):**
| Config | Model Memory (GiB) | Mem Reduction | Decode (tok/s) |
|--------|-------------------|---------------|----------------|
| BF16 baseline | 59.26 | — | 41.6 |
| FP8 (ModelOpt) | 31.41 | 47% | 39.9 |
FP8 enables the full 3-stage pipeline (thinker + talker + code2wav) on a single 64GB GPU,
which is impossible with BF16 (thinker alone requires 59.26 GiB).
## Quantization Scope
### Diffusion Models
When `--quantization fp8` is enabled for diffusion models:
| Component | What Gets Quantized | Mechanism |
|-----------|-------------------|-----------|
| **DiT (transformer)** | `nn.Linear` layers | vLLM W8A8 FP8 compute (Ada/Hopper) or weight-only (older GPUs) |
| **Text encoder** | `nn.Linear` layers | FP8 weight storage, BF16 compute |
| **VAE** | `nn.Conv2d`, `nn.Conv3d` layers | FP8 weight storage, BF16 compute |
### Multi-stage Omni Models
For pre-quantized omni model checkpoints, quantization is scoped per component:
| Component | Quantized? | Notes |
|-----------|-----------|-------|
| **Thinker (language_model)** | Yes | FP8 or NVFP4 via ModelOpt |
| **Audio encoder** | No | Stays BF16 |
| **Vision encoder** | No | Stays BF16 |
| **Talker** | No | Stays BF16 |
| **Code2Wav** | No | Stays BF16 |
!!! note
Not all models support all components. See the [FP8 supported models table](fp8.md#supported-models) for per-model details.
| Method | Guide |
|--------|-------|
| FP8 | [FP8](fp8.md) |
| Int8 | [Int8](int8.md) |
| GGUF | [GGUF](gguf.md) |
## Device Compatibility for FP8
| Device | Example Hardware | FP8 | NVFP4 |
|--------|-----------------|-----|-------|
| Blackwell GPU (SM 100+) | RTX 5090, B100, B200 | Yes | Yes (native FP4 HW) |
| Ada/Hopper GPU (SM 89+) | RTX 4090, H100, H200 | Yes (W8A8 native) | No |
| Turing/Ampere GPU (SM 75-86) | RTX 3090, A100 | Yes (weight-only Marlin) | No |
| Ascend NPU | Atlas 800T A2 (910B) | Not yet supported | No |
| GPU Generation | Example GPUs | FP8 Mode |
|---------------|-------------------|----------|
| Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
Kernel selection is automatic on CUDA GPUs.
Kernel selection is automatic.
## Device Compatibility for Int8
@@ -88,62 +24,3 @@ Kernel selection is automatic on CUDA GPUs.
|-------------|---------------|-------------------|----------|
| NVIDIA GPU | Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
| Ascend NPU | Atlas A2/Atlas A3 | Atlas 800T A2/Atlas 900 A3 | Full W8A8 with native hardware |
## Python API
The `build_quant_config()` factory accepts multiple input formats:
```python
from vllm_omni.quantization import build_quant_config
# String
config = build_quant_config("fp8")
# Dict with parameters
config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
# Per-component dict
config = build_quant_config({
"transformer": {"method": "fp8"},
"vae": None,
})
# None / "none"
config = build_quant_config(None) # returns None
```
## Architecture
```
build_quant_config(spec)
├── str → _build_single(method) → vLLM registry / _OVERRIDES
├── dict with "method" → _build_single(method, **kwargs)
├── per-component dict → ComponentQuantizationConfig
│ routes get_quant_method() by longest-prefix match
├── QuantizationConfig → passthrough
└── None / "none" → None
```
## Migration Guide
### Before (v0.14.x)
```python
# Old diffusion-specific API (removed)
from vllm_omni.diffusion.quantization import get_diffusion_quant_config
config = get_diffusion_quant_config("fp8", activation_scheme="static")
```
### After (v0.16.0+)
```python
# Unified API — delegates to vLLM's registry
from vllm_omni.quantization import build_quant_config
config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
# Per-component (new)
config = build_quant_config({
"transformer": {"method": "fp8"},
"vae": None,
})
```
+1 -1
View File
@@ -46,7 +46,7 @@ its stepwise request state machine. For normal diffusion inference, leave it
disabled unless your workflow depends on this mode.
If you are looking for general diffusion speedups, see
[Diffusion Features Overview](../diffusion_features.md).
[Diffusion Acceleration Overview](../diffusion_acceleration.md).
## Troubleshooting
+145
View File
@@ -0,0 +1,145 @@
# TeaCache Configuration Guide
TeaCache speeds up diffusion model inference by caching transformer computations when consecutive timesteps are similar. This typically provides **1.5x-2.0x speedup** with minimal quality loss.
## Quick Start
Enable TeaCache by setting `cache_backend` to `"tea_cache"`:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# Simple configuration - model_type is automatically extracted from pipeline.__class__.__name__
omni = Omni(
model="Qwen/Qwen-Image",
cache_backend="tea_cache",
cache_config={
"rel_l1_thresh": 0.2 # Optional, defaults to 0.2
}
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(
num_inference_steps=50,
),
)
```
### Using Environment Variable
You can also enable TeaCache via environment variable:
```bash
export DIFFUSION_CACHE_BACKEND=tea_cache
```
Then initialize without explicitly setting `cache_backend`:
```python
from vllm_omni import Omni
omni = Omni(
model="Qwen/Qwen-Image",
cache_config={"rel_l1_thresh": 0.2} # Optional
)
```
## Online Serving (OpenAI-Compatible)
Enable TeaCache for online serving by passing `--cache-backend tea_cache` when starting the server:
```bash
vllm serve Qwen/Qwen-Image --omni --port 8091 \
--cache-backend tea_cache \
--cache-config '{"rel_l1_thresh": 0.2}'
```
## Configuration Parameters
### `rel_l1_thresh` (float, default: `0.2`)
Controls the balance between speed and quality. Lower values prioritize quality, higher values prioritize speed.
**Recommended values:**
- `0.2` - **~1.5x speedup** with minimal quality loss (recommended)
- `0.4` - **~1.8x speedup** with slight quality loss
- `0.6` - **~2.0x speedup** with noticeable quality loss
- `0.8` - **~2.25x speedup** with significant quality loss
## Examples
### Python API
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="Qwen/Qwen-Image",
cache_backend="tea_cache",
cache_config={"rel_l1_thresh": 0.2}
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(
num_inference_steps=50,
),
)
```
## Performance Tuning
Start with the default `rel_l1_thresh=0.2` and adjust based on your needs:
- **Maximum quality**: Use `0.1-0.2`
- **Balanced**: Use `0.2-0.4` (recommended)
- **Maximum speed**: Use `0.6-0.8` (may reduce quality)
## Troubleshooting
### Quality Degradation
If you notice quality issues, lower the threshold:
```python
cache_config={"rel_l1_thresh": 0.1} # More conservative caching
```
## Supported Models
### ImageGen
<style>
th {
white-space: nowrap;
min-width: 0 !important;
}
</style>
| Architecture | Models | Example HF Models |
|--------------|--------|-------------------|
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |
### VideoGen
No VideoGen models are supported by TeaCache yet.
### Coming Soon
<style>
th {
white-space: nowrap;
min-width: 0 !important;
}
</style>
| Architecture | Models | Example HF Models |
|--------------|--------|-------------------|
| `FluxPipeline` | Flux | - |
| `CogVideoXPipeline` | CogVideoX | - |
+372
View File
@@ -0,0 +1,372 @@
# Diffusion Acceleration Overview
vLLM-Omni supports various acceleration methods to speed up diffusion model inference with minimal quality degradation. These include **cache methods** that intelligently cache intermediate computations to avoid redundant work across diffusion timesteps, **parallelism methods** that distribute the computation across multiple devices, and **quantization methods** that reduce memory footprint while preserving accuracy.
## Supported Acceleration Methods
vLLM-Omni currently supports two main cache acceleration backends:
1. **[TeaCache](diffusion/teacache.md)** - Hook-based adaptive caching that caches transformer computations when consecutive timesteps are similar
2. **[Cache-DiT](diffusion/cache_dit_acceleration.md)** - Library-based acceleration using multiple techniques:
- **DBCache** (Dual Block Cache): Caches intermediate transformer block outputs based on residual differences
- **TaylorSeer**: Uses Taylor expansion-based forecasting for faster inference
- **SCM** (Step Computation Masking): Selectively computes steps based on adaptive masking
Both methods can provide significant speedups (typically **1.5x-2.0x**) while maintaining high output quality.
vLLM-Omni also supports quantization methods:
3. **[Quantization](diffusion/quantization/overview.md)** - Reduces DiT linear layers from BF16 to FP8 or Int8, providing ~1.28x speedup with minimal quality loss. Supports per-layer skip for sensitive layers.
vLLM-Omni also supports parallelism methods for diffusion models, including:
1. [Ulysses-SP](diffusion/parallelism_acceleration.md#ulysses-sp) - splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
2. [Ring-Attention](diffusion/parallelism_acceleration.md#ring-attention) - splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results, keeping the sequence dimension sharded.
3. [CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel) - runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step.
4. [Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism) - shards DiT weights across devices to reduce per-GPU memory usage.
5. [VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism) - shards VAE decode/encode spatially across ranks to reduce VAE peak memory (and can speed up VAE decode).
6. [HSDP](diffusion/parallelism_acceleration.md#hsdp) - Hybrid Sharded Data Parallel shards model weights across GPUs to reduce per-GPU memory usage, enabling inference of large models on limited GPU memory.
## Quick Comparison
### Cache Methods
| Method | Configuration | Description | Best For |
|--------|--------------|-------------|----------|
| **TeaCache** | `cache_backend="tea_cache"` | Simple, adaptive caching with minimal configuration | Quick setup, balanced speed/quality |
| **Cache-DiT** | `cache_backend="cache_dit"` | Advanced caching with multiple techniques (DBCache, TaylorSeer, SCM) | Maximum acceleration, fine-grained control |
### Quantization Methods
| Method | Configuration | Description | Best For |
|--------|--------------|-------------|----------|
| **FP8** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup |
| **Int8** | `quantization="int8"` | Int8 W8A8 | Memory reduction, inference speedup |
## Supported Models
The following table shows which models are currently supported by each acceleration method:
### ImageGen
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel |
|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:---------------:|:------------------:|
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ (TP=2 only) | ✅ |
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
### VideoGen
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | VAE-Patch-Parallel |
|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
| **Wan2.1-T2V** | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **Wan2.1-T2V** | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ |
### Quantization
| Model | Model Identifier | FP8 | Int8 |
|-------|------------------|:---:|:---:|
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ |
| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ |
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ |
## Performance Benchmarks
The following benchmarks were measured on **Qwen/Qwen-Image** and **Qwen/Qwen-Image-Edit** models generating 1024x1024 images with 50 inference steps:
!!! note "Benchmark Disclaimer"
These benchmarks are provided for **general reference only**. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:
- Specific model and use case
- Hardware configuration
- Careful parameter tuning
- Different inference settings (e.g., number of steps, image resolution)
For optimal performance in your specific scenario, we recommend experimenting with different parameter configurations as described in the detailed guides below.
| Model | Cache Backend | Cache Config | Generation Time | Speedup | Notes |
|-------|---------------|--------------|----------------|---------|-------|
| **Qwen/Qwen-Image** | None | None | 20.0s | 1.0x | Baseline (diffusers) |
| **Qwen/Qwen-Image** | TeaCache | `rel_l1_thresh=0.2` | 10.47s | **1.91x** | Recommended default setting |
| **Qwen/Qwen-Image** | Cache-DiT | DBCache + TaylorSeer (Fn=1, Bn=0, W=8, TaylorSeer order=1) | 10.8s | **1.85x** | - |
| **Qwen/Qwen-Image** | Cache-DiT | DBCache + TaylorSeer + SCM (Fn=8, Bn=0, W=4, TaylorSeer order=1, SCM fast) | 14.0s | **1.43x** | - |
| **Qwen/Qwen-Image-Edit** | None | No acceleration | 51.5s | 1.0x | Baseline (diffusers) |
| **Qwen/Qwen-Image-Edit** | Cache-DiT | Default (Fn=1, Bn=0, W=4, TaylorSeer disabled, SCM disabled) | 21.6s | **2.38x** | - |
To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** model generating images (**2048x2048** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA H800 GPUs. `sdpa` is the attention backends.
| Configuration | Ulysses degree |Generation Time | Speedup |
|---------------|----------------|---------|---------|
| **Baseline (diffusers)** | - | 112.5s | 1.0x |
| Ulysses-SP | 2 | 65.2s | 1.73x |
| Ulysses-SP | 4 | 39.6s | 2.84x |
| Ulysses-SP | 8 | 30.8s | 3.65x |
## Quick Start
### Using TeaCache
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="Qwen/Qwen-Image",
cache_backend="tea_cache",
cache_config={"rel_l1_thresh": 0.2} # Optional, defaults to 0.2
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(
num_inference_steps=50,
),
)
```
### Using Cache-DiT
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="Qwen/Qwen-Image",
cache_backend="cache_dit",
cache_config={
"Fn_compute_blocks": 1,
"Bn_compute_blocks": 0,
"max_warmup_steps": 8,
"enable_taylorseer": True,
"taylorseer_order": 1,
}
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(
num_inference_steps=50,
),
)
```
### Using Ulysses-SP
Run text-to-image:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ulysses_degree = 2
omni = Omni(
model="Qwen/Qwen-Image",
parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)
```
Run image-to-image:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ulysses_degree = 2
omni = Omni(
model="Qwen/Qwen-Image-Edit",
parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
)
outputs = omni.generate(
{
"prompt": "turn this cat to a dog",
"multi_modal_data": {"image": input_image}
},
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
### Using Ring-Attention
Run text-to-image:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ring_degree = 2
omni = Omni(
model="Qwen/Qwen-Image",
parallel_config=DiffusionParallelConfig(ring_degree=2)
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)
```
### Using Tensor Parallelism
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="Tongyi-MAI/Z-Image-Turbo",
parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
)
outputs = omni.generate(
prompt="a cat reading a book",
num_inference_steps=9,
width=512,
height=512,
)
```
### Using VAE Patch Parallelism
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="Tongyi-MAI/Z-Image-Turbo",
parallel_config=DiffusionParallelConfig(
tensor_parallel_size=2,
vae_patch_parallel_size=2,
),
)
outputs = omni.generate(
prompt="a cat reading a book",
num_inference_steps=9,
width=1024,
height=1024,
)
```
### Using HSDP
HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
parallel_config=DiffusionParallelConfig(
use_hsdp=True,
hsdp_replicate_size=1, # Number of replica groups
hsdp_shard_size=8, # Number of GPUs to shard across (or -1 for auto)
),
)
outputs = omni.generate(
"A cat playing piano",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
### Using CFG-Parallel
Run image-to-image:
CFG-Parallel splits the CFG positive/negative branches across GPUs. Use it when you set a non-trivial `true_cfg_scale`.
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
cfg_parallel_size = 2
omni = Omni(
model="Qwen/Qwen-Image-Edit",
parallel_config=DiffusionParallelConfig(cfg_parallel_size=cfg_parallel_size)
)
outputs = omni.generate(
{
"prompt": "turn this cat to a dog",
"multi_modal_data": {"image": input_image}
},
OmniDiffusionSamplingParams(num_inference_steps=50, true_cfg_scale=4.0),
)
```
### Using FP8 Quantization
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="<your-model>",
quantization="fp8",
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
### Using Int8 Quantization
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="<your-model>",
quantization="int8",
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
## Documentation
For detailed information on each acceleration method:
- **[TeaCache Guide](diffusion/teacache.md)** - Complete TeaCache documentation, configuration options, and best practices
- **[Cache-DiT Acceleration Guide](diffusion/cache_dit_acceleration.md)** - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters
- **[Quantization Guide](diffusion/quantization/overview.md)** - Quantization for DiT models with per-layer control
- **[Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism)** - Guidance on how to enable TP for diffusion models.
- **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
- **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.
- **[VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism)** - Guidance on how to reduce VAE memory via patch/tile parallelism.
- **[HSDP](diffusion/parallelism_acceleration.md#hsdp)** - Hybrid Sharded Data Parallel for memory-efficient inference of large models.
-190
View File
@@ -1,190 +0,0 @@
# Diffusion Advanced Features
## Table of Contents
- [Overview](#overview)
- [Supported Features](#supported-features)
- [Supported Models](#supported-models)
- [Feature Compatibility](#feature-compatibility)
- [Learn More](#learn-more)
## Overview
vLLM-Omni supports various advanced features for diffusion models:
- Acceleration: **cache methods**, **parallelism methods**
- Memory optimization: **cpu offloading**, **quantization**
- Extensions: **LoRA inference**
## Supported Features
### Acceleration
#### Lossy Acceleration
Cache methods trade minimal quality for significant speedup. Quality loss is typically imperceptible with proper tuning.
| Method | Description | Best For |
|--------|-------------|----------|
| **[TeaCache](diffusion/cache_acceleration/teacache.md)** | Adaptive caching using modulated inputs | Quick setup, balanced quality/speed on single GPU |
| **[Cache-DiT](diffusion/cache_acceleration/cache_dit.md)** | Multiple caching techniques: DBCache, TaylorSeer, SCM | Fine-grained control, tunable quality-speed tradeoff |
#### Lossless Acceleration
Parallelism methods distribute computation across GPUs without quality loss (mathematically equivalent to single-GPU).
| Method | Description | Best For |
|--------|-------------|----------|
| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via all-to-all communication | High-resolution images (>1536px) or long videos with 2-8 GPUs |
| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via ring-based communication | Videos, very long sequences, memory-constrained, with 2-8 GPUs |
| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)** | Splits CFG positive/negative branches across devices | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs |
| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs |
| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
**Note:** Some acceleration methods can be combined together for optimized performance. See [Feature Compatibility Table](#feature-compatibility) and [Feature Compatibility Tutorial](feature_compatibility.md) for detailed configuration examples.
### Memory Optimization
Memory optimization methods help reduce GPU memory usage, enabling inference on resource-constrained hardware or larger models.
| Method | Description | Best For |
|--------|-------------|----------|
| **[CPU Offload](diffusion/cpu_offload_diffusion.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
| **[Quantization](diffusion/quantization/overview.md)** | Reduces DiT layers from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss |
| **[VAE Patch Parallelism](diffusion/parallelism/vae_patch_parallel.md)** | Distributes VAE decode tiling across GPUs | High-resolution generation with reduced VAE memory peak |
### Extensions
Extension methods add specialized capabilities to diffusion models beyond standard inference.
| Method | Description | Best For |
|--------|-------------|----------|
| **[LoRA Inference](diffusion/lora.md)** | Enables inference with Low-Rank Adaptation (LoRA) adapters weights | Reinforcement learning extensions |
### Quantization Methods
| Method | Configuration | Description | Best For |
|--------|--------------|-------------|----------|
| **[FP8](diffusion/quantization/fp8.md)** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup |
| **[INT8](diffusion/quantization/int8.md)** | `quantization="int8"` | INT8 weight-only, no calibration or pre-quantized checkpoint needed | Memory reduction, broad GPU compatibility |
| **[GGUF](diffusion/quantization/gguf.md)** | `quantization="gguf"` | Native GGUF transformer-only weights (Q4, Q8, etc.) | Memory reduction on consumer GPUs |
## Supported Models
The following tables show which models support each feature:
- **🔀SP (Ulysses & Ring)**: Includes both Ulysses-SP and Ring-Attention methods
- ✅ = Fully supported
- ❌ = Not supported
> Notes:
> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only.
> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](diffusion/quantization/overview.md) for per-method (FP8, GGUF, …) and per-model support details.
### ImageGen
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ |
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ |
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ |
| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ |
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ❌ | ✅ | ✅ |
> Notes:
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
> 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary.
### VideoGen
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
### AudioGen
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
| **Stable-Audio-Open** | ❌ | ❌ | ❓ | ❓ | ❌ | ❌ | ❌ | ❌ | ✅ |
## Feature Compatibility
**Legend:**
- ✅: Functionality is supported
- ❌: No support plan
- ❓: Not verified yet and Not Recommended
| | ⚡TeaCache | ⚡Cache-DiT | 🔀Ulysses-SP | 🔀Ring-Attn | 🔀CFG-Parallel | 🔀Tensor Parallel | 🔀HSDP | 🔀Expert Parallel | 💾CPU Offloading (Layerwise) | 💾CPU Offloading (Module-wise) | 💾VAE Patch Parallel | 💾FP8 Quant | 🔧LoRA Inference |
|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
| **⚡TeaCache** | | | | | | | | | | | | | |
| **⚡Cache-DiT** | ❌ | | | | | | | | | | | | |
| **🔀Ulysses-SP** | ✅ | ✅ | | | | | | | | | | | |
| **🔀Ring-Attn** | ✅ | ✅ | ✅ | | | | | | | | | | |
| **🔀CFG-Parallel** | ✅ | ✅ | ✅ | ✅ | | | | | | | | | |
| **🔀Tensor Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | |
| **🔀HSDP** | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | | | | | | |
| **🔀Expert Parallel** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | | | | | |
| **💾CPU Offloading (Layerwise)** | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | | | | |
| **💾CPU Offloading (Module-wise)** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ❌ | | | | |
| **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | |
| **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | |
| **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | |
!!! info
1. Tensor Parallel and HSDP are not compatible.
2. TeaCache and Cache-DiT are not compatible.
3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
4. CPU Offloading (Layerwise) supports single-card for now.
5. Using FP8-Quant as an example of qunatization methods.
## Learn More
**Cache Acceleration:**
- **[TeaCache Configuration Guide](diffusion/cache_acceleration/teacache.md)** - Parameter tuning, performance tips, troubleshooting
- **[Cache-DiT Advanced Guide](diffusion/cache_acceleration/cache_dit.md)** - DBCache, TaylorSeer, SCM techniques and optimization
**Parallelism Methods:**
- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, HSDP, and Expert Parallelism
**Memory Optimization:**
- **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage
- **[VAE Patch Parallelism Guide](diffusion/parallelism/vae_patch_parallel.md)** - Distribute VAE decode tiling across GPUs for high-resolution images
- **[Quantization Overview](diffusion/quantization/overview.md)** - Overview of quantization methods for diffusion models
**Extensions:**
- **[LoRA Inference Guide](diffusion/lora.md)** - Low-Rank Adaptation for style customization and fine-tuning
**Advanced Topics:**
- **[Feature Compatibility](feature_compatibility.md)** - How to combine multiple features for maximum performance
@@ -79,7 +79,19 @@ Each DAC decoder chunk is logged as it arrives. The final WAV file is written on
## Example materials
??? abstract "diag_codebook.py"
``````py
--8<-- "examples/offline_inference/fish_speech/diag_codebook.py"
``````
??? abstract "diag_reference_compare.py"
``````py
--8<-- "examples/offline_inference/fish_speech/diag_reference_compare.py"
``````
??? abstract "end2end.py"
``````py
--8<-- "examples/offline_inference/fish_speech/end2end.py"
``````
??? abstract "standalone_test.py"
``````py
--8<-- "examples/offline_inference/fish_speech/standalone_test.py"
``````
@@ -43,7 +43,8 @@ GLM-Image uses a 2-stage pipeline:
```bash
python end2end.py \
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--model-path /path/to/glm-image \
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--prompt "A beautiful sunset over the ocean with sailing boats" \
--height 1024 \
--width 1024 \
@@ -54,7 +55,8 @@ python end2end.py \
```bash
python end2end.py \
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--model-path /path/to/glm-image \
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--prompt "Transform this scene into a winter wonderland" \
--image input.png \
--output output_i2i.png
@@ -65,7 +67,7 @@ python end2end.py \
```bash
python end2end.py \
--model-path /path/to/glm-image \
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--prompt "A photorealistic cat sitting on a window sill" \
--height 1024 \
--width 1024 \
@@ -1,203 +0,0 @@
# Helios Video Generation
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/helios>.
Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants:
| Variant | Description | Key Features |
|---------|-------------|--------------|
| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` |
| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support |
| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` |
## Setup
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
## Run Examples
Get into the example folder:
```bash
cd examples/offline_inference/helios
```
### Text-to-Video (T2V)
**Helios-Base** (Stage 1 only):
```bash
python end2end.py \
--model BestWishYsh/Helios-Base \
--sample-type t2v \
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
--guidance-scale 5.0 \
--output helios_t2v_base.mp4
```
**Helios-Mid** (Stage 2 + CFG-Zero*):
```bash
python end2end.py \
--model BestWishYsh/Helios-Mid \
--sample-type t2v \
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
--guidance-scale 5.0 \
--is-enable-stage2 \
--pyramid-num-inference-steps-list 20 20 20 \
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
--output helios_t2v_mid.mp4
```
**Helios-Distilled** (Stage 2 pyramid + DMD):
```bash
python end2end.py \
--model BestWishYsh/Helios-Distilled \
--sample-type t2v \
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
--num-frames 240 \
--guidance-scale 1.0 \
--is-enable-stage2 \
--pyramid-num-inference-steps-list 2 2 2 \
--is-amplify-first-chunk \
--output helios_t2v_distilled.mp4
```
### Image-to-Video (I2V)
**Helios-Base**:
```bash
python end2end.py \
--model BestWishYsh/Helios-Base \
--sample-type i2v \
--image-path /path/to/image.jpg \
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
--guidance-scale 5.0 \
--output helios_i2v_base.mp4
```
**Helios-Mid** (Stage 2 + CFG-Zero*):
```bash
python end2end.py \
--model BestWishYsh/Helios-Mid \
--sample-type i2v \
--image-path /path/to/image.jpg \
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
--guidance-scale 5.0 \
--is-enable-stage2 \
--pyramid-num-inference-steps-list 20 20 20 \
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
--output helios_i2v_mid.mp4
```
**Helios-Distilled**:
```bash
python end2end.py \
--model BestWishYsh/Helios-Distilled \
--sample-type i2v \
--image-path /path/to/image.jpg \
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
--num-frames 240 \
--guidance-scale 1.0 \
--is-enable-stage2 \
--pyramid-num-inference-steps-list 2 2 2 \
--is-amplify-first-chunk \
--output helios_i2v_distilled.mp4
```
### Video-to-Video (V2V)
**Helios-Base**:
```bash
python end2end.py \
--model BestWishYsh/Helios-Base \
--sample-type v2v \
--video-path /path/to/video.mp4 \
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
--guidance-scale 5.0 \
--output helios_v2v_base.mp4
```
**Helios-Mid** (Stage 2 + CFG-Zero*):
```bash
python end2end.py \
--model BestWishYsh/Helios-Mid \
--sample-type v2v \
--video-path /path/to/video.mp4 \
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
--guidance-scale 5.0 \
--is-enable-stage2 \
--pyramid-num-inference-steps-list 20 20 20 \
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
--output helios_v2v_mid.mp4
```
**Helios-Distilled**:
```bash
python end2end.py \
--model BestWishYsh/Helios-Distilled \
--sample-type v2v \
--video-path /path/to/video.mp4 \
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
--num-frames 240 \
--guidance-scale 1.0 \
--is-enable-stage2 \
--pyramid-num-inference-steps-list 2 2 2 \
--is-amplify-first-chunk \
--output helios_v2v_distilled.mp4
```
## Common Parameters
| Parameter | Default | Description |
|-----------|---------|-------------|
| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path |
| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` |
| `--prompt` | — | Text prompt describing the video |
| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) |
| `--image-path` | — | Input image (required for `i2v`) |
| `--video-path` | — | Input video (required for `v2v`) |
| `--height` | `384` | Video height in pixels |
| `--width` | `640` | Video width in pixels |
| `--num-frames` | `99` | Number of output frames |
| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) |
| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) |
| `--seed` | `42` | Random seed |
| `--fps` | `16` | Output video frame rate |
| `--output` | `helios_output.mp4` | Output file path |
### Stage 2 / Pyramid Parameters
| Parameter | Default | Description |
|-----------|---------|-------------|
| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising |
| `--pyramid-num-stages` | `3` | Number of pyramid stages |
| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage |
| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) |
### CFG-Zero* Parameters
| Parameter | Default | Description |
|-----------|---------|-------------|
| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) |
| `--use-zero-init` | off | Zero init for first steps |
| `--zero-steps` | `1` | Number of zero-init steps |
### Memory & Parallelism
| Parameter | Default | Description |
|-----------|---------|-------------|
| `--vae-use-slicing` | off | Enable VAE slicing |
| `--vae-use-tiling` | off | Enable VAE tiling |
| `--enforce-eager` | off | Disable torch.compile |
| `--enable-cpu-offload` | off | CPU offloading |
| `--enable-layerwise-offload` | off | Layerwise offloading |
| `--ulysses-degree` | `1` | Ulysses SP degree |
| `--ring-degree` | `1` | Ring SP degree |
| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) |
| `--tensor-parallel-size` | `1` | Tensor parallelism size |
## Example materials
??? abstract "end2end.py"
``````py
--8<-- "examples/offline_inference/helios/end2end.py"
``````
@@ -1,35 +0,0 @@
# HunyuanImage-3.0 Image-to-Text Inference
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/hunyuan_image3>.
This example demonstrates how to run HunyuanImage-3.0 Image-to-Text with the vLLM-Omni.
## Local CLI Usage
Download the example image:
```bash
wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
```
Run example:
```bash
python image_to_text.py \
--image cherry_blossom.jpg \
--prompt "<|startoftext|>You are an assistant that understands images and outputs text.<img>Describe the content of the picture."
```
Key arguments:
- `--model`: Model used. Default is: tencent/HunyuanImage-3.0-Instruct (Optional).
- `--image`: Path to input image (required).
- `--prompt`: Text description used to guide image understanding (required).
## Example materials
??? abstract "image_to_text.py"
``````py
--8<-- "examples/offline_inference/hunyuan_image3/image_to_text.py"
``````
@@ -68,7 +68,6 @@ Key arguments:
- `--vae-use-slicing`: Enable VAE slicing for memory optimization.
- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
- `--use-hsdp`: Enable Hybrid Sharded Data Parallel to shard model weights across GPUs.
- `--hsdp-shard-size`: Number of GPUs to shard model weights across within each replica group. -1 (default) auto-calculates as world_size / replicate_size.
@@ -1,46 +0,0 @@
# MammothModa2-Preview
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/mammothmodal2_preview>.
## Run examples (MammothModa2-Preview)
Download model
```bash
hf bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
```
### Text-to-Image (T2I)
```bash
python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py \
--model ./MammothModa2-Preview \
--stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
--height 1024 \
--width 1024 \
--num-inference-steps 50 \
--text-guidance-scale 4.0 \
--out output.png
```
### Image Summary
```bash
python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
--model ./MammothModa2-Preview \
--stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
--question "Summarize this image." \
--image ./image.png
```
## Example materials
??? abstract "run_mammothmoda2_image_summarize.py"
``````py
--8<-- "examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py"
``````
??? abstract "run_mammothmoda2_t2i.py"
``````py
--8<-- "examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py"
``````
@@ -1,254 +0,0 @@
# MiMo-Audio Offline Inference
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/mimo_audio>.
This directory contains an offline demo for running MiMo-Audio models with vLLM Omni. It builds task-specific inputs and generates WAV files or text outputs locally.
## Model Overview
MiMo-Audio provides multiple task variants for audio understanding and generation:
- **tts_sft**: Basic text-to-speech generation from text input.
- **tts_sft_with_instruct**: TTS generation with explicit voice style instructions.
- **tts_sft_with_audio**: TTS generation with audio reference for voice cloning.
- **tts_sft_with_natural_instruction**: TTS generation from natural language descriptions embedded in text.
- **audio_trancribing_sft**: Transcribe audio to text (speech-to-text). (note: the upstream task name uses the spelling 'trancribing', don't fix it)
- **audio_understanding_sft**: Understand and analyze audio content with text queries.
- **audio_understanding_sft_with_thinking**: Audio understanding with reasoning chain.
- **spoken_dialogue_sft_multiturn**: Multi-turn spoken dialogue with audio input/output.
- **speech2text_dialogue_sft_multiturn**: Multi-turn dialogue converting speech to text.
- **text_dialogue_sft_multiturn**: Multi-turn text-only dialogue.
## Setup
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
### Environment Variables
The `MIMO_AUDIO_TOKENIZER_PATH` environment variable is mandatory due to the specialized architecture:
```bash
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
```
## Quick Start
Run a single sample for basic TTS:
```bash
python3 -u end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
--query-type tts_sft
```
Run batch samples for basic TTS:
```bash
python3 -u end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
--query-type tts_sft \
--num-prompts {batch_size}
```
When enabling multi-batch processing, if the total number of tokens passed to the next stage exceeds the `max_model_len` value in the `mimo_audio.yaml` configuration file, you must also synchronously update the `max_position_embeddings` value in `MiMo-Audio-7B-Instruct/config.json` to match the modified value.
Generated audio files are saved to `output_audio/` by default. `--num-prompts` also can be used to all tasks below.
## Task Usage
### tts_sft (Basic Text-to-Speech)
Generate speech from text input:
```bash
python3 -u end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
--query-type tts_sft \
--text "The weather is so nice today."
```
### tts_sft_with_instruct (TTS with Voice Instructions)
Generate speech with explicit voice style instructions:
```bash
python3 -u end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
--query-type tts_sft_with_instruct \
--text "The weather is so nice today." \
--instruct "Speak happily in a child's voice"
```
### tts_sft_with_audio (TTS with Audio Reference)
Generate speech using an audio reference for voice cloning:
```bash
python3 -u end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
--query-type tts_sft_with_audio \
--text "The weather is so nice today." \
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
```
### tts_sft_with_natural_instruction (Natural Language TTS)
Generate speech from text containing natural voice descriptions:
```bash
python3 -u end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
--query-type tts_sft_with_natural_instruction \
--text "In a panting young male voice, he said: I can't run anymore, wait for me!"
```
### audio_trancribing_sft (Speech-to-Text)
Transcribe audio to text:
```bash
python3 -u end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
--query-type audio_trancribing_sft \
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
```
### audio_understanding_sft (Audio Understanding)
Understand and analyze audio content with text queries:
```bash
python3 -u end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
--query-type audio_understanding_sft \
--text "Summarize the audio." \
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
```
### audio_understanding_sft_with_thinking (Audio Understanding with Reasoning)
Audio understanding with reasoning chain:
```bash
python3 -u end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
--query-type audio_understanding_sft_with_thinking \
--text "Summarize the audio." \
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
```
### spoken_dialogue_sft_multiturn (Multi-turn Spoken Dialogue)
Multi-turn dialogue with audio input and output:
```bash
python3 -u end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
--query-type spoken_dialogue_sft_multiturn \
--audio-path "./prompt_speech_zh_m.wav"
```
Note: This task uses hardcoded audio files in the script. The audio files used in examples are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
### speech2text_dialogue_sft_multiturn (Speech-to-Text Dialogue)
Multi-turn dialogue converting speech to text:
```bash
python3 -u end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
--query-type speech2text_dialogue_sft_multiturn
```
Note: This task uses hardcoded audio files and message lists in the script.
### text_dialogue_sft_multiturn (Text Dialogue)
Multi-turn text-only dialogue:
```bash
python3 -u end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
--query-type text_dialogue_sft_multiturn
```
Note: This task uses hardcoded message lists in the script.
## Troubleshooting
### Audio dependencies (soundfile, librosa)
This example depends on **soundfile** (read/write WAV) and **librosa** (load audio including MP3). Install the project requirements first:
```bash
pip install -r requirements/common.txt
# or at least: pip install soundfile>=0.13.1 librosa>=0.11.0
```
- **`soundfile` / libsndfile not found**
`soundfile` uses the C library **libsndfile**. On Linux, install the system package before pip:
- Debian/Ubuntu: `sudo apt-get install libsndfile1`
- For development builds: `sudo apt-get install libsndfile1-dev`
- Then: `pip install soundfile`
- **`librosa` fails to load MP3 or reports "No backend available"**
Loading MP3 (e.g. in `spoken_dialogue_sft_multiturn` with `.mp3` files) uses **ffmpeg** as the backend. Install ffmpeg:
- Debian/Ubuntu: `sudo apt-get install ffmpeg`
- macOS: `brew install ffmpeg`
- **`ImportError: No module named 'soundfile'` or `ModuleNotFoundError: ... librosa`**
Ensure you are in the same Python environment where vLLM Omni and the example dependencies are installed, and that `requirements/common.txt` (or the packages above) are installed.
### Tokenizer path
- **`MIMO_AUDIO_TOKENIZER_PATH` not set or model fails to find tokenizer**
Export the tokenizer path before running:
```bash
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
```
See [Environment Variables](#environment-variables) in Setup.
### Other
- If the model or stage config fails to load, check [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for memory and GPU settings.
- For errors when reading/writing WAV (e.g. unsupported format), ensure input files are standard WAV/MP3 and that `soundfile` is linked to a working libsndfile (see above).
## Notes
- The script uses default model paths and audio files embedded in `end2end.py`. Update them if your local cache path differs.
- Use `--output-dir` to change the output folder (default: `./output_audio`).
- Use `--num-prompts` to generate multiple prompts in one run (default: 1).
- Audio files used in multi-turn dialogue examples are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
- The script supports various configuration options for initialization timeouts, batch timeouts, and shared memory thresholds. See `--help` for details.
## Example materials
??? abstract "end2end.py"
``````py
--8<-- "examples/offline_inference/mimo_audio/end2end.py"
``````
??? abstract "message_base64_wav.json"
``````json
--8<-- "examples/offline_inference/mimo_audio/message_base64_wav.json"
``````
??? abstract "message_convert.py"
``````py
--8<-- "examples/offline_inference/mimo_audio/message_convert.py"
``````
??? abstract "process_speechdata.py"
``````py
--8<-- "examples/offline_inference/mimo_audio/process_speechdata.py"
``````
@@ -67,51 +67,6 @@ If media file paths are not provided, the script will use default assets. Suppor
- `multi_audios`: Multiple audio inputs
- `mixed_modalities`: Combination of video, image, and audio inputs
### Async-chunk (offline)
For true stage-level concurrency -- where downstream stages (Talker, Code2Wav)
start **before** the upstream stage (Thinker) finishes -- use the async_chunk
example. This requires:
1. A stage config YAML with ``async_chunk: true`` (e.g.
``qwen3_omni_moe_async_chunk.yaml``).
2. Hardware that matches the config (e.g. 2x H100 for the default 3-stage
config).
The async_chunk example uses ``AsyncOmni`` instead of the synchronous ``Omni``
class, which enables the async orchestrator to receive stage-0 intermediate
outputs and trigger downstream stages early. Chunk data flows directly between
stage workers via the in-worker ``OmniChunkTransferAdapter`` / connector,
**not** through the orchestrator.
#### Single prompt
```bash
cd examples/offline_inference/qwen3_omni
bash run_single_prompt_async_chunk.sh
```
#### Multiple prompts with concurrency control
```bash
bash run_multiple_prompts_async_chunk.sh --max-in-flight 4
```
#### Text-only output (skip audio generation)
```bash
python end2end_async_chunk.py --query-type text --modalities text
```
#### Custom stage config
```bash
python end2end_async_chunk.py \
--query-type use_audio \
--stage-configs-path /path/to/your_async_chunk.yaml
```
> **Note**: The synchronous ``end2end.py`` (using ``Omni``) is still the
> recommended entry point for non-async-chunk workflows. Only use the
> async_chunk example when you need the stage-level concurrency semantics
> described in PR #962 / #1151.
### FAQ
If you encounter error about backend of librosa, try to install ffmpeg with command below.
@@ -126,26 +81,14 @@ sudo apt install ffmpeg
``````py
--8<-- "examples/offline_inference/qwen3_omni/end2end.py"
``````
??? abstract "end2end_async_chunk.py"
``````py
--8<-- "examples/offline_inference/qwen3_omni/end2end_async_chunk.py"
``````
??? abstract "run_multiple_prompts.sh"
``````sh
--8<-- "examples/offline_inference/qwen3_omni/run_multiple_prompts.sh"
``````
??? abstract "run_multiple_prompts_async_chunk.sh"
``````sh
--8<-- "examples/offline_inference/qwen3_omni/run_multiple_prompts_async_chunk.sh"
``````
??? abstract "run_single_prompt.sh"
``````sh
--8<-- "examples/offline_inference/qwen3_omni/run_single_prompt.sh"
``````
??? abstract "run_single_prompt_async_chunk.sh"
``````sh
--8<-- "examples/offline_inference/qwen3_omni/run_single_prompt_async_chunk.sh"
``````
??? abstract "run_single_prompt_tp.sh"
``````sh
--8<-- "examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh"
@@ -90,43 +90,6 @@ Examples:
python end2end.py --query-type Base --mode-tag icl
```
## Voice and Language Control
### Supported Voices (CustomVoice)
Predefined speaker voices are set via the `speaker` (or `voice_type`) field in `additional_information`. Available speakers depend on the loaded checkpoint; check `talker_config.spk_id` in the model config for the full list. Common voices include `vivian`, `ryan`, `aiden`, `ethan`, `serena` (case-insensitive).
Pass the speaker name in your request:
```python
additional_information = {
"text": ["你好,我是通义千问"],
"task_type": ["CustomVoice"],
"speaker": ["Vivian"],
"language": ["Chinese"],
}
```
### Supported Languages
The `language` field controls the codec-level language tag. Use `"Auto"` (default) for automatic detection.
Supported values: `Auto`, `Chinese`, `English`, `Japanese`, `Korean`, `German`, `French`, `Russian`, `Portuguese`, `Spanish`, `Italian`.
```python
additional_information = {
"text": ["Hello, nice to meet you."],
"task_type": ["CustomVoice"],
"speaker": ["Aiden"],
"language": ["English"],
}
```
### VoiceDesign and Base
- **VoiceDesign**: Use `instruct` for natural-language voice description; no `speaker` needed.
- **Base**: Use `ref_audio` and `ref_text` for voice cloning; `language` is optional.
## Streaming Mode
Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
@@ -3,50 +3,14 @@
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_image>.
Generate images from text prompts using vLLM-Omni's diffusion pipeline entrypoints.
This folder provides several entrypoints for experimenting with `Qwen/Qwen-Image` `Qwen/Qwen-Image-2512` `Tongyi-MAI/Z-Image-Turbo` `stepfun-ai/NextStep-1.1` using vLLM-Omni, note that NextStep-1.1 has different architecture so we treat it differently regarding running arguments and pipeline.
- `text_to_image.py`: command-line script for single image generation with advanced options.
- `gradio_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
- `web_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
## Table of Contents
Note that when you pass in multiple independent prompts, they will be processed sequentially. Batching requests is currently not supported.
- [Overview](#overview)
- [Quick Start](#quick-start)
- [Key Arguments](#key-arguments)
- [More CLI Examples](#more-cli-examples)
- [Web UI Demo](#web-ui-demo)
## Overview
This folder provides several entrypoints for experimenting with text-to-image diffusion models using vLLM-Omni. Note that `NextStep-1.1` has a different architecture, so it is treated differently regarding running arguments and pipeline.
### Supported Models
| Model | Image Shape | Peak VRAM (GiB) * | Model Weights (GiB) |
| ----- | ----------- | ----------- | ----------------- |
| `Qwen/Qwen-Image` | 1024 x 1024 | 60.0 | 53.7 |
| `Qwen/Qwen-Image-2512` |1024 x 1024 | 60.0 | 53.7 |
| `Tongyi-MAI/Z-Image-Turbo` | 1024 x 1024 | 24.8 | 19.2 |
| `stepfun-ai/NextStep-1.1` | 512 x 512 | 71.8 | 28.1 |
| `meituan-longcat/LongCat-Image` | 1024 x 1024 | 71.2 | 27.3 |
| `AIDC-AI/Ovis-Image-7B` | 1024 x 1024 | 71.8 | 17.1 |
| `OmniGen2/OmniGen2` | 1024 x 1024 | 20.1 | 14.7 |
| `stabilityai/stable-diffusion-3.5-medium` | 1024 x 1024 | 20.1 | 15.6 |
| `black-forest-labs/FLUX.1-dev` | 1024 x 1024 | 77.6 | 31.4 |
| `black-forest-labs/FLUX.2-klein-4B` | 1024 x 1024 | 72.7 | 14.9 |
| `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 |
| `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
!!! info
*Peak VRAM: based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU.
Default model: `Qwen/Qwen-Image`
## Quick Start
### Python API
Single-prompt generation:
## Basic Usage
```python
from vllm_omni.entrypoints.omni import Omni
@@ -59,74 +23,78 @@ if __name__ == "__main__":
images[0].save("coffee.png")
```
### Local CLI Usage
Or put more than one prompt in a request.
```bash
python text_to_image.py \
--model Qwen/Qwen-Image \
--prompt "a cup of coffee on the table" \
--output coffee.png
```python
from vllm_omni.entrypoints.omni import Omni
if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
prompts = [
"a cup of coffee on a table",
"a toy dinosaur on a sandy beach",
"a fox waking up in bed and yawning",
]
outputs = omni.generate(prompts)
for i, output in enumerate(outputs):
image = output.request_output.images[0].save(f"{i}.jpg")
```
## Key Arguments
!!! info
**Common arguments:**
However, it is not currently recommended to do so
because not all models support batch inference,
and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
| Argument | Type | Default | Description |
| -------- | ---- | ------- | ----------- |
| `--prompt` | str | `"a cup of coffee on the table"` | Text description for image generation |
| `--seed` | int | `142` | Integer seed for deterministic sampling |
| `--negative-prompt` | str | `None` | Negative prompt for classifier-free conditional guidance |
| `--cfg-scale` | float | `4.0` | True CFG scale (model-specific guidance strength) |
| `--guidance-scale` | float | `1.0` | Classifier-free guidance scale |
| `--num-images-per-prompt` | int | `1` | Number of images per prompt (saved as `output`, `output_1`, ...) |
| `--num-inference-steps` | int | `50` | Diffusion sampling steps (more steps = higher quality, slower) |
| `--height` | int | `1024` | Output image height in pixels |
| `--width` | int | `1024` | Output image width in pixels |
| `--output` | str | `"qwen_image_output.png"` | Path to save the generated image |
| `--vae-use-slicing` | flag | off | Enable VAE slicing for memory optimization |
| `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization |
| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
| `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
| `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
| `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
!!! info
**NextStep-1.1 specific arguments:**
For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
batched inputs, you can [modify this configuration](https://github.com/vllm-project/vllm-omni/tree/main/configuration/stage_configs.md) to let the model accept a longer batch of prompts.
| Argument | Type | Default | Description |
| -------- | ---- | ------- | ----------- |
| `--guidance-scale-2` | float | `1.0` | Secondary guidance scale (e.g. image-level CFG) |
| `--timesteps-shift` | float | `1.0` | Timesteps shift parameter for sampling |
| `--cfg-schedule` | str | `"constant"` | CFG schedule type: `"constant"` or `"linear"` |
| `--use-norm` | flag | off | Apply layer normalization to sampled tokens |
Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
This is useful for models that support negative prompts.
> If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
```python
from vllm_omni.entrypoints.omni import Omni
> Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
outputs = omni.generate([
{
"prompt": "a cup of coffee on a table",
"negative_prompt": "low resolution"
},
{
"prompt": "a toy dinosaur on a sandy beach",
"negative_prompt": "cinematic, realistic"
}
])
for i, output in enumerate(outputs):
image = output.request_output.images[0].save(f"{i}.jpg")
```
## More CLI Examples
## Local CLI Usage
### Tongyi Models
### Qwen/Tongyi Models
```bash
python text_to_image.py \
--model Tongyi-MAI/Z-Image-Turbo \
--prompt "a cup of coffee on the table" \
--seed 42 \
--guidance-scale 0.0 \
--cfg-scale 4.0 \
--num-images-per-prompt 1 \
--num-inference-steps 9 \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output outputs/coffee.png
```
`Tongyi-MAI/Z-Image-Turbo` is a distilled version of Z-Image. Distilled diffusion models usually require less number of inference steps (4~9), and Classifier-Free Guidance (CFG) is usually NOT applied. Similar distilled models are `black-forest-labs/FLUX.2-klein-4B` and `black-forest-labs/FLUX.2-klein-9B`.
### NextStep Models
NextStep-1.1 supports extra arguments for dual-level CFG control:
NextStep-1.1 can have extra arguments
```bash
python text_to_image.py \
--model stepfun-ai/NextStep-1.1 \
@@ -141,101 +109,38 @@ python text_to_image.py \
--seed 42
```
### FLUX.2-dev Models
### Key Arguments
To run FLUX.2-dev on a single GPU, `--enable-cpu-offload` is required because the model weights exceed 80 GiB:
**Common arguments:**
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model black-forest-labs/FLUX.2-dev \
--prompt "a lovely bunny holding a sign that says 'vllm-omni'" \
--seed 42 \
--tensor-parallel-size 1 \
--num-images-per-prompt 1 \
--num-inference-steps 50 \
--guidance-scale 4.0 \
--height 1024 \
--width 1024 \
--enable-cpu-offload \
--output flux2-dev.png
```
- `--prompt`: text description (string).
- `--seed`: integer seed for deterministic sampling.
- `--cfg-scale`: true CFG scale (model-specific guidance strength).
- `--num-images-per-prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
- `--num-inference-steps`: diffusion sampling steps (more steps = higher quality, slower).
- `--height/--width`: output resolution (defaults 1024x1024).
- `--output`: path to save the generated PNG.
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion_acceleration.md#using-cfg-parallel).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
- `--guidance-scale`: classifier-free guidance scale.
### Batch Requests (Multiple Prompts)
**NextStep-1.1 specific:**
- `--guidance-scale-2`: secondary guidance scale, e.g. image-level CFG (default: 1.0).
- `--timesteps-shift`: timesteps shift parameter for sampling (default: 1.0).
- `--cfg-schedule`: CFG schedule type, "constant" or "linear" (default: "constant").
- `--use-norm`: apply layer normalization to sampled tokens.
You can pass multiple prompts in a single `generate` call.
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
```python
from vllm_omni.entrypoints.omni import Omni
> ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
prompts = [
"a cup of coffee on a table",
"a toy dinosaur on a sandy beach",
"a fox waking up in bed and yawning",
]
outputs = omni.generate(prompts)
for i, output in enumerate(outputs):
output.request_output.images[0].save(f"{i}.jpg")
```
## LoRA
!!! info
This example supports Peft-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. Pass `--lora-path` to use a LoRA adapter and optionally `--lora-scale` (default 1.0); omit it to use the base model only.
Not all models support batch inference, and batch requesting mostly does not provide significant
performance improvement. This feature is primarily for interface compatibility with vLLM and to
allow for future improvements.
!!! info
For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by
default, and the input list is sliced into single-item requests before feeding into the diffusion
pipeline. For models that do internally support batched inputs, you can
[modify this configuration](https://github.com/vllm-project/vllm-omni/tree/main/configuration/stage_configs.md) to let the model accept a
longer batch of prompts.
### Negative Prompts
vLLM-Omni supports dictionary prompts for models that accept negative prompts:
```python
from vllm_omni.entrypoints.omni import Omni
if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
outputs = omni.generate([
{
"prompt": "a cup of coffee on a table",
"negative_prompt": "low resolution"
},
{
"prompt": "a toy dinosaur on a sandy beach",
"negative_prompt": "cinematic, realistic"
}
])
for i, output in enumerate(outputs):
output.request_output.images[0].save(f"{i}.jpg")
```
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model Qwen/Qwen-Image \
--prompt "a cup of coffee on a table" \
--negative-prompt "low resolution, blurry" \
--output coffee.png
```
### Advanced Features
#### CFG Parallel
Set `--cfg-parallel-size 2` to enable CFG Parallel for faster inference on multi-GPU setups.
See more examples in the [diffusion acceleration user guide](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion_acceleration.md#using-cfg-parallel).
#### LoRA
This example supports PEFT-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. Pass `--lora-path` to use a LoRA adapter and optionally `--lora-scale` (default `1.0`); omit it to use the base model only.
### Basic usage with LoRA
```bash
python text_to_image.py \
@@ -246,7 +151,14 @@ python text_to_image.py \
--output output.png
```
LoRA adapters must be in PEFT format. A typical adapter directory structure:
### LoRA parameters
- `--lora-path`: Path to LoRA adapter folder (PEFT format). Loaded at initialization and used for generation.
- `--lora-scale`: Scale factor for LoRA weights (default: 1.0). Higher values increase the influence of the LoRA adapter.
### LoRA adapter format
LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
```
lora_adapter/
@@ -256,22 +168,13 @@ lora_adapter/
## Web UI Demo
!!! note "Gradio is an optional dependency"
The Gradio demo requires the `[demo]` extras. Install them first:
```bash
pip install 'vllm-omni[demo]'
```
Or, if installing from source: `pip install -e '.[demo]'`
Launch the Gradio demo:
Launch the gradio demo:
```bash
python gradio_demo.py --port 7862
```
Then open `http://localhost:7862/` in your local browser to interact with the web UI.
Then open `http://localhost:7862/` on your local browser to interact with the web UI.
## Example materials
@@ -3,20 +3,13 @@
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_video>.
A unified script for text-to-video generation. Supports multiple models with model-aware defaults.
The Wan T2V pipeline generates short videos from text prompts. Supported models include:
## Supported Models
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|---|---|---|---|---|---|
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
- `Wan-AI/Wan2.2-T2V-A14B-Diffusers` (14B, requires ~48GB VRAM)
- `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` (1.3B, requires ~16GB VRAM, use `--flow-shift 3.0 --boundary-ratio 0.0`)
## Local CLI Usage
### Wan2.2 (default)
```bash
python text_to_video.py \
--prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
@@ -32,119 +25,22 @@ python text_to_video.py \
--output t2v_out.mp4
```
LTX2 example:
Key arguments:
```bash
python text_to_video.py \
--model "Lightricks/LTX-2" \
--prompt "A cinematic close-up of ocean waves at golden hour." \
--negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
--height 512 \
--width 768 \
--num-frames 121 \
--num-inference-steps 40 \
--guidance-scale 4.0 \
--frame-rate 24 \
--output ltx2_out.mp4
```
### HunyuanVideo-1.5 (480p)
```bash
python text_to_video.py \
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
--prompt "A cat walks through a sunlit garden, flowers swaying gently in the breeze." \
--height 480 \
--width 832 \
--num-frames 121 \
--guidance-scale 6.0 \
--flow-shift 5.0 \
--num-inference-steps 50 \
--fps 24 \
--output hunyuan_video_15_output.mp4
```
### HunyuanVideo-1.5 (720p)
```bash
python text_to_video.py \
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v \
--prompt "A serene lakeside sunrise with mist over the water." \
--height 720 \
--width 1280 \
--num-frames 121 \
--guidance-scale 6.0 \
--flow-shift 9.0 \
--num-inference-steps 50 \
--fps 24 \
--output hunyuan_720p.mp4
```
### HunyuanVideo-1.5 with FP8 Quantization
```bash
python text_to_video.py \
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
--prompt "A dog running across a field of golden wheat." \
--quantization fp8 \
--height 480 --width 832 --num-frames 121 \
--guidance-scale 6.0 --flow-shift 5.0 \
--output hunyuan_fp8.mp4
```
Quick test (smaller resolution, fewer frames):
```bash
python text_to_video.py \
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
--prompt "A serene lakeside sunrise with mist over the water." \
--height 320 --width 576 --num-frames 17 --num-inference-steps 30 \
--flow-shift 5.0 \
--output quick_test.mp4
```
## Key Arguments
### Common
- `--model`: Diffusers model ID or local path.
- `--prompt`: text description (string).
- `--height/--width`: output resolution. Default depends on model.
- `--num-frames`: number of frames. Default depends on model.
- `--guidance-scale`: CFG scale. Default depends on model.
- `--num-inference-steps`: sampling steps. Default depends on model.
- `--fps`: frames per second for the saved MP4.
- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
- `--num-frames`: Number of frames (Wan default is 81).
- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high).
- `--negative-prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
- `--boundary-ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
- `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
- `--output`: path to save the generated video.
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
- `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
- `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio).
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
- `--flow-shift`: scheduler flow_shift parameter.
### Wan2.2-specific
- `--negative-prompt`: artifacts to suppress.
- `--guidance-scale-high`: separate CFG scale for high-noise stage.
- `--boundary-ratio`: boundary split for low/high DiT (default 0.875).
- `--flow-shift`: scheduler flow_shift (5.0 for 720p, 12.0 for 480p).
- `--cache-backend`: `cache_dit` for acceleration.
### HunyuanVideo-1.5 Optimal Configs
| Variant | flow_shift | guidance_scale | steps |
|---------|-----------|----------------|-------|
| 480p T2V | 5.0 | 6.0 | 50 |
| 720p T2V | 9.0 | 6.0 | 50 |
| 480p I2V | 5.0 | 6.0 | 50 |
| 720p I2V | 7.0 | 6.0 | 50 |
| CFG-distilled | (same) | 1.0 | 50 |
> If you encounter OOM errors, try `--vae-use-slicing`, `--vae-use-tiling`, `--enable-cpu-offload`, or `--quantization fp8`.
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
## Example materials
@@ -1,74 +0,0 @@
# Voxtral TTS Offline Inference
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/voxtral_tts>.
`end2end.py` runs Voxtral TTS end-to-end offline inference using vLLM. It supports both blocking (`Omni`) and streaming (`AsyncOmni`) generation, batched prompts with configurable concurrency, and voice selection via preset name or reference audio file.
When `mistral_common` has `SpeechRequest` support, prompt token IDs are built via `encode_speech_request`. Otherwise, the script falls back to manual token construction.
## Usage Examples
```bash
# Basic single-prompt with cheerful_female voice preset
python3 examples/offline_inference/voxtral_tts/end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
--write-audio --voice cheerful_female \
--model mistralai/Voxtral-4B-TTS-2603 \
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
# 32 replicate prompts with cheerful_female voice preset
python3 examples/offline_inference/voxtral_tts/end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
--num-prompts 32 --write-audio --voice cheerful_female \
--model mistralai/Voxtral-4B-TTS-2603 \
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
# Streaming with neutral_female voice preset
python3 examples/offline_inference/voxtral_tts/end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
--streaming --write-audio --voice neutral_female \
--model mistralai/Voxtral-4B-TTS-2603 \
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
# 32 prompts, 8 concurrent requests per wave, streaming with neutral_female voice
python3 examples/offline_inference/voxtral_tts/end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
--num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \
--model mistralai/Voxtral-4B-TTS-2603 \
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
# Short debug prompt with reference audio
# Note: Reference audio capability is not yet released.
python3 examples/offline_inference/voxtral_tts/end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
--write-audio \
--model mistralai/Voxtral-4B-TTS-2603 \
--text "This is a test message." \
--audio-path path/to/reference_audio.wav
```
## Arguments
| Argument | Description |
|---|---|
| `--model PATH` | HuggingFace repo ID or local directory path (default: `mistralai/Voxtral-4B-TTS-2603`) |
| `--text TEXT` | Text to synthesize (default: `"This is a test message."`) |
| `--audio-path PATH` | Path to reference audio file for voice cloning |
| `--output-dir DIR` | Directory to write output WAV files (default: `output_audio`) |
| `--stage-configs-path PATH` | Path to stage configs YAML (currently it must be set for VoxtralTTS) |
| `--num-prompts N` | Number of replicate prompts to run for measuring performance (default: 1) |
| `--streaming` | Use streaming generation via `AsyncOmni` (default: blocking `Omni`) |
| `--concurrency N` | Max concurrent requests per wave (must be used with `--streaming`, must evenly divide `--num-prompts`) |
| `--voice NAME` | Voice preset to use instead of reference audio file (e.g., casual_female, casual_male, cheerful_female, neutral_female, neutral_male) |
| `--write-audio` | Write generated audio to WAV files |
| `--profiling-mode` | Enable profiling mode (reduces max tokens to 50) |
| `--log-stats` | Enable detailed statistics logging |
## Example materials
??? abstract "end2end.py"
``````py
--8<-- "examples/offline_inference/voxtral_tts/end2end.py"
``````
@@ -1,8 +1,5 @@
# X-To-Video-Audio
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/x_to_video_audio>.
The `DreamID-Omni` pipeline generates short videos from text, image and video.
## Local CLI Usage
@@ -59,14 +56,3 @@ Key arguments:
- `--video-negative-prompt`: negative prompt for video generation.
- `--audio-negative-prompt`: negative prompt for audio generation.
- `--enable-cpu-offload`: enable CPU offload (defaults False).
## Example materials
??? abstract "download_dreamid_omni.py"
``````py
--8<-- "examples/offline_inference/x_to_video_audio/download_dreamid_omni.py"
``````
??? abstract "x_to_video_audio.py"
``````py
--8<-- "examples/offline_inference/x_to_video_audio/x_to_video_audio.py"
``````
@@ -1,226 +0,0 @@
# vLLM-Omni Helm Chart
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/chart-helm>.
Helm chart for deploying [vLLM-Omni](https://github.com/vllm-project/vllm-omni) on Kubernetes. vLLM-Omni extends vLLM with omni-modality model serving, supporting text-to-image, multimodal chat, text-to-speech, and more.
## Prerequisites
- Kubernetes 1.24+
- Helm 3.x
- NVIDIA GPU nodes with [NVIDIA Device Plugin](https://github.com/NVIDIA/k8s-device-plugin)
## Quick Start
```bash
helm install my-release ./chart-helm \
--set model=Tongyi-MAI/Z-Image-Turbo
```
## Configuration
### Model Selection
Set the `model` value to any supported HuggingFace model ID:
| Model | Type | GPUs | Notes |
|-------|------|------|-------|
| `Tongyi-MAI/Z-Image-Turbo` | text-to-image | 1 | Small, fast (default) |
| `stabilityai/stable-diffusion-3.5-medium` | text-to-image | 1 | ~6GB VRAM |
| `Qwen/Qwen-Image` | text-to-image | 1 | Large, ~40GB+ VRAM |
| `Qwen/Qwen2.5-Omni-7B` | multimodal | 2 | Text + audio + image + video |
| `Qwen/Qwen3-Omni-7B-Chat` | multimodal | 2 | Latest omni model |
| `Qwen/Qwen3-TTS` | text-to-speech | 1 | TTS |
### HuggingFace Token
For gated models that require authentication:
```bash
helm install my-release ./chart-helm \
--set model=Qwen/Qwen2.5-Omni-7B \
--set hfToken=hf_xxxxx \
--set resources.requests."nvidia\.com/gpu"=2 \
--set resources.limits."nvidia\.com/gpu"=2
```
### Omni-Specific Flags
Enable VAE memory optimizations for diffusion models:
```bash
helm install my-release ./chart-helm \
--set model=Qwen/Qwen-Image \
--set omniArgs.vaeUseSlicing=true \
--set omniArgs.vaeUseTiling=true
```
Enable CPU offloading:
```bash
helm install my-release ./chart-helm \
--set model=Qwen/Qwen-Image \
--set omniArgs.enableCpuOffload=true
```
Pass additional raw CLI flags:
```bash
helm install my-release ./chart-helm \
--set model=Qwen/Qwen-Image \
--set omniArgs.extraArgs[0]="--enable-layerwise-offload"
```
### Model Cache
By default, a PersistentVolumeClaim is created for the HuggingFace model cache to avoid re-downloading models on pod restarts:
```yaml
modelCache:
enabled: true
storageSize: "50Gi"
storageClassName: ""
```
To use an ephemeral volume instead:
```bash
helm install my-release ./chart-helm \
--set modelCache.enabled=false
```
### Custom Command Override
To fully override the container command:
```bash
helm install my-release ./chart-helm \
--set image.command[0]=vllm \
--set image.command[1]=serve \
--set image.command[2]=my-model \
--set image.command[3]=--omni \
--set image.command[4]=--host \
--set image.command[5]=0.0.0.0
```
## API Endpoints
Once deployed, vLLM-Omni exposes the following OpenAI-compatible endpoints:
| Endpoint | Method | Description |
|----------|--------|-------------|
| `/health` | GET | Health check |
| `/v1/models` | GET | List available models |
| `/v1/chat/completions` | POST | Chat completions (text/multimodal) |
| `/v1/images/generations` | POST | Image generation |
| `/v1/images/edits` | POST | Image editing |
| `/v1/audio/speech` | POST | Text-to-speech |
## Files
| File | Description |
|------|-------------|
| `Chart.yaml` | Chart metadata (name, version, maintainers) |
| `values.yaml` | Default configuration values |
| `values.schema.json` | JSON schema for validating values |
| `templates/_helpers.tpl` | Helper templates for common configurations |
| `templates/deployment.yaml` | Kubernetes Deployment |
| `templates/service.yaml` | Kubernetes Service (ClusterIP) |
| `templates/secrets.yaml` | Secrets (generic + HuggingFace token) |
| `templates/pvc.yaml` | PersistentVolumeClaim for model cache |
| `templates/configmap.yaml` | Optional ConfigMap |
| `templates/hpa.yaml` | HorizontalPodAutoscaler |
| `templates/poddisruptionbudget.yaml` | PodDisruptionBudget |
| `templates/custom-objects.yaml` | Custom Kubernetes objects |
## Running Tests
This chart includes unit tests using [helm-unittest](https://github.com/helm-unittest/helm-unittest). Install the plugin and run tests:
```bash
# Install plugin
helm plugin install https://github.com/helm-unittest/helm-unittest
# Run tests
helm unittest .
```
## Example materials
??? abstract ".helmignore"
``````
--8<-- "examples/online_serving/chart-helm/.helmignore"
``````
??? abstract "Chart.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/Chart.yaml"
``````
??? abstract "ct.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/ct.yaml"
``````
??? abstract "lintconf.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/lintconf.yaml"
``````
??? abstract "templates/_helpers.tpl"
``````tpl
--8<-- "examples/online_serving/chart-helm/templates/_helpers.tpl"
``````
??? abstract "templates/configmap.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/templates/configmap.yaml"
``````
??? abstract "templates/custom-objects.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/templates/custom-objects.yaml"
``````
??? abstract "templates/deployment.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/templates/deployment.yaml"
``````
??? abstract "templates/hpa.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/templates/hpa.yaml"
``````
??? abstract "templates/ingress.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/templates/ingress.yaml"
``````
??? abstract "templates/poddisruptionbudget.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/templates/poddisruptionbudget.yaml"
``````
??? abstract "templates/pvc.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/templates/pvc.yaml"
``````
??? abstract "templates/secrets.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/templates/secrets.yaml"
``````
??? abstract "templates/service.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/templates/service.yaml"
``````
??? abstract "tests/deployment_test.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/tests/deployment_test.yaml"
``````
??? abstract "tests/ingress_test.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/tests/ingress_test.yaml"
``````
??? abstract "tests/pvc_test.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/tests/pvc_test.yaml"
``````
??? abstract "tests/secrets_test.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/tests/secrets_test.yaml"
``````
??? abstract "values.yaml"
``````yaml
--8<-- "examples/online_serving/chart-helm/values.yaml"
``````
@@ -11,15 +11,6 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serv
## Gradio Demo
!!! note "Gradio is an optional dependency"
The Gradio demo requires the `[demo]` extras. Install them first:
```bash
pip install 'vllm-omni[demo]'
```
Or, if installing from source: `pip install -e '.[demo]'`
An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support.
```bash
@@ -73,42 +73,90 @@ The default yaml configuration deploys AR on GPU 0 and DiT on GPU 1. You can use
### Text-to-Image
Generate images from text prompts:
**Using Python client**
```bash
python openai_chat_client.py \
--prompt "A photorealistic mountain landscape at sunset" \
--height 1024 \
--width 1024 \
--output landscape.png
```
# Or use the curl script:
**Using curl**
```bash
curl -s http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"messages": [
{"role": "user", "content": "A beautiful sunset over the ocean with sailing boats"}
],
"extra_body": {
"height": 1024,
"width": 1024,
"num_inference_steps": 50,
"guidance_scale": 1.5,
"seed": 42
}
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
```
Or use the script:
```bash
bash run_curl_text_to_image.sh "A futuristic city skyline at night"
```
### Image-to-Image (Image Editing)
Edit images with text instructions:
**Using Python client**
```bash
python openai_chat_client.py \
--prompt "Convert this image to watercolor style" \
--image input.png \
--output watercolor.png
```
# Or use the curl script:
**Using curl**
```bash
IMG_B64=$(base64 < input.png | tr -d '\n')
curl -s http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d @- <<EOF | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
{
"messages": [{
"role": "user",
"content": [
{"type": "text", "text": "Convert this image to watercolor style"},
{"type": "image_url", "image_url": {"url": "data:image/png;base64,'$IMG_B64'"}}
]
}],
"extra_body": {
"height": 1024,
"width": 1024,
"num_inference_steps": 50,
"guidance_scale": 1.5,
"seed": 42
}
}
EOF
```
Or use the script:
```bash
bash run_curl_image_edit.sh input.png "Convert to watercolor style"
```
For general-purpose request methods (curl, OpenAI SDK, Python `requests`), see
the [Text-to-Image](text_to_image.md) and [Image-to-Image](image_to_image.md)
guides.
## Generation Parameters
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
When using the dedicated [`/v1/images/generations`](../../../../serving/image_generation_api.md)
or [`/v1/images/edits`](../../../../serving/image_edit_api.md) endpoints, pass
the supported generation controls as top-level fields directly. For image
dimensions and count, use `size` and `n` rather than `height` or `width`.
## Generation Parameters (extra_body)
| Parameter | Type | Default | Description |
| --------------------- | ----- | ------- | ----------------------------------- |
@@ -116,7 +164,7 @@ dimensions and count, use `size` and `n` rather than `height` or `width`.
| `width` | int | 1024 | Image width in pixels |
| `num_inference_steps` | int | 50 | Number of diffusion denoising steps |
| `guidance_scale` | float | 1.5 | Classifier-free guidance scale |
| `seed` | int | None | Optional random seed; `/v1/images/*` generates one server-side if omitted |
| `seed` | int | 42 | Random seed for reproducibility |
| `negative_prompt` | str | None | Negative prompt |
## Response Format
@@ -69,49 +69,10 @@ cat <<EOF > request.json
}
EOF
curl -s http://localhost:8092/v1/chat/completions \
-H "Content-Type: application/json" \
-d @request.json \
| jq -r '.choices[0].message.content[0].image_url.url' \
| cut -d',' -f2 | base64 -d > output.png
curl -s http://localhost:8092/v1/chat/completions -H "Content-Type: application/json" -d @request.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png
```
### Method 2: Using OpenAI Python SDK
```python
import base64
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8092/v1", api_key="none")
with open("input.png", "rb") as f:
img_b64 = base64.b64encode(f.read()).decode()
response = client.chat.completions.create(
model="Qwen/Qwen-Image-Edit",
messages=[{
"role": "user",
"content": [
{"type": "text", "text": "Convert to watercolor style"},
{"type": "image_url", "image_url": {
"url": f"data:image/png;base64,{img_b64}"
}},
],
}],
extra_body={
"num_inference_steps": 50,
"guidance_scale": 1,
"seed": 42,
},
)
img_url = response.choices[0].message.content[0].image_url.url
_, b64_data = img_url.split(",", 1)
with open("output.png", "wb") as f:
f.write(base64.b64decode(b64_data))
```
### Method 3: Using Python Client Script
### Method 2: Using Python Client
```bash
python openai_chat_client.py --input input.png --prompt "Convert to oil painting style" --output output.png
@@ -120,16 +81,7 @@ python openai_chat_client.py --input input.png --prompt "Convert to oil painting
python openai_chat_client.py --input input1.png input2.png --prompt "Combine these images into a single scene" --output output.png
```
### Method 4: Using Gradio Demo
!!! note "Gradio is an optional dependency"
The Gradio demo requires the `[demo]` extras. Install them first:
```bash
pip install 'vllm-omni[demo]'
```
Or, if installing from source: `pip install -e '.[demo]'`
### Method 3: Using Gradio Demo
```bash
python gradio_demo.py
@@ -172,7 +124,7 @@ python gradio_demo.py
### Image Editing with Parameters
Wrap generation parameters inside `extra_body` in the request JSON:
Use `extra_body` to pass generation parameters:
```json
{
@@ -195,149 +147,6 @@ Wrap generation parameters inside `extra_body` in the request JSON:
}
```
!!! tip "Using the OpenAI SDK"
When using the OpenAI Python SDK, pass these parameters via the `extra_body`
keyword argument. The SDK merges them into the top-level request body automatically:
```python
client.chat.completions.create(
model="Qwen/Qwen-Image-Edit",
messages=[...],
extra_body={"num_inference_steps": 50, "guidance_scale": 7.5, "seed": 42},
)
```
For details on how generation parameters are handled across different clients, see the
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md).
### Layered Image Generation (Qwen-Image-Layered)
Qwen-Image-Layered generates multiple decomposed layers from a reference image and a text prompt.
Start the server with:
```bash
vllm serve Qwen/Qwen-Image-Layered --omni --port 8093
```
=== "curl"
```bash
IMG_B64=$(base64 -w0 input.png)
curl -sS http://localhost:8093/v1/chat/completions \
-H "Content-Type: application/json" \
-d "$(jq -n --arg img "$IMG_B64" '{
messages: [{
role: "user",
content: [
{type: "image_url", image_url: {url: ("data:image/png;base64," + $img)}},
{type: "text", text: "a rabbit"}
]
}],
extra_body: {
num_inference_steps: 50,
cfg_scale: 4.0,
seed: 0,
layers: 4,
resolution: 640
}
}')" \
| jq -r '.choices[0].message.content[] | .image_url.url | split(",")[1]' \
| while IFS= read -r b64; do
((i++)); echo "$b64" | base64 -d > "layer_${i}.png"
done
```
=== "OpenAI SDK"
```python
import base64
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8093/v1", api_key="none")
with open("input.png", "rb") as f:
img_b64 = base64.b64encode(f.read()).decode()
response = client.chat.completions.create(
model="Qwen/Qwen-Image-Layered",
messages=[{
"role": "user",
"content": [
{"type": "image_url", "image_url": {
"url": f"data:image/png;base64,{img_b64}"
}},
{"type": "text", "text": "a rabbit"},
],
}],
extra_body={
"num_inference_steps": 50,
"cfg_scale": 4.0,
"seed": 0,
"layers": 4,
"resolution": 640,
},
)
for i, item in enumerate(response.choices[0].message.content):
_, b64_data = item.image_url.url.split(",", 1)
with open(f"layer_{i}.png", "wb") as f:
f.write(base64.b64decode(b64_data))
```
=== "Python requests"
```python
import base64
import requests
with open("input.png", "rb") as f:
img_b64 = base64.b64encode(f.read()).decode()
payload = {
"messages": [{
"role": "user",
"content": [
{"type": "image_url", "image_url": {
"url": f"data:image/png;base64,{img_b64}"
}},
{"type": "text", "text": "a rabbit"},
],
}],
"extra_body": {
"num_inference_steps": 50,
"cfg_scale": 4.0,
"seed": 0,
"layers": 4,
"resolution": 640,
},
}
resp = requests.post(
"http://localhost:8093/v1/chat/completions",
json=payload,
timeout=600,
)
data = resp.json()
for i, item in enumerate(data["choices"][0]["message"]["content"]):
_, b64_data = item["image_url"]["url"].split(",", 1)
with open(f"layer_{i}.png", "wb") as f:
f.write(base64.b64decode(b64_data))
```
The response contains multiple images in `choices[0].message.content` — one per generated layer.
#### Qwen-Image-Layered Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `layers` | int | 4 | Number of layers to decompose |
| `resolution` | int | 640 | Resolution for dimension calculation (640 or 1024) |
| `cfg_scale` | float | 4.0 | Classifier-free guidance scale (alias for `true_cfg_scale`) |
| `num_inference_steps` | int | 50 | Number of denoising steps |
| `seed` | int | None | Random seed for reproducibility |
### Multi-Image Editing (Qwen-Image-Edit-2509)
Provide multiple images in `content` (order matters):
@@ -357,15 +166,7 @@ Provide multiple images in `content` (order matters):
}
```
## Generation Parameters
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
When using the dedicated [`/v1/images/edits`](../../../../serving/image_edit_api.md)
endpoint, pass the supported generation controls as top-level form fields
directly. For image dimensions and count, use `size` and `n` rather than
`height`, `width`, or `num_outputs_per_prompt`.
## Generation Parameters (extra_body)
| Parameter | Type | Default | Description |
| ------------------------ | ----- | ------- | ------------------------------------- |
@@ -373,12 +174,10 @@ directly. For image dimensions and count, use `size` and `n` rather than
| `width` | int | None | Output image width in pixels |
| `size` | str | None | Output image size (e.g., "1024x1024") |
| `num_inference_steps` | int | 50 | Number of denoising steps |
| `guidance_scale` | float | 1.0 | CFG guidance scale |
| `guidance_scale` | float | 7.5 | CFG guidance scale |
| `seed` | int | None | Random seed (reproducible) |
| `negative_prompt` | str | None | Negative prompt |
| `num_outputs_per_prompt` | int | 1 | Number of images to generate |
| `layers` | int | 4 | Number of layers (Qwen-Image-Layered) |
| `resolution` | int | 640 | Resolution, 640 or 1024 (Qwen-Image-Layered) |
## Response Format
Loaded 100 of 425 files, more files were not shown because too many files have changed in this diff. Show more