Compare commits

...
Author SHA1 Message Date
Samit 26cef6e1ff Update latest news section in README.md
Signed-off-by: Samit <285365963@qq.com>
2026-08-14 11:22:30 +08:00
596c16a550 [Kernel] Fuse Q/K RMSNorm and RoPE (#5990)
Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Signed-off-by: Mengjie Zhao <zmj0129@gmail.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-08-14 00:46:59 +00:00
Lucas(Ruijie) 90a08c0eda [Model][TTS][WIP] Add dots.tts (rednote-hilab) — continuous-AR 48kHz TTS (#4765)
Signed-off-by: Moore-Z <zhumoore9661996@gmail.com>
2026-08-13 17:58:33 +00:00
Sy03 dbc0dd6d85 [Bugfix][Higgs-Audio-V3] Disable XQA decode (#6068)
Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: Sysy <1370724210@qq.com>
2026-08-13 10:00:50 -07:00
Mu GuanLin 83d1f11449 [Feat] LTX Standard Two-Stage Pipeline (#5500)
Signed-off-by: mglyn <1203789601@qq.com>
2026-08-14 00:13:14 +08:00
Yukim1andHongsheng Liu 4b66f62e10 [Refactor] Remove internal stage_configs_path plumbing (#5741)
Signed-off-by: zwhzzz0821 <2831474076@qq.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 21:29:03 +08:00
Weiming Liao f41451f64c [NPU] Upgrade to v0.27.0 (#6096)
Signed-off-by: Weiming Liao <liaowm5@gmail.com>
2026-08-13 20:01:01 +08:00
Hongsheng Liu bbe6ccc512 [Doc] Refresh community documentation (#6141)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 17:04:06 +08:00
Canlin Guo bbb709cfb5 [BugFix] Use MOSS-TTS-Local official sample params (#6156)
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
2026-08-13 08:56:26 +00:00
vraitiandClaude Opus 4.6 c98517ce7f [Metrics] Add Diffusion Metrics (#4755)
Signed-off-by: vraiti <vraiti@redhat.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-08-13 06:59:13 +00:00
NATUREandAlicia d77a35a32b [Bugfix] Scope Bagel FP8 config to diffusion stage (#6085)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-08-13 11:23:47 +08:00
liwenerandamy-why-3459 877beeaa49 [Perf][MiniCPM-O-4.5]Add CUDA Graph For CFM DiT estimator (#6082)
Signed-off-by: liwener <liwener.c@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-13 09:36:17 +08:00
Hongsheng Liu 8b8d3f56db [Misc] Consolidate LingBot text-to-video runner (#6076)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 08:42:37 +08:00
Hongsheng Liu c57178c0e1 [Dependency] Upgrade Cache-DiT to 1.5.0 (#6065)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 08:41:59 +08:00
mjZhaoElaineandClaude Opus 5 9a59475015 [CI/Build] Fix ReadTheDocs build under RTD's seeded pip 23.1 (#6129)
Signed-off-by: Mengjie Zhao <zmj0129@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-08-13 06:52:01 +08:00
Nick Cao 12f8168179 [Tests] Enable tiny model testing for Qwen-Image Edit and EditPlus (#5656)
Signed-off-by: Nick Cao <ncao@redhat.com>
2026-08-12 16:03:49 -04:00
Zhenghan Song 72a02b492a [Bugfix] Fix diffusion TTS adapter lookup in speech serving (#6121)
Signed-off-by: HaningZS <resossr0v0@gmail.com>
2026-08-12 15:20:12 -04:00
Juan Pablo ZuluagaandYueqian Lin 5358f5db01 [Perf][TTS] omnivoice hot path: D2H batching, mask caching, cached text embeddings (#5174)
Signed-off-by: JuanPZuluaga <juanz9312@gmail.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-08-12 17:15:35 +00:00
MaciejBalaNV e677360168 Add cosmos-guardrail dependency and update error message (#6107)
Signed-off-by: Maciej Bala <mbala@nvidia.com>
2026-08-12 17:04:07 +00:00
Juan Pablo ZuluagaandYueqian Lin 208c6776a8 [Perf][TTS] glm_tts: hoist loop-invariant text embed, RoPE, mask and CFG batch out of Euler loop (#5068)
Signed-off-by: JuanPZuluaga <juanz9312@gmail.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-08-12 17:02:02 +00:00
162 changed files with 9664 additions and 1338 deletions

No files matched your search

+13 -7
View File
@@ -345,7 +345,7 @@ steps:
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
@@ -361,7 +361,7 @@ steps:
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
@@ -377,7 +377,7 @@ steps:
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
@@ -393,7 +393,7 @@ steps:
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
@@ -409,7 +409,7 @@ steps:
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
@@ -456,7 +456,7 @@ steps:
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
- pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx2_expansion.py -m "full_model and cuda" --run-level "full_model"
- pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx.py -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test"
@@ -477,10 +477,16 @@ steps:
commands:
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_one_stage_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · LTX Accuracy Test"
timeout_in_minutes: 180
commands:
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx23_two_stage_layer_fused_i2v]' -m full_model --run-level full_model
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2V · Perf Test"
key: nightly-diffusion-x2v-performance
timeout_in_minutes: 180
+1 -1
View File
@@ -115,7 +115,7 @@ steps:
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export CACHE_DIT_VERSION=1.3.0
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
+7 -2
View File
@@ -11,10 +11,15 @@ build:
post_checkout:
- git fetch --unshallow || true
pre_install:
# `pre_install` runs before RTD upgrades pip, and the seeded pip 23.1
# rejects the `typing_extensions` wheel torch needs (fixed in pip 23.2).
- python -m pip install --upgrade pip
# Pre-install CPU torch so the subsequent `pip install .[docs]` does not
# pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
# 1 GB memory cap and time out the mkdocs build.
- python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
# 1 GB memory cap and time out the mkdocs build. Wheel-only: `--index-url`
# replaces PyPI, so an sdist fallback cannot reach its build deps. The
# quotes keep YAML from reading `:all:` as a mapping key.
- "python -m pip install --only-binary=:all: --index-url https://download.pytorch.org/whl/cpu torch"
mkdocs:
configuration: mkdocs.yml
+1
View File
@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
---
*Latest News* 🔥
- [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0).
- [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support.
- [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
- [2026/06] Starting with [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0), vLLM-Omni publishes a stable release aligned with every even-numbered upstream vLLM minor version. [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0), [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0), [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0), and [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) continued this cadence, expanding omni and world-model support with [NVIDIA Cosmos3](recipes/cosmos3/Cosmos3-Nano.md) and DreamZero, adding models such as MiniCPM-o 4.5, MOSS-TTS, and Lance, and advancing TTS, diffusion, distributed execution, quantization, RL integration through [VeRL-Omni](https://github.com/verl-project/verl-omni), and CUDA/ROCm/MUSA/NPU/XPU coverage.
@@ -167,9 +167,11 @@ def main() -> None:
]
native_cmd = [
sys.executable,
"examples/offline_inference/text_to_video/text_to_video_lingbot.py",
"examples/offline_inference/text_to_video/text_to_video.py",
"--model",
args.model,
"--model-class-name",
"LingBotVideoPipeline",
"--prompt",
args.prompt,
"--output",
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.26.0
ARG VLLM_ASCEND_TAG=v0.27.0
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
# WORKDIR /vllm-workspace/vllm
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.26.0-a3
ARG VLLM_ASCEND_TAG=v0.27.0-a3
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
# WORKDIR /vllm-workspace/vllm
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.26.0
ARG VLLM_ASCEND_TAG=v0.27.0
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
ARG APP_DIR=/vllm-workspace/vllm-omni
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.26.0-a3
ARG VLLM_ASCEND_TAG=v0.27.0-a3
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
ARG APP_DIR=/vllm-workspace/vllm-omni
+2 -1
View File
@@ -163,7 +163,8 @@ nav:
- CLI Reference: cli
- Community:
- Governance: community/governance.md
- community/*
- Contact Us: community/contact_us.md
- Community Meetings: community/meetings.md
- Slack: https://slack.vllm.ai
- Blog: https://blog.vllm.ai
- Forum: https://discuss.vllm.ai
+7 -3
View File
@@ -1,5 +1,9 @@
# Contact Us
- For technical questions and feature requests, please use GitHub [Issues](https://github.com/vllm-project/vllm-omni/issues)
- For coordinating contributions and development and discussing with other users and developers, please join `sig-omni` channel in our [Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/)
- For security disclosures, please use GitHub's [Security Advisories](https://github.com/vllm-project/vllm-omni/security/advisories) feature
- For bugs and feature requests, search or open a GitHub [issue](https://github.com/vllm-project/vllm-omni/issues).
- For technical and usage questions, join the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/).
- For coordinating contributions and development, use `sig-omni`, the forum, or join a [community meeting](meetings.md).
## Security Disclosures
Do not share vulnerability details in public GitHub issues, pull requests, discussions, or Slack channels. To arrange a private report, send a direct message to one of the [lead maintainers](governance.md#lead-maintainers) in [vLLM Slack](https://slack.vllm.ai/). If you cannot send a direct message, you may post a detail-free request in `sig-omni` asking a lead maintainer to contact you privately; do not include any vulnerability details in that request.
+3 -23
View File
@@ -52,31 +52,11 @@ Committers have write access and merge rights. They typically have deep expertis
- [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
- [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
## Diffusion Workload Division
## Reviewer Routing
This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review.
For path-based reviewer routing, consult [`.github/CODEOWNERS`](https://github.com/vllm-project/vllm-omni/blob/main/.github/CODEOWNERS). Detailed subsystem boundaries and safe-change guidance live in the [module design documents](../design/index.md#module-design-documents).
### Runtime-related
- Diffusion models: @RuixiangMa, @wtomin, @xuechendi
### Optimization
- Parallel: @RuixiangMa, @wtomin
- Attention backend: @ZJY0516
- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
- Offload: @yuanheng-zhao
- Intel GPU support: @xuechendi
- Quantization: @lishunyang12, @david6666666
- Kernel fusion / communication-computation: @ZJY0516
### Docs & Test
- Diffusion docs: @RuixiangMa, @wtomin
## Meetings
Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
`CODEOWNERS` is operational routing metadata, not a governance or approval policy. Being listed does not grant committer status or merge rights, and it does not mean every listed reviewer must approve a change. The maintainer roster above is authoritative for project roles and merge rights.
## Committer Nomination Process
+8
View File
@@ -0,0 +1,8 @@
# Community Meetings
vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. Schedules and access details can change, so use these maintained documents as the source of truth:
- [Chinese-language meeting document](https://docs.google.com/document/d/1pdUBiS_7mdOUNDtdwy-9OUf7jsMWN324BIbps5olDME/edit?tab=t.0)
- [English-language meeting document](https://docs.google.com/document/d/19AlD08T8uJ7wqfs8VtYCi0iGMTwlTYVVlAxFZ3S7OIg/edit?tab=t.0)
Each document contains the current schedule, access details, agendas, and past meeting notes. For broader vLLM community events, see the [vLLM events calendar](https://vllm.ai/events/).
-1
View File
@@ -1 +0,0 @@
# Meetups
-12
View File
@@ -1,12 +0,0 @@
# Volunteers for Bugfix and CI
We encourage you to check current docs and [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions for your questions. If non of these can solve it, please propose an issue to describe your questions about bug or CI problems for developing.
If you have urgent need for locating and solving bugfix or CI problems, please find community volunteers below.
| Dec 4-Dec 12 | Dec 15-Dec 19 | Dec 22-Dec 26 | Dec 29- Jan 2, 2026| Jan 5-Jan 9 | Jan 12-Jan 16 |
|----------|----------|----------|----------|----------|----------|
| <a href="https://github.com/congw729">Conw729</a> | <a href="https://github.com/yinpeiqi">yinpeiqi</a> | <a href="https://github.com/tzhouam">tzhouam</a> | <a href="https://github.com/SamitHuang">SamitHuang</a> | <a href="https://github.com/gcanlin">gcanlin</a> | <a href="https://github.com/natureofnature">natureofnature</a> |
| <a href="https://github.com/david6666666">david6666666</a> | <a href="https://github.com/R2-Y">R2-Y</a> | <a href="https://github.com/hsliuustc0106">hsliuustc0106</a> | <a href="https://github.com/Gaohan123">Gaohan123</a> | <a href="https://github.com/ZJY0516">ZJY0516</a> | <a href="https://github.com/qibaoyuan">qibaoyuan</a> |
We kindly welcome more contributors to fix bugs and contribute new features!
+2 -2
View File
@@ -3,7 +3,7 @@
Thank you for your interest in contributing to vLLM-Omni! This document provides guidelines and instructions for contributing.
!!! note
We host weekly developer-facing online meetings to discuss milestones and updates **every Tuesday at 19:30 PDT**. Meeting link as well as the past meeting notes can be found [here](https://tinyurl.com/vllm-omni-meeting).
vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. See [Community Meetings](../community/meetings.md) for current schedules, access details, agendas, and past notes.
## Getting Started
@@ -122,7 +122,7 @@ Repository skills guide the agent, but they do not replace contributor judgment,
If you encounter a bug or have a feature request, please search existing issues first to see if it has already been reported. If not, please file a new issue, providing as much relevant information as possible.
!!! important
If you discover a security vulnerability, please report it by creating a GitHub issue with the `security` label.
Do not report suspected security vulnerabilities through a public issue, pull request, discussion, or Slack channel. Follow the [security disclosure instructions](../community/contact_us.md#security-disclosures) to arrange a private report.
## Pull Requests & Code Reviews
+1 -1
View File
@@ -1,6 +1,6 @@
# Prometheus Metrics Design
This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. This iteration covers the text and audio surface; diffusion / image / video families land in a follow-up.
This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics.
## Objectives
+3 -1
View File
@@ -42,7 +42,8 @@ th {
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — |
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | — |
| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
| `LTX2DistilledPipeline` | LTX-2 / LTX-2.3 full-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + `Lightricks/LTX-2.3` upsampler | ✅︎ | ✅︎ | | | — |
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -58,6 +59,7 @@ th {
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
| `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+3 -3
View File
@@ -10,12 +10,12 @@ A: Refer to [GPU memory calculation and configuration](../configuration/gpu_memo
> Q: I encounter some bugs or CI problems, which is urgent. How can I solve it?
A: At first, you can check current [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions. If none of these satisfy your demand and it is urgent, please find these [volunteers](https://docs.vllm.ai/projects/vllm-omni/en/latest/community/volunteers/) for help.
A: Search the current [issues](https://github.com/vllm-project/vllm-omni/issues) first. If no existing issue matches a bug or CI failure, open a [bug report](https://github.com/vllm-project/vllm-omni/issues/new?template=400-bug-report.yml) with a minimal reproducer, environment details, and relevant logs. For technical or usage questions, use the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or the [vLLM Forum](https://discuss.vllm.ai/).
> Q: Does vLLM-Omni support AWQ or any other quantization?
A: We plan to introduce GGUF FP8 prequantized models and online FP8 quantization in version 0.16.0. Support for other quantization types will follow in future releases. For details, please see our [Q1 quantization roadmap](https://github.com/vllm-project/vllm-omni/issues/1057).
A: AWQ is available through the inherited vLLM quantization registry, and vLLM-Omni supports several additional methods. Actual support and validation vary by model, component, hardware, and checkpoint. See the [quantization overview](../user_guide/quantization/overview.md) for validated combinations and method-specific guides.
> Q: Does vLLM-Omni support multimodal streaming input and output?
A: Not yet. We already put it on the [Roadmap](https://github.com/vllm-project/vllm-omni/issues/165). Please stay tuned!
A: Yes, for specific models and serving paths. Current capabilities include [streaming video input](../serving/video_stream_api.md), [streaming diffusion output](../user_guide/diffusion/execution_modes.md#streaming-output), and experimental full-duplex real-time audio described in the [project overview](../README.md). Check the model and API documentation because input and output support is model-specific.
+50 -26
View File
@@ -1,23 +1,23 @@
# Production Metrics
vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. This page covers the text and audio surface; diffusion / image / video metrics are tracked in a follow-up PR.
vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server.
```bash
vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats
curl http://localhost:8000/metrics
```
**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm:omni_*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm_omni:*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
## Metric Namespaces
| Prefix | Source | Present when |
|--------|--------|--------------|
| `vllm:omni_` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
| `vllm_omni:` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
| `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage |
| `http_` / `process_` | Uvicorn / Python runtime | Always |
## Pipeline-Level Metrics (`vllm:omni_`)
## Pipeline-Level Metrics (`vllm_omni:`)
Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline.
@@ -25,42 +25,66 @@ Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the
| Metric | Type | Labels | Description |
|--------|------|--------|-------------|
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
| `vllm_omni:num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
| `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
| `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
### Latency
| Metric | Type | Labels | Description |
|--------|------|--------|-------------|
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
| `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
## Audio Modality Metrics (`vllm:omni_`)
## Audio Modality Metrics (`vllm_omni:`)
Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label.
| Metric | Type | Extra label | Description |
|--------|------|-------------|-------------|
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
| `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
| `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
| `vllm_omni:audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
| `vllm_omni:audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
| `vllm_omni:audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
| `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
| `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition.
## Cross-Stage Transfer Metrics (`vllm:omni_`)
## Diffusion Metrics (`vllm_omni:`)
Per-request timing breakdowns for diffusion (image/video) stages. Emitted at request finalize when `stage_metrics.diffusion_metrics` is present. All carry `{model_name, stage, replica}`.
| Metric | Type | Description |
|--------|------|-------------|
| `vllm_omni:diffusion_exec_s` | Histogram | DiT forward pass execution time per request in seconds |
| `vllm_omni:diffusion_exec_per_step_s` | Histogram | DiT forward pass execution time per denoising step in seconds (`exec_s / num_inference_steps`) |
| `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time per request in seconds |
| `vllm_omni:diffusion_postprocess_s` | Histogram | Diffusion output postprocessing (VAE decode) time per request in seconds |
`diffusion_exec_s` uses `SECONDS_BUCKETS` (50 ms → 300 s); the other three use `SECONDS_FAST_BUCKETS` (1 ms → 60 s). `diffusion_exec_per_step_s` is skipped when `num_inference_steps` is not set.
### Example PromQL
```promql
# P99 DiT exec time by stage
histogram_quantile(0.99, rate(vllm_omni:diffusion_exec_s_bucket[5m]))
# Ratio of preprocess to total exec wall time
rate(vllm_omni:diffusion_preprocess_s_sum[5m])
/ rate(vllm_omni:diffusion_exec_s_sum[5m])
```
## Cross-Stage Transfer Metrics (`vllm_omni:`)
Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed.
| Metric | Type | Description |
|--------|------|-------------|
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
| `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
| `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
| `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
| `vllm_omni:transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
## vLLM Engine Metrics (`vllm:`)
@@ -83,9 +107,9 @@ For the full list of upstream metrics, see [the vLLM docs](https://github.com/vl
| Metric group | Multi-stage LLM (Qwen3-Omni) |
|---|---|
| `vllm:omni_` request tracking + latency | With `--log-stats` |
| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage |
| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
| `vllm_omni:` request tracking + latency | With `--log-stats` |
| `vllm_omni:` audio modality | With `--log-stats`, if pipeline has a talker stage |
| `vllm_omni:` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
| `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` |
| `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` |
@@ -94,5 +118,5 @@ For the full list of upstream metrics, see [the vLLM docs](https://github.com/vl
- All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values.
- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
- Sizes use the `_bytes` suffix.
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label.
- All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm_omni:`.
- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm_omni:audio_ttfp_s` for audio) rather than a single metric with a `modality` label.
@@ -12,6 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|---|---|---|---|---|---|
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) |
| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
| `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 |
@@ -35,6 +36,25 @@ python text_to_video.py \
--output t2v_out.mp4
```
### LingBot-Video
The shared runner recognizes both official LingBot checkpoint IDs and selects
the native `LingBotVideoPipeline` automatically:
```bash
python text_to_video.py \
--model robbyant/lingbot-video-dense-1.3b \
--prompt "a robotic arm picks up a red block" \
--height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
--guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
--output lingbot_t2v.mp4
```
Use `robbyant/lingbot-video-moe-30b-a3b` for the MoE checkpoint. For local
paths without `lingbot` in the directory name, add
`--model-class-name LingBotVideoPipeline`. Model-only options use
`--extra-body`, for example `'{"batch_cfg": true, "output_type": "np"}'`.
### LTX-2
```bash
+3 -1
View File
@@ -80,12 +80,14 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
## Model Type Support
### Diffusion Model (Qwen-Image, Wan2.2)
### Diffusion Models
| Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization |
|-------|-----------|:-------:|:------:|----------------|------------------|------------------|
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | |
| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | |
| LTX-2 | `Lightricks/LTX-2`, `rootonchair/LTX-2-19b-distilled` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers`, `diffusers/LTX-2.3-Distilled-Diffusers` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None | ✅︎ |
| FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None | |
| FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | Yes | Not validated | All layers | None | ✅︎ |
@@ -59,7 +59,7 @@ def main(args):
omni = Omni(
model=args.model_name,
stage_configs_path=args.stage_configs_path,
deploy_config=args.deploy_config,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
batch_timeout=args.batch_timeout,
@@ -167,10 +167,10 @@ def parse_args():
help="Sampling rate for audio loading (default: 16000).",
)
parser.add_argument(
"--stage-configs-path",
"--deploy-config",
type=str,
default=None,
help="Path to stage configs YAML file.",
help="Path to a deploy config YAML file.",
)
parser.add_argument(
"--log-stats",
@@ -876,7 +876,7 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
parser.add_argument("--task", type=str, required=True, choices=TASK_CHOICES)
parser.add_argument("--model", type=str, required=True, help="HF repo id or local model directory.")
parser.add_argument(
"--stage-config-path",
"--deploy-config",
type=str,
default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
help="Path to deploy config yaml.",
@@ -1394,8 +1394,8 @@ def main() -> None:
from vllm_omni.entrypoints.omni import Omni
stage_config_path = str(Path(args.stage_config_path).expanduser())
omni = Omni(model=model_source, stage_configs_path=stage_config_path, dtype=args.dtype)
deploy_config_path = str(Path(args.deploy_config).expanduser())
omni = Omni(model=model_source, deploy_config=deploy_config_path, dtype=args.dtype)
sampling_params_list = [
SamplingParams(max_tokens=int(args.max_tokens_per_stage), temperature=0.0, top_p=1.0, detokenize=False)
for _ in range(omni.num_stages)
@@ -73,7 +73,6 @@ def parse_args():
)
parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
parser.add_argument(
"--stream",
action="store_true",
@@ -146,13 +145,7 @@ def main():
else:
bot_task = args.bot_task
if args.deploy_config is not None and args.stage_configs_path is not None:
raise ValueError("--deploy-config and --stage-configs-path are mutually exclusive.")
deploy_config = args.deploy_config
stage_configs_path = args.stage_configs_path
if deploy_config is None and stage_configs_path is None:
deploy_config = _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
deploy_config = args.deploy_config or _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
omni_kwargs = {
"model": args.model,
@@ -168,10 +161,7 @@ def main():
if additional_config is not None:
omni_kwargs["additional_config"] = additional_config
if deploy_config is not None:
omni_kwargs["deploy_config"] = deploy_config
else:
omni_kwargs["stage_configs_path"] = stage_configs_path
omni_kwargs["deploy_config"] = deploy_config
omni = Omni(**omni_kwargs)
@@ -272,10 +262,7 @@ def main():
print(f" Modality: {args.modality}")
print(f" Prompt task: {task}")
print(f" Bot task: {bot_task}")
if deploy_config is not None:
print(f" Deploy config: {deploy_config}")
else:
print(f" Stage config: {stage_configs_path}")
print(f" Deploy config: {deploy_config}")
print(f" Num stages: {omni.num_stages}")
if args.modality in ("text2img", "img2img"):
print(f" Inference steps: {args.steps}")
@@ -452,12 +452,6 @@ def parse_args():
default=None,
help="Path to a .txt file with one prompt per line (preferred).",
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="Path to a stage configs file (deprecated; prefer --deploy-config).",
)
parser.add_argument(
"--deploy-config",
type=str,
@@ -484,10 +484,10 @@ def parse_args():
help="Path to a .txt file with one prompt per line (preferred).",
)
parser.add_argument(
"--stage-configs-path",
"--deploy-config",
type=str,
default=None,
help="Path to a stage configs file.",
help="Path to a deploy config YAML.",
)
parser.add_argument(
"--video-path",
@@ -232,15 +232,8 @@ def main(args):
raise ValueError(f"Unknown query type: {args.query_type}")
# Initialize vLLM-Omni with Step-Audio2
# Resolve deploy config path. Keep explicit stage_configs_path for legacy
# custom configs, but use bundled deploy configs by default.
deploy_config_path = None
stage_config_path = None
if args.stage_configs_path and args.deploy_config:
raise ValueError("--stage-configs-path and --deploy-config are mutually exclusive")
if args.stage_configs_path:
stage_config_path = args.stage_configs_path
elif args.deploy_config:
# Resolve an explicit deploy config or use the bundled task default.
if args.deploy_config:
deploy_config_path = args.deploy_config
else:
configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/deploy"
@@ -262,10 +255,7 @@ def main(args):
"ray_address": args.ray_address,
"trust_remote_code": True,
}
if stage_config_path is not None:
omni_kwargs["stage_configs_path"] = stage_config_path
else:
omni_kwargs["deploy_config"] = deploy_config_path
omni_kwargs["deploy_config"] = deploy_config_path
omni_llm = Omni(**omni_kwargs)
@@ -418,12 +408,6 @@ def parse_args():
" - HuggingFace ID: stepfun-ai/Step-Audio-2-mini"
),
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="Path to a legacy stage config YAML file (default: use bundled deploy config)",
)
parser.add_argument(
"--deploy-config",
type=str,
@@ -92,12 +92,6 @@ def parse_args() -> argparse.Namespace:
"meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
"stabilityai/stable-diffusion-3.5-medium, Tongyi-MAI/Z-Image-Turbo and etc.",
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="[Deprecated] Path to a legacy stage_args-format YAML. Prefer --deploy-config.",
)
parser.add_argument(
"--deploy-config",
type=str,
@@ -475,8 +469,6 @@ def main():
omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
if args.enforce_eager is not None:
omni_kwargs["enforce_eager"] = args.enforce_eager
if args.stage_configs_path:
omni_kwargs["stage_configs_path"] = args.stage_configs_path
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
if use_nextstep:
@@ -515,8 +507,8 @@ def main():
print(f" Image size: {args.width}x{args.height}")
if args.lora_path:
print(f" LoRA: scale={args.lora_scale}")
if args.stage_configs_path:
print(f" stage-configs-path: {args.stage_configs_path}")
if args.deploy_config:
print(f" deploy-config: {args.deploy_config}")
print(f"{'=' * 60}\n")
# Build LoRA request when --lora-path is set
@@ -23,6 +23,7 @@ list of supported architectures across all modalities, see
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
| dots.tts | `rednote-hilab/dots.tts-soar` | single (native AR) | — (not wired yet) | — | — | 48 kHz |
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz |
| IndexTTS-2.5 | native `checkpoints/` bundle | 2 (AR talker + EnhancedCodec + S2Mel DiT + BigVGAN) | ✓ (required) | — | multilingual (`--lang`) + emotion control | 22.05 kHz |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
@@ -430,6 +431,31 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [
---
## dots.tts
Single-stage native AR TTS at 48 kHz (rednote-hilab). Pipeline: `Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR loopback → AudioVAE (streaming decode)`. Same "vLLM-native base LM + side-path computation" pattern as VoxCPM2, with a plain Qwen2 backbone instead of MiniCPM4 and no FSQ / residual LM stage.
### Quick start
```bash
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
--model rednote-hilab/dots.tts-soar \
--text "Hello, this is a test of dots TTS running on vLLM Omni."
```
### Voice cloning
Not wired in this release — generation is zero-shot only. The CAM++ x-vector speaker encoder weights load, but `end2end.py` has no `--ref-audio`/`--ref-text` flags yet.
### Streaming
The AudioVAE decoder has an internal streaming path (`init_stream_state` / `stream_step` / `stream_flush`) used to avoid boundary artifacts between 160 ms patches, but it is not yet exposed through an online serving endpoint or example.
### Notes
- Output: 48 kHz mono WAV.
- Deploy config: `vllm_omni/deploy/dots_tts.yaml` (auto-loaded by HF `model_type`).
- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default. `dots.tts-base` shares the same architecture but is unvalidated in this repo. `dots.tts-mf` (MeanFlow, 2-4 step) is not supported yet.
- Known limitation: no CUDA graph capture and no batched side-path computation yet, so concurrent requests do not currently scale (each request's DiT Euler steps run serially). See `recipes/rednote-hilab/dots.tts.md` for details and the roadmap.
---
## IndexTTS-2 and IndexTTS-2.5
Both versions use a 2-stage TTS pipeline and produce 22.05 kHz mono audio. Stage 0 is a GPT AR talker. IndexTTS-2 Stage 1 uses RepCodec semantic embeddings, GPT latent, S2Mel CFM/DiT, and BigVGAN. IndexTTS-2.5 replaces RepCodec with EnhancedCodec and uses the official code-only path (`use_gpt_latent=false`). Every request requires reference audio for zero-shot voice cloning.
@@ -0,0 +1,125 @@
"""Offline dots.tts inference example (single-stage native AR pipeline).
Uses the single-stage native AR config (vllm_omni/deploy/dots_tts.yaml).
Reference checkpoint: rednote-hilab/dots.tts-soar.
"""
from __future__ import annotations
import os
import time
from pathlib import Path
import soundfile as sf
import torch
from vllm_omni import Omni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
REPO_ROOT = Path(__file__).resolve().parents[4]
SAMPLE_RATE = 48_000
def parse_args():
parser = TrackingArgumentParser(description="Offline dots.tts native AR inference")
parser.add_argument(
"--model",
type=str,
default="rednote-hilab/dots.tts-soar",
help="dots.tts model path or HuggingFace repo ID.",
)
parser.add_argument(
"--text",
type=str,
default="Hello, this is a test of dots TTS running on vLLM Omni.",
help="Text to synthesize.",
)
parser.add_argument(
"--output-dir",
type=str,
default="output_audio",
help="Directory for output WAV files.",
)
parser.add_argument(
"--deploy-config",
type=str,
default=None,
help="Override the deploy config path. If unset, auto-loads "
"vllm_omni/deploy/dots_tts.yaml based on the HF model_type.",
)
return parser.parse_args()
def extract_audio(multimodal_output: dict) -> torch.Tensor:
"""Extract the final complete audio tensor from multimodal output.
The output processor concatenates per-step delta tensors under
``model_outputs``. Falls back to ``audio`` for backwards compat.
"""
audio = multimodal_output.get("model_outputs")
if audio is None:
audio = multimodal_output.get("audio")
if audio is None:
raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}")
if isinstance(audio, list):
valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
if not valid:
raise ValueError("Audio list is empty or all elements are None.")
return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
return torch.as_tensor(audio).float().cpu().reshape(-1)
def main():
args = parse_args()
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
engine = Omni(
model=args.model,
deploy_config=args.deploy_config,
)
from transformers import AutoTokenizer
from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import (
build_dots_tts_prompt,
)
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
prompt = build_dots_tts_prompt(
tokenizer=tokenizer,
text=args.text,
)
print(f"Model : {args.model}")
print(f"Text : {args.text}")
print(f"Output dir : {output_dir}")
print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens")
t_start = time.perf_counter()
outputs = engine.generate([prompt])
elapsed = time.perf_counter() - t_start
request_output = outputs[0]
mm = request_output.outputs[0].multimodal_output
audio = extract_audio(mm)
duration = audio.numel() / SAMPLE_RATE
rtf = elapsed / duration if duration > 0 else float("inf")
output_path = output_dir / "output.wav"
sf.write(str(output_path), audio.numpy(), SAMPLE_RATE, format="WAV")
print(f"Saved : {output_path}")
print(f"Duration : {duration:.2f}s")
print(f"Inference : {elapsed:.2f}s")
print(f"RTF : {rtf:.3f}")
if __name__ == "__main__":
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
main()
@@ -38,7 +38,7 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
logger = logging.getLogger(__name__)
DEFAULT_MODEL = "fishaudio/s2-pro"
DEFAULT_STAGE_CONFIG = os.path.join(
DEFAULT_DEPLOY_CONFIG = os.path.join(
os.path.dirname(__file__),
"..",
"..",
@@ -152,10 +152,7 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
else:
omni_kwargs["stage_configs_path"] = args.stage_configs_path or DEFAULT_STAGE_CONFIG
omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
return omni_kwargs
@@ -278,12 +275,6 @@ def parse_args():
default=None,
help="Deploy config YAML path (under vllm_omni/deploy/).",
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="Path to stage configs YAML (legacy, prefer --deploy-config).",
)
parser.add_argument(
"--output-dir",
default="output_audio",
@@ -145,7 +145,7 @@ def main(args):
omni = Omni(
model=args.model,
stage_configs_path=deploy_config_path,
deploy_config=deploy_config_path,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
@@ -490,10 +490,10 @@ def parse_args():
help="Path to a .txt file with one prompt per line (preferred).",
)
parser.add_argument(
"--stage-configs-path",
"--deploy-config",
type=str,
default=None,
help="Path to a stage configs file.",
help="Path to a deploy config YAML.",
)
parser.add_argument(
"--audio-path",
@@ -8,6 +8,7 @@ A unified script for text-to-video generation. Supports multiple models with mod
|---|---|---|---|---|---|
| `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) |
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) |
| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
| `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 |
@@ -57,18 +58,31 @@ python text_to_video.py \
### LingBot-Video
The dedicated LingBot example explicitly selects the video output modality:
The shared runner recognizes the official dense and MoE checkpoint IDs, selects
`LingBotVideoPipeline`, and builds the canonical video request envelope:
```bash
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
python text_to_video.py \
--model robbyant/lingbot-video-dense-1.3b \
--prompt "a robotic arm picks up a red block" \
--height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
--guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
--output lingbot_t2v.mp4
```
Use the same command with `robbyant/lingbot-video-moe-30b-a3b` for the MoE
checkpoint. Requested frame counts are rounded upward to the causal
VAE's `4n+1` grid.
VAE's `4n+1` grid. For a local checkpoint path whose name does not contain
`lingbot`, pass `--model-class-name LingBotVideoPipeline`.
LingBot-only options such as `batch_cfg` and `output_type` use the shared
model-extra channel, for example:
```bash
python text_to_video.py \
--model robbyant/lingbot-video-dense-1.3b \
--extra-body '{"batch_cfg": true, "output_type": "np"}'
```
### LTX-2
@@ -16,7 +16,7 @@ from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.lora.request import LoRARequest
from vllm_omni.lora.utils import stable_lora_int_id
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name, get_output_tensor_range
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
@@ -80,6 +80,17 @@ _MODEL_PRESETS = {
"fps": 16,
"output": "helios_output.mp4",
},
"lingbot": {
"model_class_name": "LingBotVideoPipeline",
"height": 192,
"width": 320,
"num_frames": 9,
"num_inference_steps": 2,
"guidance_scale": 3.0,
"fps": 24,
"flow_shift": 3.0,
"output": "lingbot_video_output.mp4",
},
"ltx2": {
"height": 512,
"width": 768,
@@ -110,27 +121,51 @@ _MODEL_PRESETS = {
def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
model_lower = model.lower()
class_lower = (model_class_name or "").lower()
if "distilled" in class_lower or "distilled" in model_lower:
return _MODEL_PRESETS["ltx2_distilled"]
if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
return _MODEL_PRESETS["ltx23"]
if "ltx2" in class_lower or "ltx-2" in model_lower or "ltx_2" in model_lower:
if "lingbot" in model_lower or "lingbotvideo" in class_lower:
return _MODEL_PRESETS["lingbot"]
if "ltx" in class_lower or "ltx" in model_lower:
if "distilled" in class_lower or "distilled" in model_lower:
return _MODEL_PRESETS["ltx2_distilled"]
if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
return _MODEL_PRESETS["ltx23"]
return _MODEL_PRESETS["ltx2"]
if "vace" in model_lower:
if "vace" in model_lower or "vace" in class_lower:
return _MODEL_PRESETS["vace"]
# Edge must be matched before the generic cosmos branch (its "cosmos" substring would
# otherwise pick up the Nano/Super 720p / gs 6.0 / flow_shift 10.0 preset).
if ("cosmos" in model_lower or "cosmos" in class_lower) and ("edge" in model_lower or "edge" in class_lower):
return _MODEL_PRESETS["cosmos3_edge"]
if "cosmos" in model_lower:
if "cosmos" in model_lower or "cosmos" in class_lower:
return _MODEL_PRESETS["cosmos"]
if "hunyuan" in model_lower:
if "hunyuan" in model_lower or "hunyuan" in class_lower:
return _MODEL_PRESETS["hunyuan"]
if "helios" in model_lower:
if "helios" in model_lower or "helios" in class_lower:
return _MODEL_PRESETS["helios"]
return _MODEL_PRESETS["wan"]
def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]:
"""Build the canonical request envelope for the shared T2V example."""
result: dict[str, Any] = {
"prompt": prompt,
"modalities": ["video"],
}
if negative_prompt is not None:
result["negative_prompt"] = negative_prompt
return result
def _normalize_float_tensor(tensor: torch.Tensor, source_range: str) -> torch.Tensor:
"""Normalize decoded frames according to the pipeline's output contract."""
if not tensor.is_floating_point():
return tensor
if source_range == "negative_one_to_one":
return tensor.clamp(-1, 1) * 0.5 + 0.5
if source_range == "zero_to_one":
return tensor.clamp(0, 1)
raise ValueError(f"Unsupported floating-point tensor range: {source_range!r}")
def parse_profiler_config(value: str) -> dict[str, Any]:
try:
config = json.loads(value)
@@ -160,7 +195,7 @@ def parse_extra_body(value: str) -> dict[str, Any]:
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Generate a video from a text prompt. "
"Supports Wan2.2, HunyuanVideo-1.5, Helios, and other text-to-video models."
"Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, and other text-to-video models."
)
parser.add_argument(
"--model",
@@ -408,6 +443,7 @@ def main():
for key, default_val in preset.items():
if getattr(args, key.replace("-", "_"), None) is None:
setattr(args, key.replace("-", "_"), default_val)
model_class_name = args.model_class_name
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
# Cache-dit config (Wan2.2 only)
@@ -480,6 +516,7 @@ def main():
omni = Omni(**omni_kwargs)
model_class_name = get_model_class_name(omni) or model_class_name
declared_extra_body_params = get_extra_body_params(model_class_name)
output_tensor_range = get_output_tensor_range(model_class_name)
if profiler_enabled:
print("[Profiler] Starting profiling...")
@@ -513,12 +550,11 @@ def main():
lora_path=lora_path,
)
prompt_dict = {"prompt": args.prompt}
if args.negative_prompt is not None:
prompt_dict["negative_prompt"] = args.negative_prompt
elif preset not in (_MODEL_PRESETS["ltx2"], _MODEL_PRESETS["ltx23"]):
negative_prompt = args.negative_prompt
if negative_prompt is None and all(preset is not _MODEL_PRESETS[name] for name in ("lingbot", "ltx2", "ltx23")):
# Preserve the historical empty-prompt behavior for non-LTX examples.
prompt_dict["negative_prompt"] = ""
negative_prompt = ""
prompt_dict = build_text_to_video_prompt(args.prompt, negative_prompt)
extra_args = {}
if lora_request:
@@ -630,8 +666,7 @@ def main():
frame_tensor = frame_tensor[0]
if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4):
frame_tensor = frame_tensor.permute(1, 2, 0)
if frame_tensor.is_floating_point():
frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5
frame_tensor = _normalize_float_tensor(frame_tensor, output_tensor_range)
return frame_tensor.float().numpy()
if isinstance(frame, np.ndarray):
frame_array = frame
@@ -683,8 +718,7 @@ def main():
video_tensor = video_tensor[0]
elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
video_tensor = video_tensor.permute(1, 2, 3, 0)
if video_tensor.is_floating_point():
video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
video_tensor = _normalize_float_tensor(video_tensor, output_tensor_range)
video_array = video_tensor.float().numpy()
elif isinstance(frames, np.ndarray):
video_array = frames
@@ -1,131 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import argparse
import time
from pathlib import Path
from typing import Any
import numpy as np
import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate text-to-video with a LingBot-Video checkpoint.")
parser.add_argument("--model", default="robbyant/lingbot-video-dense-1.3b")
parser.add_argument("--prompt", default="a robotic arm picks up a red block")
parser.add_argument("--negative-prompt", default=None)
parser.add_argument("--output", default="lingbot_video_output.mp4")
parser.add_argument("--height", type=int, default=192)
parser.add_argument("--width", type=int, default=320)
parser.add_argument("--num-frames", type=int, default=9)
parser.add_argument("--num-inference-steps", type=int, default=2)
parser.add_argument("--guidance-scale", type=float, default=3.0)
parser.add_argument("--flow-shift", type=float, default=3.0)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--fps", type=int, default=24)
parser.add_argument("--output-type", choices=["pt", "np"], default="pt")
parser.add_argument("--batch-cfg", action="store_true")
return parser.parse_args()
def _peak_memory_mb(output: Any) -> float:
result = output[0] if isinstance(output, list) and output else output
return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
def _extract_video(output: Any) -> Any:
if isinstance(output, list):
output = output[0] if output else None
if isinstance(output, OmniRequestOutput):
if output.final_output_type != "image":
raise ValueError(f"Unexpected output type {output.final_output_type!r}.")
if not output.images:
raise ValueError("No video frames found in OmniRequestOutput.")
return output.images[0] if len(output.images) == 1 else output.images
return output
def _to_frame_array(video: Any) -> np.ndarray:
if isinstance(video, torch.Tensor):
tensor = video.detach().cpu()
if tensor.dim() == 5:
tensor = tensor[0]
if tensor.dim() == 4 and tensor.shape[-1] not in (3, 4):
if tensor.shape[0] in (3, 4):
tensor = tensor.permute(1, 2, 3, 0)
elif tensor.shape[1] in (3, 4):
tensor = tensor.permute(0, 2, 3, 1)
if tensor.is_floating_point():
if float(tensor.min()) < 0.0:
tensor = tensor.clamp(-1, 1) * 0.5 + 0.5
else:
tensor = tensor.clamp(0, 1)
return tensor.float().numpy()
if isinstance(video, np.ndarray):
array = video[0] if video.ndim == 5 else video
if np.issubdtype(array.dtype, np.integer):
return array.astype(np.float32) / 255.0
return array.astype(np.float32, copy=False)
if isinstance(video, list):
return np.stack([_to_frame_array(frame) for frame in video], axis=0)
return np.asarray(video, dtype=np.float32)
def main() -> None:
args = parse_args()
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
omni = Omni(
model=args.model,
model_class_name="LingBotVideoPipeline",
flow_shift=args.flow_shift,
parallel_config=DiffusionParallelConfig(),
)
prompt: dict[str, Any] = {
"prompt": args.prompt,
"modalities": ["video"],
"num_frames": args.num_frames,
}
if args.negative_prompt is not None:
prompt["negative_prompt"] = args.negative_prompt
sampling_params = OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
num_frames=args.num_frames,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
generator=generator,
output_type=args.output_type,
extra_args={"batch_cfg": args.batch_cfg, "shift": args.flow_shift},
)
start = time.perf_counter()
output = omni.generate(prompt, sampling_params)
elapsed = time.perf_counter() - start
print(f"Total generation time: {elapsed:.4f} seconds ({elapsed * 1000:.2f} ms)")
peak_mb = _peak_memory_mb(output)
if peak_mb:
print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
video = _to_frame_array(_extract_video(output))
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
from diffusers.utils import export_to_video
export_to_video(list(video), str(output_path), fps=args.fps)
print(f"Saved generated video to {output_path}")
omni.close()
if __name__ == "__main__":
main()
+2
View File
@@ -73,6 +73,8 @@ plugins:
- awesome-nav
- redirects:
redirect_maps:
community/meetups.md: community/meetings.md
community/volunteers.md: community/contact_us.md
user_guide/diffusion/request_batching.md: user_guide/diffusion/execution_modes.md
user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
+72 -46
View File
@@ -4,11 +4,20 @@
## Pipelines
| Model | Checkpoint | Task | `--model-class-name` | Batching |
|---|---|---|---|---|
| LTX-2 | `Lightricks/LTX-2` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
| LTX-2 distilled | `rootonchair/LTX-2-19b-distilled` | Two-stage T2V/I2V | `LTX2DistilledPipeline` | No |
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
| `--model-class-name` | Task | Required checkpoint repositories |
|---|---|---|
| `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` |
| `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` |
| `LTX2DistilledPipeline` | LTX-2 full-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
| `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` |
| `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers`<br>`Lightricks/LTX-2.3` |
| `LTX2DistilledPipeline` | LTX-2.3 full-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |
Repositories in the table are download units. A full pipeline repository
contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler,
and tokenizer; an additional repository supplies LoRA or upsampler sidecars.
The runtime first searches the model root for sidecars under their official
filenames, then downloads them from the matching Lightricks Hub repository.
`LTX2Pipeline` is the unified one-stage entry. Checkpoint metadata selects the
LTX-2 or LTX-2.3 profile; omitting an image selects T2V, while one initial
@@ -16,10 +25,10 @@ image selects I2V. Both one-stage repositories declare this class, so
`--model-class-name` is optional. LTX-2.3 requires the Diffusers checkpoint;
the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable.
`LTX2DistilledPipeline` is the unified distilled entry: omit `image` for T2V
or provide one initial image for I2V. As in the official LTX pipelines,
distilled mode is selected explicitly rather than inferred from the checkpoint
name or metadata; always pass `--model-class-name LTX2DistilledPipeline`.
`LTX2TwoStagePipeline` samples the regular model at half resolution, upsamples,
then refines with the distilled LoRA. `LTX2DistilledPipeline` uses a fully
merged distilled Transformer in both stages. Both entries support T2V and I2V;
select their class explicitly.
## API Migration
@@ -68,24 +77,28 @@ The model recipe also supplies the default negative prompt. Top-level
API defaults `num_frames` to `1`, so set it explicitly; the offline LTX scripts
default to `121`.
## Distilled Two-Stage Defaults
## Two-Stage Defaults
| Parameter | Value |
|---|---:|
| Final width × height | 1536 × 1024 |
| Stage 1 width × height | 768 × 512 |
| Frames / frame rate | 121 / 24 |
| Stage 1 / Stage 2 denoise steps | 8 / 3 |
| Guidance | Fixed positive-only |
| Parameter | Ordinary | Full-distilled |
|---|---:|---:|
| Final width × height | 1536 × 1024 | 1536 × 1024 |
| Stage 1 width × height | 768 × 512 | 768 × 512 |
| Frames / frame rate | 121 / 24 | 121 / 24 |
| Stage 1 / Stage 2 steps | 40 (LTX-2) or 30 (LTX-2.3) / 3 | 8 / 3 |
| Guidance | Stage 1 guided; Stage 2 positive-only | Fixed positive-only |
The dimensions passed to the API are the final output dimensions and must be
divisible by 64. Both one- and two-stage requests require `num_frames = 8k+1`.
The distilled recipe fixes both sigma schedules; `num_inference_steps` must be
`8` when supplied.
API dimensions are final dimensions and must be divisible by 64. All LTX
requests require `num_frames = 8k+1`. Ordinary Stage 1 uses the LTX-2 or
LTX-2.3 one-stage defaults shown above. Distilled schedules are fixed, so
`num_inference_steps`, when supplied, must be `8`. Both entries reject custom
sigmas and input latents.
Ordinary two-stage uses layer-fused LoRA for an unquantized BF16 Transformer
and automatically switches to dynamic LoRA when quantization is enabled.
## Serving
Start either checkpoint:
Start either one-stage checkpoint:
```bash
vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
@@ -95,11 +108,19 @@ vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
vllm serve diffusers/LTX-2.3-Diffusers --omni --stage-init-timeout 600
```
Start the distilled checkpoint with its two-stage entry:
Start a two-stage checkpoint with its explicit class:
```bash
vllm serve rootonchair/LTX-2-19b-distilled --omni \
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
# LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent
vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
# Ordinary LTX-2.3 two-stage
vllm serve diffusers/LTX-2.3-Diffusers --omni \
--model-class-name LTX2TwoStagePipeline \
--enable-layerwise-offload \
--stage-init-timeout 600
```
The same server handles T2V and I2V. A T2V request using the selected recipe's
@@ -135,8 +156,9 @@ together with `input_reference`.
## Guidance
LTX supports independent video/audio CFG, spatio-temporal guidance (STG),
cross-modality guidance, and rescaling:
One-stage and ordinary Stage 1 support independent video/audio CFG,
spatio-temporal guidance (STG), cross-modality guidance, and rescaling.
Distilled stages and ordinary Stage 2 are fixed positive-only.
| Parameter | Default | Effect | Alias |
|---|---:|---|---|
@@ -172,9 +194,9 @@ complete guidance plan in parallel. Despite the legacy `cfg` name, it covers
text CFG, STG, and cross-modality guidance passes; guidance rescaling is
applied after all pass predictions are gathered.
The default one-stage recipe has four Transformer passes per denoise step:
`cond`, `uncond`, `ptb` (STG), and `mod` (cross-modality). The useful balanced
configurations are therefore:
The default one-stage recipe and ordinary Stage 1 have four Transformer passes
per denoise step: `cond`, `uncond`, `ptb` (STG), and `mod`
(cross-modality). The useful balanced configurations are therefore:
| `--cfg-parallel-size` | Passes per rank | Guidance-slot utilization | Notes |
|---:|---:|---:|---|
@@ -195,9 +217,8 @@ vllm serve diffusers/LTX-2.3-Diffusers --omni \
```
The total device count is the product of `cfg_parallel_size` and the other
configured parallel dimensions. Distilled pipelines use positive-only
guidance, so `cfg_parallel_size > 1` does not provide guidance speedup for
them.
configured parallel dimensions. Positive-only distilled phases do not benefit
from `cfg_parallel_size > 1`.
### Python API
@@ -242,22 +263,22 @@ noted below.
| `image` | image or batch, `None` | Direct value wins over request images; no image selects T2V. I2V accepts one image per prompt, and a batch cannot mix T2V/I2V. |
| `prompt` | string or list, `None` | Positive-text fallback; request prompts win. Mutually exclusive with `prompt_embeds`. |
| `negative_prompt` | string or list, `None` | Fallback after request values and before the recipe default; mutually exclusive with negative embeddings. |
| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 distilled. |
| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 two-stage. |
| `width` | `int`, `None` | Same precedence and alignment as `height`. |
| `num_frames` | `int`, `None` | Request → direct value → recipe default; must be `8k+1` and also determines audio duration with `frame_rate`. |
| `frame_rate` | `float`, `None` | Request `frame_rate` → request `fps` → direct value → recipe default. |
| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage and ordinary Stage 1, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
| `sigmas` | list of float, `None` | One-stage only. Request values win; every request in a fused batch must use the same schedule. |
| `timesteps` | list of int, `None` | Compatibility slot; LTX accepts only `None`. Use `sigmas`. |
| `guidance_scale` | `float`, `None` | One-stage common video/audio CFG fallback; an explicit request value wins. Distilled two-stage uses fixed positive-only guidance. |
| `guidance_scale` | `float`, `None` | One-stage and ordinary Stage 1 common video/audio CFG fallback; an explicit request value wins. Distilled phases use fixed positive-only guidance. |
| `guidance_rescale` | `float`, `None` | Accepts only `None` or `0.0`; use the modality rescale fields. |
| `noise_scale` | `float`, `0.0` | Compatibility slot; LTX accepts only `0.0`. |
| `num_videos_per_prompt` | `int`, `1` | Output-count fallback; positive request `num_outputs_per_prompt` wins. |
| `generator` | generator or list, `None` | Explicit RNG; otherwise request generators/seeds are collated. Lists must match the effective output batch. |
| `latents` | tensor, `None` | Request tensors win. One-stage accepts packed `[B, S, C]` and validated unpacked video layouts. |
| `audio_latents` | tensor, `None` | One-stage initial audio latents; request tensors win and are collated. |
| `latents` | tensor, `None` | One-stage only. Request tensors win; packed `[B, S, C]` and validated unpacked video layouts are accepted. |
| `audio_latents` | tensor, `None` | One-stage only. Request tensors win and are collated. |
| `prompt_embeds` | tensor, `None` | Precomputed positive conditioning; requires `prompt_attention_mask` and cannot accompany `prompt`. |
| `negative_prompt_embeds` | tensor, `None` | One-stage precomputed negative conditioning; requires its mask and cannot accompany `negative_prompt`. |
| `negative_prompt_embeds` | tensor, `None` | One-stage and ordinary Stage 1 conditioning; requires its mask and cannot accompany `negative_prompt`. |
| `prompt_attention_mask` | tensor, `None` | Mask for positive embeddings; request `prompt_attention_mask` or `attention_mask` wins. |
| `negative_prompt_attention_mask` | tensor, `None` | Mask for negative embeddings; request negative mask fields win. |
| `decode_timestep` | float or list, `0.0` | Video-VAE decode timestep; request value wins. Lists may match 1, prompt batch, or output batch. |
@@ -273,13 +294,13 @@ request prompt payload; LTX guidance fields live in sampling `extra_args`.
### Recipe-Specific Request Capabilities
| Override | One-stage | Distilled two-stage |
|---|---|---|
| Guidance | Supported | Fixed positive-only |
| Negative prompt/embeddings | Supported | Rejected |
| `num_inference_steps` | Supported | Fixed at 8 for Stage 1; Stage 2 uses 3 |
| Custom `sigmas` | Supported | Rejected; both phases use fixed schedules |
| Video/audio latents | Supported | Rejected |
| Override | One-stage | Ordinary two-stage | Distilled two-stage |
|---|---|---|---|
| Guidance | Supported | Stage 1 only; Stage 2 is positive-only | Fixed positive-only |
| Negative prompt/embeddings | Supported | Supported by Stage 1 | Rejected |
| `num_inference_steps` | Supported | Controls Stage 1; Stage 2 uses 3 | Fixed at 8 for Stage 1; Stage 2 uses 3 |
| Custom `sigmas` | Supported | Rejected | Rejected; both phases use fixed schedules |
| Video/audio latents | Supported | Rejected | Rejected |
These capability checks apply equally to direct `forward` keywords and
values in `OmniDiffusionSamplingParams`; unsupported values fail instead of
@@ -309,6 +330,10 @@ bundled offline CLI do not currently expose `sigmas`.
`--max-num-seqs 1`.
- `--cfg-parallel-size` shards the complete LTX guidance plan, including STG,
modality guidance, and rescale-compatible prediction gathering.
- Sequence parallelism may pad audio latents. Pure Ulysses masks the padding;
Ring cannot, so audio length must be SP-divisible. Use `ring_degree=1` or a
divisible request shape.
- Cache-DiT is one-stage only; multi-stage configurations are rejected.
## Operational Notes
@@ -323,14 +348,15 @@ bundled offline CLI do not currently expose `sigmas`.
request parameter.
- For benchmarks, use `tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with
`tests/dfx/perf/scripts/run_diffusion_benchmark.py`.
- Current two-stage support is limited to the distilled checkpoint; official
LTX two-stage/HQ execution remains out of scope.
- Ordinary and full-distilled two-stage T2V/I2V are supported; HQ execution
remains out of scope.
## References
- <https://huggingface.co/Lightricks/LTX-2>
- <https://huggingface.co/Lightricks/LTX-2.3>
- <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
- <https://huggingface.co/diffusers/LTX-2.3-Distilled-Diffusers>
- [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md)
- [Diffusion execution modes](../../docs/user_guide/diffusion/execution_modes.md)
- [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md)
+3 -3
View File
@@ -27,7 +27,7 @@ parallel, cache, quantized, or expert-kernel backends.
- Dense checkpoint: <https://huggingface.co/robbyant/lingbot-video-dense-1.3b>
- MoE checkpoint: <https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b>
- Upstream project: <https://github.com/Robbyant/lingbot-video>
- Related offline example: [`examples/offline_inference/text_to_video/text_to_video_lingbot.py`](../../examples/offline_inference/text_to_video/text_to_video_lingbot.py)
- Related offline example: [`examples/offline_inference/text_to_video/text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py)
- Related online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
## Hardware Support
@@ -52,7 +52,7 @@ The MoE path is validated with BF16 expert weights.
```bash
CUDA_VISIBLE_DEVICES=0 \
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
python examples/offline_inference/text_to_video/text_to_video.py \
--model robbyant/lingbot-video-dense-1.3b \
--prompt "a robotic arm picks up a red block" \
--output lingbot_t2v.mp4 \
@@ -70,7 +70,7 @@ python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
```bash
CUDA_VISIBLE_DEVICES=0 \
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
python examples/offline_inference/text_to_video/text_to_video.py \
--model robbyant/lingbot-video-moe-30b-a3b \
--prompt "a robotic arm picks up a red block" \
--output lingbot_moe_t2v.mp4 \
+1 -1
View File
@@ -82,7 +82,7 @@ example and the DiT sections below.
- Optional environment variables:
```bash
export CACHE_DIT_VERSION=1.3.0
export CACHE_DIT_VERSION=1.5.0
```
HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA`
+152
View File
@@ -0,0 +1,152 @@
# dots.tts
> Offline continuous-AR TTS at 48 kHz (rednote-hilab)
## Summary
- Vendor: rednote-hilab
- Model: `rednote-hilab/dots.tts-soar`
- Task: Text-to-speech, zero-shot synthesis only
- Mode: Offline end-to-end example only (not yet wired for online serving —
see [Known limitations](#known-limitations))
- Maintainer: Community
## When to use this recipe
Use this recipe as a known-good starting point for running
`rednote-hilab/dots.tts-soar` offline on vLLM-Omni on consumer-class GPUs.
dots.tts is a ~1.7B-parameter continuous-AR TTS model (Qwen2.5-1.5B base LM
+ 344M DiT flow-matching head + 180M AudioVAE) that emits 48 kHz mono audio.
It follows the same "vLLM-native base LM + side-path computation" pattern as
VoxCPM2 — single-stage pipeline
`Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR
loopback → AudioVAE (streaming decode)` — with a plain Qwen2 backbone
instead of MiniCPM4, and no FSQ / residual-LM stage.
This is an early integration (see [Known limitations](#known-limitations)
below before depending on it for anything beyond offline experimentation).
## References
- Offline end-to-end script:
[`examples/offline_inference/text_to_speech/dots_tts/end2end.py`](../../examples/offline_inference/text_to_speech/dots_tts/end2end.py)
- Example guide:
[`examples/offline_inference/text_to_speech/README.md`](../../examples/offline_inference/text_to_speech/README.md#dotstts)
- Default deploy config:
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
- Talker / pipeline source:
[`vllm_omni/model_executor/models/dots_tts/`](../../vllm_omni/model_executor/models/dots_tts/)
- Upstream: [rednote-hilab/dots.tts](https://github.com/rednote-hilab/dots.tts)
## Hardware Support
This recipe documents one tested 16 GB consumer-GPU configuration. Other
vendor sections (ROCm, NPU) and larger-VRAM configurations are welcome as
community validation lands.
## GPU
### 1 x RTX 5080 16GB (Single GPU, Minimum Recommended)
dots.tts (~1.7B params across the base LM + DiT + AudioVAE + CAM++
speaker encoder, bfloat16) fits comfortably on a single 16 GB GPU. The
bundled default config at
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
(`gpu_memory_utilization: 0.8`, `max_num_seqs: 4`, `enforce_eager: true`,
`enable_prefix_caching: false`) loads cleanly with ~5.1 GiB for model
weights and ~0.3 GiB peak activation; the remainder of the configured
budget is available for KV cache. Total resident footprint at idle is
roughly **7-8 GiB / 16 GB** — the only tight spot in the full CUDA-Graph
roadmap would be step 8's graph capture (not implemented yet; this
release runs `enforce_eager: true`, so it doesn't apply today).
#### Environment
- OS: Linux (WSL2)
- Python: 3.12
- Driver / runtime: NVIDIA driver 595.95
- torch: 2.11.0+cu130
- vLLM: 0.26.0
- vLLM-Omni: 0.22.1.dev (current `main`)
#### Command
```bash
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
--model rednote-hilab/dots.tts-soar \
--text "Hello, this is a test of dots TTS running on vLLM Omni."
```
The deploy config at
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
is loaded automatically by the model registry (HF `model_type=dots_tts`).
Pass `--deploy-config <path>` to override.
#### Verification
**T1 — offline zero-shot synthesis**:
```bash
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
--model rednote-hilab/dots.tts-soar \
--text "Hello, this is a test of dots TTS running on vLLM Omni."
```
Observed: `output_audio/output.wav`, 3.52 s @ 48 kHz mono. Single
one-shot process (init → one `generate()` → exit), so the reported
numbers include engine init and first-request warmup, not just steady-state
per-step throughput — same caveat as VoxCPM2's recipe. `Inference: 5.52s`,
`RTF: 1.569`.
Weight-loading breakdown from the same run (all tensors matched, no
missing/extra keys): 951/951 AudioVAE, 244/244 DiT, 270/270
patch_encoder, 198 Qwen2, 938/938 CAM++ speaker encoder.
Whisper transcription of the output matched the input text with no
dropped leading word (confirms the streaming-vocoder patch-boundary fix
described in [Known limitations](#known-limitations)).
#### Notes
- Output: 48 kHz mono WAV.
- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default
used throughout this recipe. `rednote-hilab/dots.tts-base` shares the
same architecture but is unvalidated in this repo. `rednote-hilab/dots.tts-mf`
(MeanFlow, 2-4 step) is not supported — see below.
- `enforce_eager: true` and `enable_prefix_caching: false` in the deploy
config are load-bearing, not just conservative defaults: with prefix
caching enabled, vLLM-Omni's prefix-cache multimodal-output merge path
does not preserve the `sparse_audio` marker this model relies on to
route audio output correctly, and generation silently truncates to a
single ~160 ms patch. Do not override `enable_prefix_caching` for this
model until that framework-level gap is fixed.
## Known limitations
- **No online serving yet.** dots.tts is not registered in
`vllm_omni/entrypoints/openai/tts_adapters/` (the framework other TTS
models in this repo use for `/v1/audio/speech`), so `vllm serve
rednote-hilab/dots.tts-soar --omni` will start but the OpenAI-compatible
speech endpoint does not know how to build request params for it yet.
Offline `Omni()` (as used in `end2end.py`) is the only supported path
today.
- **Voice cloning is not wired.** The CAM++ x-vector speaker encoder
weights load and are exercised by `load_weights()`, but nothing in the
prompt builder consumes a reference audio yet — generation is zero-shot
only. `end2end.py` has no `--ref-audio`/`--ref-text` flags.
- **`dots.tts-mf` (MeanFlow, 2-4 step) checkpoint is not supported.** Only
the fixed 10-step Euler DiT sampler used by `dots.tts-soar` /
`dots.tts-base` is implemented.
- **No CUDA graph capture.** The talker runs fully eager. voxcpm2's three
captured graphs (base LM decode, CFM solver, VAE decode) have no
dots.tts equivalent yet.
- **Concurrent requests do not scale.** Each request's 10-step DiT Euler
integration runs serially in the side path (no cross-request batching,
unlike voxcpm2's `enable_batched_cfm`). A community review of this
integration measured no throughput gain at `c=4` concurrent requests
versus `c=1`.
- **`SamplingParams.seed` does not control audio-generation randomness.**
The DiT's flow-matching noise is deterministically derived per-request
from a fixed internal seed (reproducible run-to-run, matching
voxcpm2's `deterministic_cfm_seed` convention) rather than from the
caller-supplied `seed` field — the same limitation voxcpm2 has today.
+3 -1
View File
@@ -16,7 +16,7 @@ diffusers==0.38.0
safetensors>=0.8.0
accelerate==1.12.0
soundfile>=0.13.1
cache-dit==1.3.0
cache-dit==1.5.0
tqdm>=4.66.0
torchsde>=0.2.6
openai-whisper>=20250625
@@ -34,3 +34,5 @@ pydantic-settings>=2.0.0
# base image does not ship it, so declare it explicitly. The loader also guards
# this import so a missing gguf no longer breaks all diffusion model loading.
gguf>=0.10.0
# Required for running the Cosmos3 model with guardrails enabled (default behavior).
cosmos-guardrail>=0.3.1
+4 -48
View File
@@ -6,7 +6,7 @@ Unit tests for StageConfigFactory and related classes.
import importlib
import warnings
from dataclasses import dataclass, fields
from dataclasses import dataclass
from pathlib import Path
from unittest.mock import patch
@@ -34,9 +34,8 @@ from vllm_omni.config.stage_config import (
load_deploy_config,
merge_pipeline_deploy,
pipeline_cfg_resolver,
strip_parent_engine_args,
)
from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys
from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -475,7 +474,7 @@ class TestStageConfigFactory:
cli_overrides = {
"gpu_memory_utilization": 0.9,
"model": "some_model", # Internal
"stage_configs_path": "/path", # Internal
"deploy_config": "/path", # Internal
"batch_timeout": 10, # Internal
}
@@ -483,7 +482,7 @@ class TestStageConfigFactory:
assert overrides["gpu_memory_utilization"] == 0.9
assert "model" not in overrides
assert "stage_configs_path" not in overrides
assert "deploy_config" not in overrides
assert "batch_timeout" not in overrides
def test_per_stage_override_excludes_internal_keys(self):
@@ -525,49 +524,6 @@ class TestStageResolutionHelpers:
assert "model" not in overrides
assert "parallel_config" not in overrides
def test_strip_parent_engine_args_reports_only_surprising_parent_overrides(self):
parent_fields = {f.name: f for f in fields(EngineArgs)}
filtered, overridden = strip_parent_engine_args(
{
"model": "some/model",
"stage_configs_path": "/tmp/stages.yaml",
"tensor_parallel_size": 4,
"worker_extension_cls": "some.Extension",
"custom_pipeline_args": {"pipeline_class": "demo.Pipeline"},
},
parent_fields=parent_fields,
keep_keys={"worker_extension_cls"},
strip_keys={"stage_configs_path"},
no_warn_keys={"model"},
)
assert filtered == {
"worker_extension_cls": "some.Extension",
"custom_pipeline_args": {"pipeline_class": "demo.Pipeline"},
}
assert overridden == ["tensor_parallel_size"]
def test_strip_parent_engine_args_keeps_allowed_media_access_controls(self):
parent_fields = {f.name: f for f in fields(EngineArgs)}
filtered, overridden = strip_parent_engine_args(
{
"model": "some/model",
"stage_configs_path": "/tmp/stages.yaml",
"allowed_local_media_path": "/data/qwentts",
"allowed_media_domains": ["example.com"],
},
parent_fields=parent_fields,
keep_keys={"allowed_local_media_path", "allowed_media_domains"},
strip_keys={"stage_configs_path"},
no_warn_keys={"model"},
)
assert filtered == {
"allowed_local_media_path": "/data/qwentts",
"allowed_media_domains": ["example.com"],
}
assert overridden == []
class TestPipelineDiscovery:
"""Tests for the central pipeline registry (``OMNI_PIPELINES``)."""
+1
View File
@@ -471,6 +471,7 @@ def test_sub_config_fields_match_structured_scopes():
"trust_remote_code",
"dtype",
"attention_backend",
"attention_config",
"moe_backend",
"hf_overrides",
"limit_mm_per_prompt",
@@ -0,0 +1,82 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import torch.nn.functional as F
from vllm.triton_utils import HAS_TRITON
pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion]
_HEAD_DIM = 128
_ROTARY_DIM = 96
_EPS = 1e-5
def _reference(q, k, q_weight, k_weight, rope_table):
q = F.rms_norm(q, (_HEAD_DIM,), q_weight, _EPS)
k = F.rms_norm(k, (_HEAD_DIM,), k_weight, _EPS)
half = _ROTARY_DIM // 2
cos = rope_table[..., :half].unsqueeze(1)
sin = rope_table[..., half:].unsqueeze(1)
def apply(x):
first = x[..., :half]
second = x[..., half:_ROTARY_DIM]
return torch.cat(
(
first * cos - second * sin,
second * cos + first * sin,
x[..., _ROTARY_DIM:],
),
dim=-1,
)
return apply(q), apply(k)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
@pytest.mark.parametrize("seq_len", [1, 257, 1024])
def test_fused_qk_norm_rope_matches_bf16_reference(seq_len):
from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
fused_qk_norm_rope,
)
torch.manual_seed(17)
heads = 14
qkv = torch.randn(
seq_len,
heads * _HEAD_DIM * 3,
device="cuda",
dtype=torch.bfloat16,
)
q = qkv[:, : heads * _HEAD_DIM].view(seq_len, heads, _HEAD_DIM)
k = qkv[:, heads * _HEAD_DIM : 2 * heads * _HEAD_DIM].view(
seq_len,
heads,
_HEAD_DIM,
)
q_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
k_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
freqs = torch.randn(seq_len, _ROTARY_DIM // 2, device="cuda")
rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16)
expected_q, expected_k = _reference(
q,
k,
q_weight,
k_weight,
rope_table,
)
actual_q, actual_k = fused_qk_norm_rope(
q,
k,
q_weight,
k_weight,
rope_table,
_EPS,
)
torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)
@@ -23,7 +23,7 @@ def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1):
return pipeline
def test_prepare_video_latents_samples_directly_in_packed_token_space():
def test_prepare_video_latents_matches_official_values_and_token_major_layout():
pipeline = _make_pipeline(LTX2Pipeline)
pipeline.vae_spatial_compression_ratio = 8
pipeline.vae_temporal_compression_ratio = 8
@@ -44,6 +44,7 @@ def test_prepare_video_latents_samples_directly_in_packed_token_space():
)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
assert actual.stride()[1:] == (1, actual.shape[1])
def test_prepare_audio_latents_samples_directly_in_packed_token_space():
@@ -81,6 +82,42 @@ def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
assert original_num_frames == 1
assert padded_num_frames == 2
assert latents.shape == (1, 2, 128)
torch.testing.assert_close(latents[:, 1:], torch.zeros_like(latents[:, 1:]))
def test_prepare_audio_latents_request_rng_is_invariant_to_sp_padding():
pipeline_sp1 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=1)
pipeline_sp4 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
generator_sp1 = torch.Generator().manual_seed(42)
generator_sp4 = torch.Generator().manual_seed(42)
latents_sp1, _, _ = pipeline_sp1.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=8,
audio_latent_length=3,
dtype=torch.float32,
device=torch.device("cpu"),
generator=generator_sp1,
)
latents_sp4, _, _ = pipeline_sp4.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=8,
audio_latent_length=3,
dtype=torch.float32,
device=torch.device("cpu"),
generator=generator_sp4,
)
torch.testing.assert_close(latents_sp4[:, :3], latents_sp1, rtol=0, atol=0)
torch.testing.assert_close(latents_sp4[:, 3:], torch.zeros_like(latents_sp4[:, 3:]))
torch.testing.assert_close(
torch.randn(8, generator=generator_sp4),
torch.randn(8, generator=generator_sp1),
rtol=0,
atol=0,
)
def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
@@ -137,7 +174,9 @@ def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
assert original_num_frames == 10
assert padded_num_frames == 12
assert padded.shape == (1, 12, 8)
torch.testing.assert_close(padded, latent_ops.pack_audio_latents(latents))
packed = latent_ops.pack_audio_latents(latents)
torch.testing.assert_close(padded[:, :10], packed[:, :10])
torch.testing.assert_close(padded[:, 10:], torch.zeros_like(padded[:, 10:]))
def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
@@ -12,6 +12,19 @@ import torch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker):
from vllm_omni.diffusion.models.ltx2.ltx2_guidance import velocity_from_x0
sigma = mocker.MagicMock()
sigma.to.return_value.item.return_value = 0.5
actual = velocity_from_x0(torch.tensor([1.0]), torch.tensor([0.0]), sigma)
sigma.to.assert_called_once_with(torch.float32)
sigma.to.return_value.item.assert_called_once_with()
torch.testing.assert_close(actual, torch.tensor([2.0]))
class TestCFGParallelHelpers:
"""Test LTX-2.3 CFG helper math without loading model weights."""
@@ -199,6 +212,7 @@ class TestCFGParallelHelpers:
),
attention_kwargs=None,
audio_scheduler=SimpleNamespace(sigmas=torch.stack([audio_sigma])),
original_audio_num_frames=state.audio.shape[1],
)
actual_video, actual_audio = executor.predict_parallel_guidance(
@@ -0,0 +1,483 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the fixed LTX refinement-phase adapter."""
from types import SimpleNamespace
import pytest
import torch
from safetensors.torch import save_file
from vllm_omni.diffusion.models.ltx2 import ltx2_phase_adapter
from vllm_omni.diffusion.models.ltx2.ltx2_adapter_parser import (
AdapterTarget,
parse_ltx_adapter,
)
from vllm_omni.diffusion.models.ltx2.ltx2_components import LTX23_TWO_STAGE_COMPONENT_PROFILE
from vllm_omni.diffusion.models.ltx2.ltx2_phase_adapter import LTXPhaseAdapterRuntime
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
LTX2_TWO_STAGE_RECIPE,
LTX_DISTILLED_ADAPTER_SLOT,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _target(
lora_a,
lora_b,
*,
source_module="proj",
module="proj",
target_slice=None,
):
return AdapterTarget(
source_module=source_module,
module=module,
slice=target_slice,
a_key="a",
b_key="b",
a_shape=tuple(lora_a.shape),
b_shape=tuple(lora_b.shape),
)
def _wrap(
layer,
lora_a,
lora_b,
*,
source_module="proj",
module="proj",
target_slice=None,
layer_fused=False,
):
target = _target(
lora_a,
lora_b,
source_module=source_module,
module=module,
target_slice=target_slice,
)
wrapper = ltx2_phase_adapter._PhaseAdapterLinear(
layer,
[target],
adapter_dtype=lora_a.dtype,
layer_fused=layer_fused,
)
wrapper.load_target(target, lora_a, lora_b)
wrapper.set_enabled(True)
return wrapper
class _LoRAAttention(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.to_qkv = torch.nn.Identity()
class _LoRATransformer(torch.nn.Module):
stacked_params_mapping = ((".attn1.to_qkv", ".attn1.to_q", "q"),)
def __init__(self) -> None:
super().__init__()
self.proj_in = torch.nn.Linear(4, 6, bias=False)
block = torch.nn.Module()
block.attn1 = _LoRAAttention()
self.transformer_blocks = torch.nn.ModuleList([block])
def test_ltx_adapter_parser_maps_official_and_packed_module_names(tmp_path):
adapter_path = tmp_path / "distilled.safetensors"
save_file(
{
"diffusion_model.text_embedding_projection.aggregate_embed.lora_A.weight": torch.randn(2, 4),
"diffusion_model.patchify_proj.lora_A.weight": torch.randn(2, 4),
"diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2),
"diffusion_model.transformer_blocks.0.attn1.to_q.lora_A.weight": torch.randn(2, 4),
"diffusion_model.transformer_blocks.0.attn1.to_q.lora_B.weight": torch.randn(4, 2),
},
adapter_path,
)
manifest = parse_ltx_adapter(_LoRATransformer(), str(adapter_path))
assert {(target.source_module, target.module, target.slice) for target in manifest.targets} == {
("proj_in", "proj_in", None),
("transformer_blocks.0.attn1.to_q", "transformer_blocks.0.attn1.to_qkv", "q"),
}
def test_ltx_adapter_parser_rejects_unpaired_b_tensor(tmp_path):
adapter_path = tmp_path / "distilled.safetensors"
save_file(
{"diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2)},
adapter_path,
)
with pytest.raises(ValueError, match="Missing paired LoRA A tensors"):
parse_ltx_adapter(_LoRATransformer(), str(adapter_path))
def _pipeline(*, dtype, quantization_config):
return SimpleNamespace(
pipeline_recipe=LTX2_TWO_STAGE_RECIPE,
component_profile=LTX23_TWO_STAGE_COMPONENT_PROFILE,
transformer=object(),
od_config=SimpleNamespace(
model="unused",
lora_path=None,
dtype=dtype,
quantization_config=quantization_config,
),
)
@pytest.mark.parametrize(
("dtype", "quantization_config", "layer_fused"),
[
(torch.float32, object(), False),
(torch.bfloat16, None, True),
],
)
def test_ltx_phase_adapter_selects_execution_from_quantization(
monkeypatch,
dtype,
quantization_config,
layer_fused,
):
pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config)
manifest = object()
captured = {}
def resolve_artifact(*_args, **kwargs):
assert not kwargs
return "adapter"
def parse(transformer, path):
assert transformer is pipeline.transformer
assert path == "adapter"
return manifest
class Runtime:
def __init__(self, transformer, parsed_manifest, **kwargs):
captured.update(transformer=transformer, manifest=parsed_manifest, **kwargs)
monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", resolve_artifact)
monkeypatch.setattr(ltx2_phase_adapter, "parse_ltx_adapter", parse)
monkeypatch.setattr(ltx2_phase_adapter, "LTXPhaseAdapterRuntime", Runtime)
runtime = ltx2_phase_adapter.build_ltx_phase_adapter(pipeline)
assert isinstance(runtime, Runtime)
assert captured == {
"transformer": pipeline.transformer,
"manifest": manifest,
"dtype": dtype,
"layer_fused": layer_fused,
}
def test_ltx_phase_adapter_rejects_non_bfloat16_layer_fused(monkeypatch):
pipeline = _pipeline(dtype=torch.float16, quantization_config=None)
monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", lambda *_args, **_kwargs: "adapter")
with pytest.raises(ValueError, match="bfloat16"):
ltx2_phase_adapter.build_ltx_phase_adapter(pipeline)
def test_ltx_phase_adapter_keeps_one_transformer_and_switches_a_fixed_slot(tmp_path):
class TinyTransformer(torch.nn.Module):
stacked_params_mapping = ()
def __init__(self) -> None:
super().__init__()
self.proj = torch.nn.Linear(2, 2, bias=False)
adapter_path = tmp_path / "distilled.safetensors"
save_file(
{
"diffusion_model.proj.lora_A.weight": torch.tensor([[1.0, 2.0]]),
"diffusion_model.proj.lora_B.weight": torch.tensor([[3.0], [4.0]]),
},
adapter_path,
)
transformer = TinyTransformer()
runtime = LTXPhaseAdapterRuntime(
transformer,
parse_ltx_adapter(transformer, str(adapter_path)),
dtype=torch.float32,
)
assert runtime.base_parameter_name("proj.weight") == "proj.base_layer.weight"
assert {name for name, _ in transformer.named_parameters()} == {"proj.base_layer.weight"}
with torch.no_grad():
transformer.proj.base_layer.weight.copy_(torch.eye(2))
runtime.finalize()
x = torch.tensor([[1.0, 1.0]])
runtime.activate(None)
torch.testing.assert_close(transformer.proj(x), torch.tensor([[1.0, 1.0]]))
runtime.activate(LTX_DISTILLED_ADAPTER_SLOT)
torch.testing.assert_close(transformer.proj(x), torch.tensor([[10.0, 13.0]]))
def test_ltx_layer_fused_adapter_matches_official_bf16_weight_fusion():
layer = torch.nn.Linear(4, 3, bias=True, dtype=torch.bfloat16)
base_weight = torch.tensor(
[
[0.5, -0.25, 0.125, 0.75],
[-0.375, 0.625, 0.25, -0.5],
[0.875, -0.125, -0.75, 0.375],
],
dtype=torch.bfloat16,
)
lora_a = torch.tensor(
[
[0.1015625, -0.203125, 0.3046875, -0.40625],
[0.5078125, -0.609375, 0.7109375, -0.8125],
],
dtype=torch.bfloat16,
)
lora_b = torch.tensor(
[[0.9140625, -0.1171875], [-0.21875, 0.3203125], [0.421875, -0.5234375]],
dtype=torch.bfloat16,
)
with torch.no_grad():
layer.weight.copy_(base_weight)
layer.bias.copy_(torch.tensor([0.125, -0.25, 0.5], dtype=torch.bfloat16))
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
expected_weight = torch.matmul(lora_b, lora_a)
expected_weight.add_(base_weight)
x = torch.tensor([[0.6015625, -0.703125, 0.8046875, -0.90625]], dtype=torch.bfloat16)
expected = torch.nn.functional.linear(x, expected_weight, layer.bias)
assert torch.equal(wrapper(x), expected)
assert torch.equal(layer.weight, base_weight)
def test_ltx_layer_fused_row_parallel_uses_one_collective(monkeypatch):
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
class FakeRow(torch.nn.Module):
input_size = 4
output_size = 2
input_size_per_partition = 2
input_is_parallel = True
tp_rank = 0
tp_size = 2
reduce_results = True
skip_bias_add = False
return_bias = False
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(torch.eye(2, dtype=torch.bfloat16))
self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16))
self.quant_method = UnquantizedLinearMethod()
monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", FakeRow)
reductions = []
def all_reduce(value):
reductions.append(value.clone())
return value + 1
monkeypatch.setattr(ltx2_phase_adapter, "tensor_model_parallel_all_reduce", all_reduce)
layer = FakeRow()
lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16)
lora_b = torch.tensor([[0.25], [-0.5]], dtype=torch.bfloat16)
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
expected_weight = torch.matmul(lora_b, lora_a[:, :2])
expected_weight.add_(layer.weight)
x = torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16)
expected_local = torch.nn.functional.linear(x, expected_weight, layer.bias)
actual = wrapper(x)
assert len(reductions) == 1
assert torch.equal(reductions[0], expected_local)
assert torch.equal(actual, expected_local + 1)
def test_ltx_layer_fused_column_parallel_uses_local_weight_shard(monkeypatch):
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
class FakeColumn(torch.nn.Module):
input_size = 3
output_size = 4
output_size_per_partition = 2
tp_rank = 1
tp_size = 2
gather_output = False
skip_bias_add = False
return_bias = False
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(
torch.tensor([[0.5, -0.25, 0.125], [-0.375, 0.625, 0.25]], dtype=torch.bfloat16)
)
self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16))
self.quant_method = UnquantizedLinearMethod()
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn)
layer = FakeColumn()
base_weight = layer.weight.detach().clone()
lora_a = torch.tensor([[0.5, -0.25, 0.75]], dtype=torch.bfloat16)
lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16)
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
expected_weight = torch.matmul(lora_b[2:], lora_a)
expected_weight.add_(base_weight)
x = torch.tensor([[1.0, -0.5, 0.25]], dtype=torch.bfloat16)
expected = torch.nn.functional.linear(x, expected_weight, layer.bias)
assert torch.equal(wrapper(x), expected)
assert torch.equal(layer.weight, base_weight)
def test_ltx_layer_fused_packed_qkv_updates_only_the_target_slice(monkeypatch):
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
class FakeColumn(torch.nn.Module):
pass
class FakeQKV(FakeColumn):
input_size = 4
head_size = 2
v_head_size = 2
total_num_heads = 2
total_num_kv_heads = 2
num_heads = 1
num_kv_heads = 1
num_kv_head_replicas = 1
tp_rank = 1
tp_size = 2
gather_output = False
skip_bias_add = True
return_bias = True
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(
torch.arange(24, dtype=torch.float32).reshape(6, 4).to(torch.bfloat16) / 16
)
self.bias = torch.nn.Parameter(torch.arange(6, dtype=torch.float32).to(torch.bfloat16) / 8)
self.quant_method = UnquantizedLinearMethod()
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn)
monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", FakeQKV)
layer = FakeQKV()
base_weight = layer.weight.detach().clone()
lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16)
lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16)
wrapper = _wrap(
layer,
lora_a,
lora_b,
source_module="to_k",
module="to_qkv",
target_slice="k",
layer_fused=True,
)
expected_weight = base_weight.clone()
expected_k = torch.matmul(lora_b[2:], lora_a)
expected_k.add_(base_weight[2:4])
expected_weight[2:4].copy_(expected_k)
x = torch.tensor([[1.0, -0.5, 0.25, -0.125]], dtype=torch.bfloat16)
expected_output = torch.nn.functional.linear(x, expected_weight)
actual_output, returned_bias = wrapper(x)
assert torch.equal(actual_output, expected_output)
assert returned_bias is layer.bias
assert torch.equal(layer.weight, base_weight)
@pytest.mark.parametrize(
("layer_factory", "target_slice", "expected_a", "expected_b", "expected_output"),
[
(
lambda: type(
"FakeColumn",
(torch.nn.Module,),
{
"input_size": 4,
"output_size": 4,
"output_size_per_partition": 2,
"tp_rank": 1,
"gather_output": False,
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)),
},
)(),
None,
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
torch.tensor([[1.0], [2.0], [3.0], [4.0]]),
torch.tensor([[12963.0, 17284.0]]),
),
(
lambda: type(
"FakeRow",
(torch.nn.Module,),
{
"input_size": 4,
"output_size": 2,
"input_size_per_partition": 2,
"tp_rank": 1,
"reduce_results": False,
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)),
},
)(),
None,
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
torch.tensor([[2.0], [3.0]]),
torch.tensor([[8600.0, 12900.0]]),
),
(
lambda: type(
"FakeQKV",
(torch.nn.Module,),
{
"input_size": 4,
"head_size": 2,
"v_head_size": 2,
"total_num_heads": 2,
"total_num_kv_heads": 2,
"num_heads": 1,
"num_kv_heads": 1,
"num_kv_head_replicas": 1,
"tp_rank": 1,
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 6)),
},
)(),
"k",
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
torch.tensor([[1.0], [2.0], [3.0], [4.0]]),
torch.tensor([[0.0, 0.0, 12963.0, 17284.0, 0.0, 0.0]]),
),
],
)
def test_ltx_phase_adapter_uses_local_tp_lora_slices(
monkeypatch,
layer_factory,
target_slice,
expected_a,
expected_b,
expected_output,
):
layer = layer_factory()
layer.register_buffer("device_anchor", torch.empty(0))
layer_type = type(layer)
if target_slice == "k":
monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", layer_type)
elif "Column" in layer_type.__name__:
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", layer_type)
else:
monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", layer_type)
wrapper = _wrap(layer, expected_a, expected_b, target_slice=target_slice)
torch.testing.assert_close(wrapper(torch.tensor([[1.0, 2.0, 3.0, 4.0]])), expected_output)
expected_local_a_width = getattr(layer, "input_size_per_partition", layer.input_size)
assert wrapper.pieces[0].lora_a.shape[-1] == expected_local_a_width
+362 -10
View File
@@ -13,13 +13,21 @@ import pytest
import torch
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.models.ltx2 import ltx2_components
from vllm_omni.diffusion.models.ltx2 import (
ltx2_components,
ltx2_transformer,
)
from vllm_omni.diffusion.models.ltx2.ltx2_components import (
LTX2_COMPONENT_PROFILE,
LTX2_DISTILLED_COMPONENT_PROFILE,
LTX2_TWO_STAGE_COMPONENT_PROFILE,
LTX23_COMPONENT_PROFILE,
LTX23_DISTILLED_COMPONENT_PROFILE,
LTX23_TWO_STAGE_COMPONENT_PROFILE,
_install_connector_attention,
detect_ltx_model_version,
resolve_ltx_artifact,
resolve_ltx_component_profile,
)
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
@@ -33,6 +41,7 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
LTXGuidancePlan,
LTXGuidanceSpec,
LTXModalityGuidance,
_repeat_batch,
build_perturbation_kwargs,
combine_guided_x0,
)
@@ -40,18 +49,29 @@ from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
LTX2_DISTILLED_TWO_STAGE_RECIPE,
LTX2_ONE_STAGE_RECIPE,
LTX2_TWO_STAGE_RECIPE,
LTX23_DISTILLED_TWO_STAGE_RECIPE,
LTX23_ONE_STAGE_RECIPE,
LTX23_TWO_STAGE_RECIPE,
LTX_DEFAULT_NEGATIVE_PROMPT,
LTX_POSITIVE_ONLY_RECIPE,
resolve_ltx_pipeline_recipe,
)
from vllm_omni.diffusion.models.ltx2.ltx2_request import (
LTXRequestInputs,
validate_ltx_checkpoint,
validate_pipeline_request,
)
from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs, validate_pipeline_request
from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2I2VDMD2Pipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
LTX2DistilledPipeline,
LTX2TwoStagePipeline,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -104,10 +124,22 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries():
)
assert issubclass(LTX2Pipeline, LTXRuntime)
assert issubclass(LTX2TwoStagePipeline, LTXRuntime)
assert issubclass(LTX2DistilledPipeline, LTXRuntime)
assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE
assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE
assert LTX2TwoStagePipeline.component_profile is LTX2_TWO_STAGE_COMPONENT_PROFILE
assert LTX2TwoStagePipeline.pipeline_recipe is LTX2_TWO_STAGE_RECIPE
assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE
for pipeline_cls, profile in (
(LTX2Pipeline, LTX2_COMPONENT_PROFILE),
(LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE),
(LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE),
):
assert pipeline_cls._dit_modules == list(profile.dit_modules)
assert pipeline_cls._encoder_modules == list(profile.encoder_modules)
assert pipeline_cls._vae_modules == list(profile.vae_modules)
assert pipeline_cls._resident_modules == list(profile.resident_modules)
assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
@@ -123,6 +155,72 @@ def test_ltx_checkpoint_version_detection_uses_metadata(tmp_path):
assert detect_ltx_model_version(str(tmp_path)) == "2"
def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch):
filename = "ltx-sidecar.safetensors"
calls = []
def fake_download(**kwargs):
calls.append(kwargs)
return "/cache/ltx-sidecar.safetensors"
monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
assert (
resolve_ltx_artifact(str(tmp_path / "model"), "Lightricks/test", filename) == "/cache/ltx-sidecar.safetensors"
)
assert calls == [{"repo_id": "Lightricks/test", "filename": filename}]
def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
filename = "ltx-sidecar.safetensors"
expected = tmp_path / filename
expected.write_bytes(b"sidecar")
monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup"))
assert resolve_ltx_artifact(str(tmp_path), "Lightricks/test", filename) == str(expected)
@pytest.mark.parametrize(
("expected_kind", "scheduler_config"),
[
("regular", {"use_dynamic_shifting": True, "shift_terminal": 0.1}),
("distilled", {"use_dynamic_shifting": False, "shift_terminal": None}),
],
)
def test_ltx_checkpoint_validation_accepts_matching_scheduler_metadata(expected_kind, scheduler_config):
validate_ltx_checkpoint(
scheduler_config,
expected_kind=expected_kind,
pipeline_name="TestLTXPipeline",
)
@pytest.mark.parametrize(
("expected_kind", "scheduler_config", "error"),
[
("regular", {"use_dynamic_shifting": False, "shift_terminal": None}, "regular non-distilled"),
("distilled", {"use_dynamic_shifting": True, "shift_terminal": 0.1}, "merged distilled"),
],
)
def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expected_kind, scheduler_config, error):
with pytest.raises(ValueError, match=error):
validate_ltx_checkpoint(
scheduler_config,
expected_kind=expected_kind,
pipeline_name="TestLTXPipeline",
)
def test_ltx_rejects_advanced_uaa_before_component_initialization():
od_config = SimpleNamespace(
model="unused",
parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"),
)
with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"):
LTX2Pipeline(od_config=od_config)
def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
@@ -144,16 +242,104 @@ def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, m
assert pipe.reports_stage_durations
def test_ltx_one_stage_entry_exposes_batch_image_and_distributed_decode_contracts():
def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]}))
initialized = False
def stub_components(*_args, **_kwargs):
nonlocal initialized
initialized = True
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
with pytest.raises(ValueError, match="only one-stage LTX recipes"):
LTX2TwoStagePipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
cache_backend="cache_dit",
)
)
assert not initialized
def test_ltx23_checkpoint_selects_version_specific_two_stage_profiles(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]}))
def stub_components(pipe, od_config):
pipe.od_config = od_config
pipe.vae_spatial_compression_ratio = 32
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda *_args: SimpleNamespace())
pipe = LTX2TwoStagePipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
enable_diffusion_pipeline_profiler=False,
lora_path=None,
model_config={},
)
)
assert pipe.model_version == "2.3"
assert pipe.component_profile is LTX23_TWO_STAGE_COMPONENT_PROFILE
assert pipe.pipeline_recipe is LTX23_TWO_STAGE_RECIPE
def test_ltx23_checkpoint_selects_full_distilled_profile_and_recipe(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]}))
def stub_components(pipe, od_config):
pipe.od_config = od_config
pipe.vae_spatial_compression_ratio = 32
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
pipe = LTX2DistilledPipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
enable_diffusion_pipeline_profiler=False,
)
)
assert pipe.model_version == "2.3"
assert pipe.component_profile is LTX23_DISTILLED_COMPONENT_PROFILE
assert pipe.pipeline_recipe is LTX23_DISTILLED_TWO_STAGE_RECIPE
assert pipe._phase_adapter is None
assert pipe.preserve_sp_padded_audio_duration
assert pipe.reports_stage_durations
def test_ltx_entries_expose_batch_image_and_distributed_decode_contracts():
assert LTX2Pipeline.supports_request_batch
assert LTX2Pipeline.support_image_input
assert LTX2Pipeline.unified_text_image_entry
assert LTX2Pipeline.distributed_video_decode
assert not LTX2TwoStagePipeline.supports_request_batch
assert LTX2TwoStagePipeline.support_image_input
assert LTX2TwoStagePipeline.unified_text_image_entry
assert not LTX2DistilledPipeline.supports_request_batch
assert LTX2DistilledPipeline.support_image_input
assert LTX2DistilledPipeline.unified_text_image_entry
def test_ltx_two_stage_resolves_request_embedded_images():
pipe = object.__new__(LTX2TwoStagePipeline)
request_inputs = SimpleNamespace(latents=None)
image = torch.zeros(3, 8, 8)
req = SimpleNamespace(prompts=[{"prompt": "image conditioned", "multi_modal_data": {"image": image}}])
assert pipe._resolve_request_image(req, None, request_inputs) is image
def test_ltx_dmd2_entries_retain_their_task_boundaries():
assert not LTX2T2VDMD2Pipeline.support_image_input
assert LTX2I2VDMD2Pipeline.support_image_input
@@ -330,6 +516,15 @@ def test_ltx_guidance_uses_official_close_to_disabled_semantics():
assert LTXGuidancePlan.build(LTXGuidanceSpec(video=guidance)).names == ("cond",)
def test_ltx_positive_only_guidance_preserves_token_major_layout():
tensor = torch.arange(24).view(1, 3, 8).transpose(1, 2).contiguous().transpose(1, 2)
repeated = _repeat_batch(tensor, 1)
assert repeated is tensor
assert repeated.stride() == tensor.stride()
def test_ltx_guidance_rescale_is_invariant_to_request_batching():
cond = torch.tensor(
[
@@ -380,6 +575,32 @@ def test_ltx_guidance_rescale_handles_zero_variance_prediction():
torch.testing.assert_close(actual, torch.zeros_like(cond))
def test_ltx_audio_guidance_rescale_ignores_sp_padding_tokens():
logical_cond = torch.tensor([[[1.0, 2.0], [3.0, 5.0]]])
logical_uncond = torch.tensor([[[0.5, -1.0], [2.0, 1.0]]])
guidance = LTXModalityGuidance(cfg_scale=3.0, rescale_scale=1.0)
expected = combine_guided_x0(
cond=logical_cond,
uncond_text=logical_uncond,
uncond_perturbed=0.0,
uncond_modality=0.0,
guidance=guidance,
)
padded_cond = torch.cat([logical_cond, torch.tensor([[[1000.0, -1000.0]]])], dim=1)
padded_uncond = torch.cat([logical_uncond, torch.tensor([[[-500.0, 500.0]]])], dim=1)
actual = combine_guided_x0(
cond=padded_cond,
uncond_text=padded_uncond,
uncond_perturbed=0.0,
uncond_modality=0.0,
guidance=guidance,
rescale_token_count=2,
)
torch.testing.assert_close(actual[:, :2], expected)
def test_ltx_guidance_builds_per_sample_stg_and_modality_masks():
plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance)
perturbations = build_perturbation_kwargs(plan, batch_size=2, reference=torch.ones(8, 1, 1))
@@ -547,11 +768,19 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
def test_ltx2_distilled_pipeline_shares_recipe_runtime():
assert issubclass(LTX2DistilledPipeline, LTXI2VConditioningMixin)
assert issubclass(LTX2DistilledPipeline, LTX2Pipeline)
assert not issubclass(LTX2DistilledPipeline, LTX2Pipeline)
assert LTX2DistilledPipeline._run_recipe is LTXRuntime._run_recipe
assert "_forward_request" not in LTX2DistilledPipeline.__dict__
def test_ltx_two_stage_entry_is_a_thin_pipeline_variant():
assert "_run_recipe" not in LTX2TwoStagePipeline.__dict__
assert "_forward_request" not in LTX2TwoStagePipeline.__dict__
assert "load_weights" not in LTX2TwoStagePipeline.__dict__
assert issubclass(LTX2TwoStagePipeline, LTXI2VConditioningMixin)
assert not issubclass(LTX2TwoStagePipeline, LTX2Pipeline)
def test_ltx_variants_share_denoise_loop_and_i2v_conditioning():
assert "_denoise_loop" not in LTX2Pipeline.__dict__
assert issubclass(LTX2Pipeline, LTXI2VConditioningMixin)
@@ -594,7 +823,8 @@ def test_denoise_executor_owns_progress_and_interrupt():
torch.testing.assert_close(state.audio, torch.tensor(11.0))
def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
request_inputs = LTXRequestInputs(
prompt="prompt",
negative_prompt="",
@@ -671,12 +901,15 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
prompt_context_sentinel = prompt_context
pipeline = object.__new__(LTX2DistilledPipeline)
pipeline = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8
pipeline.latent_upsampler = FakeUpsampler()
activated_slots = []
if pipeline_cls is LTX2TwoStagePipeline:
pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append)
object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
object.__setattr__(pipeline, "run_phase", run_phase)
object.__setattr__(pipeline, "decode_phase", decode_phase)
@@ -686,14 +919,18 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
assert len(phase_calls) == 2
assert phase_calls[1][2] is prompt_context_sentinel
if pipeline_cls is LTX2TwoStagePipeline:
assert activated_slots == [None, "ltx_distilled"]
torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0))
torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0))
expected_audio = 2.0 if pipeline_cls is LTX2TwoStagePipeline else 4.0
torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), expected_audio))
def test_ltx2_distilled_stage2_reapplies_final_resolution_i2v_conditioning():
@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
def test_ltx_two_stage_stage2_reapplies_final_resolution_i2v_conditioning(pipeline_cls):
from vllm_omni.diffusion.models.ltx2 import ltx2_latents
pipeline = object.__new__(LTX2DistilledPipeline)
pipeline = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipeline)
pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8
@@ -750,6 +987,18 @@ def test_ltx2_distilled_two_stage_recipe_is_fixed_positive_only():
assert (LTX2_DISTILLED_TWO_STAGE_RECIPE.height, LTX2_DISTILLED_TWO_STAGE_RECIPE.width) == (1024, 1536)
def test_ltx23_full_distilled_route_uses_merged_weights_without_lora():
assert resolve_ltx_component_profile("distilled_two_stage", "2.3") is LTX23_DISTILLED_COMPONENT_PROFILE
assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.3") is LTX23_DISTILLED_TWO_STAGE_RECIPE
assert LTX23_DISTILLED_TWO_STAGE_RECIPE is LTX2_DISTILLED_TWO_STAGE_RECIPE
assert all(phase.adapter_slot is None for phase in LTX23_DISTILLED_TWO_STAGE_RECIPE.phases)
assert LTX23_DISTILLED_COMPONENT_PROFILE.distilled_lora_filename is None
assert LTX23_DISTILLED_COMPONENT_PROFILE.vocoder_cls is LTX23_COMPONENT_PROFILE.vocoder_cls
assert LTX23_DISTILLED_COMPONENT_PROFILE.latent_upsampler_filename == (
"ltx-2.3-spatial-upscaler-x2-1.1.safetensors"
)
@pytest.mark.parametrize(
("request_inputs", "error"),
[
@@ -958,6 +1207,98 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting():
torch.testing.assert_close(timesteps, expected[:-1] * scheduler.config.num_train_timesteps)
@pytest.mark.parametrize(
("recipe", "steps", "stg_block"),
[
(LTX2_TWO_STAGE_RECIPE, 40, 29),
(LTX23_TWO_STAGE_RECIPE, 30, 28),
],
)
def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, steps, stg_block):
stage1, stage2 = recipe.phases
assert (recipe.width, recipe.height) == (1536, 1024)
assert recipe.num_inference_steps == steps
assert stage1.spatial_downscale == 2
assert stage1.adapter_slot is None
assert stage1.guidance.video.stg_blocks == (stg_block,)
assert stage1.guidance.audio.stg_blocks == (stg_block,)
assert stage2.input_transform == "spatial_upsample"
assert stage2.adapter_slot == "ltx_distilled"
assert stage2.guidance == LTXGuidanceSpec.positive_only()
assert stage2.num_inference_steps == 3
assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0)
# Official Stage 2 refines video only; its audio result is discarded.
assert (recipe.video_output_phase, recipe.audio_output_phase) == (1, 0)
def test_ltx_two_stage_entries_select_the_official_adapter_slots():
phase_switches = []
ordinary_pipeline = object.__new__(LTX2TwoStagePipeline)
ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append)
for phase in LTX2_TWO_STAGE_RECIPE.phases:
ordinary_pipeline._enter_phase(phase)
assert phase_switches == [None, "ltx_distilled"]
distilled_pipeline = object.__new__(LTX2DistilledPipeline)
for phase in LTX2_DISTILLED_TWO_STAGE_RECIPE.phases:
distilled_pipeline._enter_phase(phase)
def test_ltx23_gate_projection_declares_its_global_column_parallel_layout(monkeypatch):
class FakeColumn(torch.nn.Module):
def __init__(
self,
input_size,
output_size,
*,
bias,
gather_output,
return_bias,
quant_config,
prefix,
):
super().__init__()
self.input_size = input_size
self.output_size = output_size
self.bias_enabled = bias
self.gather_output = gather_output
self.return_bias = return_bias
self.quant_config = quant_config
self.prefix = prefix
class FakeRow(torch.nn.Module):
def __init__(self, *args, **kwargs):
super().__init__()
class FakeAttention(torch.nn.Module):
def __init__(self, **kwargs):
super().__init__()
monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumn)
monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRow)
monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeAttention)
monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2)
attention = ltx2_transformer.LTX2Attention(
query_dim=16,
heads=4,
dim_head=4,
qk_norm="rms_norm",
apply_gated_attention=True,
pack_qkv=False,
prefix="transformer_blocks.0.attn1",
)
gate = attention.to_gate_logits
assert isinstance(gate, FakeColumn)
assert gate.input_size == 16
assert gate.output_size == 4
assert gate.gather_output is False
assert gate.return_bias is False
assert gate.quant_config is None
assert gate.prefix == "transformer_blocks.0.attn1.to_gate_logits"
class TestLTXRequestParsing:
def test_request_resolves_official_and_explicit_negative_prompts_per_sample(self):
from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -1432,6 +1773,7 @@ class TestPipelineComponents:
def test_distilled_profile_keeps_upsampler_in_managed_resident_components(self):
assert LTX2_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler")
assert LTX23_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler")
class TestLTX23DecodeConditioning:
@@ -1477,6 +1819,11 @@ class TestRegistryIntegration:
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
assert _DIFFUSION_MODELS["LTX2Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX2Pipeline")
assert _DIFFUSION_MODELS["LTX2TwoStagePipeline"] == (
"ltx2",
"pipeline_ltx2_two_stage",
"LTX2TwoStagePipeline",
)
assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == (
"ltx2",
"pipeline_ltx2_two_stage",
@@ -1488,6 +1835,9 @@ class TestRegistryIntegration:
"LTX23ImageToVideoPipeline",
"LTX2TwoStagesPipeline",
"LTX2ImageToVideoTwoStagesPipeline",
"LTX2ImageToVideoTwoStagePipeline",
"LTX23TwoStagePipeline",
"LTX23ImageToVideoTwoStagePipeline",
}
assert removed_entries.isdisjoint(_DIFFUSION_MODELS)
@@ -1497,6 +1847,7 @@ class TestRegistryIntegration:
expected = [
"LTX2Pipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
@@ -1509,6 +1860,7 @@ class TestRegistryIntegration:
"model_class_name",
[
"LTX2Pipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
@@ -1,13 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import pytest
import torch
from torch import nn
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel, _make_rms_norm
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import (
LTX2AudioVideoAttnProcessor,
LTX2VideoTransformer3DModel,
_LTX2ParallelAttention,
_make_rms_norm,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -66,3 +73,83 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
assert video_timestep.expected_dims == audio_timestep.expected_dims == 2
assert not video_timestep.split_output
assert not audio_timestep.split_output
def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):
monkeypatch.setattr(LTX2AudioVideoAttnProcessor, "_is_sp_enabled", staticmethod(lambda: True))
processor = LTX2AudioVideoAttnProcessor()
hidden_states = torch.zeros(1, 3, 4)
encoder_hidden_states = torch.zeros(1, 4, 4)
key_padding_mask = torch.tensor([[True, True, True, False]])
actual = processor._prepare_attention_mask(
SimpleNamespace(),
hidden_states,
encoder_hidden_states,
key_padding_mask,
batch_size=1,
sequence_length=4,
)
torch.testing.assert_close(actual, key_padding_mask[:, None, None, :])
def _fake_ltx_parallel_attention():
calls = []
class Backend:
@staticmethod
def get_name():
return "FLASH_ATTN"
@staticmethod
def supports_attention_mask():
return True
def run(name):
def forward(query, _key, _value, _metadata):
calls.append(name)
return torch.zeros_like(query)
return forward
layer = object.__new__(_LTX2ParallelAttention)
nn.Module.__init__(layer)
layer.attn_backend = Backend
layer.attention = SimpleNamespace(forward=run("native"))
layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa"))
layer.backend_pref = "FLASH_ATTN"
return layer, calls
@pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)])
def test_ltx_masked_cross_attention_uses_sdpa(query_length, key_length):
layer, calls = _fake_ltx_parallel_attention()
query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16)
key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16)
key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]])
metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
layer._run_local_attention(query, key, value, metadata)
assert calls == ["sdpa"]
def test_ltx_masked_self_attention_uses_sdpa():
layer, calls = _fake_ltx_parallel_attention()
query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
key_padding_mask = torch.tensor([[True, True, True, False]])
metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
layer._run_local_attention(query, key, value, metadata)
assert calls == ["sdpa"]
def test_ltx_unmasked_cross_attention_keeps_native_backend():
layer, calls = _fake_ltx_parallel_attention()
query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
layer._run_local_attention(query, key, value, None)
assert calls == ["native"]
@@ -189,7 +189,7 @@ def _generate_single_stage_video(
def _generate_bagel_image(
quantization_config: str | None = None,
diffusion_quantization_config: str | None = None,
num_inference_steps: int = 15,
) -> tuple[Any, float]:
"""Generate an image with BAGEL (multi-stage: LLM + Diffusion).
@@ -199,11 +199,11 @@ def _generate_bagel_image(
config_path = get_deploy_config_path("ci/bagel.yaml")
omni_kwargs: dict[str, Any] = {
"model": "ByteDance-Seed/BAGEL-7B-MoT",
"stage_configs_path": config_path,
"deploy_config": config_path,
"stage_init_timeout": 300,
}
if quantization_config:
omni_kwargs["quantization_config"] = quantization_config
if diffusion_quantization_config:
omni_kwargs["diffusion_quantization_config"] = diffusion_quantization_config
model_name = omni_kwargs.pop("model")
with OmniRunner(model_name, **omni_kwargs) as runner:
@@ -380,21 +380,21 @@ def test_single_stage_ltx2_fp8_uses_less_memory():
@hardware_test(res={"cuda": "H100"})
def test_bagel_fp8_generates_image():
"""BAGEL with FP8 quantization_config generates a valid image.
"""BAGEL with diffusion-stage FP8 generates a valid image.
FP8 should only apply to the diffusion stage (Stage-1), not the
LLM stage (Stage-0). We verify this by checking:
1. Image is generated successfully
2. LLM stage finish_reason is 'stop' (not 'length' from garbled output)
"""
image, _ = _generate_bagel_image(quantization_config="fp8")
image, _ = _generate_bagel_image(diffusion_quantization_config="fp8")
image.save("test_bagel_fp8.png")
@hardware_test(res={"cuda": "H100"})
def test_bagel_bf16_generates_image():
"""BAGEL without quantization generates a valid image (baseline)."""
image, _ = _generate_bagel_image(quantization_config=None)
image, _ = _generate_bagel_image(diffusion_quantization_config=None)
image.save("test_bagel_bf16.png")
@@ -95,6 +95,7 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata(
"resolution": 512,
"width": None,
"height": None,
"num_inference_steps": 1,
}
@@ -314,7 +314,7 @@ def test_bagel_text2img_mooncake_connector(run_level):
temp_config_file = _resolve_deploy_config(temp_config_file, run_level)
with OmniRunner(
"ByteDance-Seed/BAGEL-7B-MoT",
stage_configs_path=temp_config_file,
deploy_config=temp_config_file,
stage_init_timeout=300,
) as runner:
generated_image = _generate_bagel_image(runner.omni)
@@ -251,7 +251,7 @@ def test_bagel_img2img_shared_memory_connector(run_level):
config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level)
with OmniRunner(
"ByteDance-Seed/BAGEL-7B-MoT",
stage_configs_path=config_path,
deploy_config=config_path,
) as runner:
generated_image = _generate_bagel_img2img(runner.omni, input_image)
if run_level == "advanced_model":
@@ -266,7 +266,7 @@ def test_bagel_text2img_shared_memory_connector(run_level):
config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level)
with OmniRunner(
"ByteDance-Seed/BAGEL-7B-MoT",
stage_configs_path=config_path,
deploy_config=config_path,
) as runner:
generated_image = _generate_bagel_text2img(runner.omni)
if run_level == "advanced_model":
+117 -50
View File
@@ -20,6 +20,11 @@ from PIL import Image
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", choices=("official", "omni"), required=True)
parser.add_argument(
"--pipeline-kind",
choices=("one_stage", "distilled", "two_stage"),
default="one_stage",
)
parser.add_argument("--request", type=Path, required=True)
parser.add_argument("--output-dir", type=Path, required=True)
parser.add_argument("--model")
@@ -27,6 +32,8 @@ def _parse_args() -> argparse.Namespace:
parser.add_argument("--official-root", type=Path)
parser.add_argument("--checkpoint", type=Path)
parser.add_argument("--gemma-root", type=Path)
parser.add_argument("--spatial-upsampler", type=Path)
parser.add_argument("--distilled-lora", type=Path)
parser.add_argument("--enable-layerwise-offload", action="store_true")
return parser.parse_args()
@@ -91,9 +98,13 @@ def _configure_official_sdpa(pipeline: Any) -> None:
attention=attention,
masked_attention=attention,
)
owners_and_attributes = (
(pipeline.stage, "_transformer_builder"),
owners_and_attributes = [
(pipeline.prompt_encoder, "_embeddings_processor_builder"),
]
owners_and_attributes.extend(
(stage, "_transformer_builder")
for name in ("stage", "stage_1", "stage_2")
if (stage := getattr(pipeline, name, None)) is not None
)
for owner, attribute in owners_and_attributes:
builder = getattr(owner, attribute)
@@ -104,6 +115,12 @@ def _configure_official_sdpa(pipeline: Any) -> None:
)
def _require_path(path: Path | None, description: str) -> str:
if path is None or not path.is_file():
raise ValueError(f"Official {description} is required and must exist: {path}")
return str(path)
@torch.inference_mode()
def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
if args.official_root is None or args.checkpoint is None or args.gemma_root is None:
@@ -113,16 +130,54 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
torch.backends.cuda.enable_cudnn_sdp(False)
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_pipelines.distilled import DistilledPipeline
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
from ltx_pipelines.utils.args import ImageConditioningInput
from ltx_pipelines.utils.types import OffloadMode
pipeline = TI2VidOneStagePipeline(
checkpoint_path=str(args.checkpoint),
gemma_root=str(args.gemma_root),
loras=(),
offload_mode=OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE,
)
if request.get("pipeline_kind", "one_stage") != args.pipeline_kind:
raise ValueError(
f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}"
)
checkpoint = _require_path(args.checkpoint, "checkpoint")
gemma_root = str(args.gemma_root)
if not Path(gemma_root).is_dir():
raise ValueError(f"Official Gemma root is required and must be a directory: {gemma_root}")
offload_mode = OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE
pipeline: Any
if args.pipeline_kind == "one_stage":
pipeline = TI2VidOneStagePipeline(
checkpoint_path=checkpoint,
gemma_root=gemma_root,
loras=(),
offload_mode=offload_mode,
)
elif args.pipeline_kind == "distilled":
pipeline = DistilledPipeline(
distilled_checkpoint_path=checkpoint,
spatial_upsampler_path=_require_path(args.spatial_upsampler, "spatial upsampler"),
gemma_root=gemma_root,
loras=(),
offload_mode=offload_mode,
)
else:
distilled_lora = [
LoraPathStrengthAndSDOps(
path=_require_path(args.distilled_lora, "distilled LoRA"),
strength=1.0,
sd_ops=LTXV_LORA_COMFY_RENAMING_MAP,
)
]
pipeline = TI2VidTwoStagesPipeline(
checkpoint_path=checkpoint,
distilled_lora=distilled_lora,
spatial_upsampler_path=_require_path(args.spatial_upsampler, "spatial upsampler"),
gemma_root=gemma_root,
loras=(),
offload_mode=offload_mode,
)
_configure_official_sdpa(pipeline)
image_path = request.get("image")
images = (
@@ -140,34 +195,45 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
)
]
)
video, audio = pipeline(
prompt=request["prompt"],
negative_prompt=request["negative_prompt"],
seed=request["seed"],
height=request["height"],
width=request["width"],
num_frames=request["num_frames"],
frame_rate=request["fps"],
num_inference_steps=request["num_inference_steps"],
video_guider_params=MultiModalGuiderParams(
cfg_scale=request["video_cfg_scale"],
stg_scale=request["video_stg_scale"],
rescale_scale=request["video_rescale_scale"],
modality_scale=request["video_modality_scale"],
skip_step=0,
stg_blocks=request["video_stg_blocks"],
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=request["audio_cfg_scale"],
stg_scale=request["audio_stg_scale"],
rescale_scale=request["audio_rescale_scale"],
modality_scale=request["audio_modality_scale"],
skip_step=0,
stg_blocks=request["audio_stg_blocks"],
),
images=images,
max_batch_size=4,
)
if args.pipeline_kind == "distilled":
video, audio = pipeline(
prompt=request["prompt"],
seed=request["seed"],
height=request["height"],
width=request["width"],
num_frames=request["num_frames"],
frame_rate=request["fps"],
images=images,
)
else:
video, audio = pipeline(
prompt=request["prompt"],
negative_prompt=request["negative_prompt"],
seed=request["seed"],
height=request["height"],
width=request["width"],
num_frames=request["num_frames"],
frame_rate=request["fps"],
num_inference_steps=request["num_inference_steps"],
video_guider_params=MultiModalGuiderParams(
cfg_scale=request["video_cfg_scale"],
stg_scale=request["video_stg_scale"],
rescale_scale=request["video_rescale_scale"],
modality_scale=request["video_modality_scale"],
skip_step=0,
stg_blocks=request["video_stg_blocks"],
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=request["audio_cfg_scale"],
stg_scale=request["audio_stg_scale"],
rescale_scale=request["audio_rescale_scale"],
modality_scale=request["audio_modality_scale"],
skip_step=0,
stg_blocks=request["audio_stg_blocks"],
),
images=images,
max_batch_size=4,
)
video_tensor = torch.cat([chunk.detach().cpu() for chunk in video], dim=0)
_save_outputs(
args.output_dir,
@@ -179,6 +245,7 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
"attention_backend": "torch_sdpa",
"official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"),
"checkpoint": str(args.checkpoint),
"pipeline_kind": args.pipeline_kind,
},
)
@@ -193,15 +260,9 @@ def _unwrap_omni_output(output: Any) -> tuple[Any, Any, int]:
multimodal_output = frames.multimodal_output or {}
audio = multimodal_output.get("audio")
audio_sample_rate = multimodal_output.get("audio_sample_rate")
if frames.is_pipeline_output and isinstance(frames.request_output, OmniRequestOutput):
frames = frames.request_output
multimodal_output = frames.multimodal_output or {}
audio = multimodal_output.get("audio", audio)
audio_sample_rate = multimodal_output.get("audio_sample_rate", audio_sample_rate)
if isinstance(frames, OmniRequestOutput):
if not frames.images:
raise ValueError("No video frames found in OmniRequestOutput")
frames = frames.images
if not frames.images:
raise ValueError("No video frames found in OmniRequestOutput")
frames = frames.images
if isinstance(frames, list) and len(frames) == 1:
frames = frames[0]
@@ -253,6 +314,10 @@ def _canonical_video(video: Any) -> torch.Tensor:
def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
if args.model is None or args.model_class_name is None:
raise ValueError("Omni backend requires --model and --model-class-name")
if request.get("pipeline_kind", "one_stage") != args.pipeline_kind:
raise ValueError(
f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}"
)
attention_config = {"default": {"backend": "TORCH_SDPA"}}
@@ -273,7 +338,8 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
parallel_config=DiffusionParallelConfig(),
)
try:
model_class_name = get_model_class_name(omni)
detected_model_class_name = get_model_class_name(omni)
model_class_name = args.model_class_name
sampling_params = OmniDiffusionSamplingParams(
height=request["height"],
width=request["width"],
@@ -289,10 +355,9 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
key: value for key, value in request.items() if key.startswith("video_") or key.startswith("audio_")
}
apply_declared_extra_args(sampling_params, get_extra_body_params(model_class_name), guidance)
prompt: dict[str, Any] = {
"prompt": request["prompt"],
"negative_prompt": request["negative_prompt"],
}
prompt: dict[str, Any] = {"prompt": request["prompt"]}
if request.get("negative_prompt"):
prompt["negative_prompt"] = request["negative_prompt"]
image_path = request.get("image")
if image_path is not None:
with Image.open(str(image_path)) as source_image:
@@ -312,6 +377,8 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
"attention_backend": "torch_sdpa",
"model": args.model,
"model_class_name": model_class_name,
"model_config_class_name": detected_model_class_name,
"pipeline_kind": args.pipeline_kind,
},
)
finally:
@@ -3,10 +3,9 @@
"""E2E accuracy guard against a pinned Lightricks LTX pipeline revision.
The comparison runs both runtimes through PyTorch SDPA and uses
``max_batch_size=4`` in the official reference to match Omni's fused guidance
batch. Video and audio guidance use the official non-HQ one-stage defaults;
only the generation shape and step count are reduced for CI runtime.
The original reduced one-stage guards remain unchanged. Four complementary
default-shape cases cover both model versions, both two-stage weight families,
and T2V/I2V without expanding every combination into a duplicate golden.
"""
from __future__ import annotations
@@ -17,8 +16,9 @@ import shutil
import subprocess
import sys
import time
from dataclasses import dataclass
from dataclasses import asdict, dataclass, replace
from pathlib import Path
from typing import Literal
import numpy as np
import pytest
@@ -53,74 +53,260 @@ NEGATIVE_PROMPT = (
# Both runtimes use PyTorch SDPA with the current Torch dispatch defaults.
ATTENTION_BACKEND = "torch_sdpa"
VIDEO_SSIM_MEAN_THRESHOLD = 0.95
VIDEO_SSIM_MIN_THRESHOLD = 0.90
VIDEO_PSNR_MEAN_THRESHOLD = 30.0
AUDIO_RELATIVE_L2_THRESHOLD = 0.2
AUDIO_COSINE_THRESHOLD = 0.95
@dataclass(frozen=True)
class LTXArtifact:
repo_id: str
filename: str
revision: str
env: str
repo_type: Literal["model", "dataset"] = "model"
@dataclass(frozen=True)
class LTXAccuracyThresholds:
video_ssim_mean: float
video_ssim_min: float
video_psnr_mean_db: float
# Zero disables gating for waveform-sensitive audio that is not expected
# to align with the official output.
audio_relative_l2: float
audio_cosine_similarity: float
STRICT_THRESHOLDS = LTXAccuracyThresholds(
video_ssim_mean=0.95,
video_ssim_min=0.90,
video_psnr_mean_db=30.0,
audio_relative_l2=0.20,
audio_cosine_similarity=0.95,
)
@dataclass(frozen=True)
class LTXAccuracyCase:
name: str
pipeline_kind: Literal["one_stage", "distilled", "two_stage"]
model_id: str
model_revision: str
model_env: str
model_class_name: str
checkpoint_repo: str
checkpoint_filename: str
checkpoint_revision: str
checkpoint_env: str
stg_block: int
checkpoint: LTXArtifact
width: int
height: int
num_frames: int
num_inference_steps: int
seed: int
stg_block: int | None
thresholds: LTXAccuracyThresholds
prompt: str = PROMPT
image_repo: str | None = None
image_filename: str | None = None
image_revision: str | None = None
gemma_model_id: str | None = None
gemma_model_revision: str | None = None
gemma_model_env: str | None = None
image: LTXArtifact | None = None
spatial_upsampler: LTXArtifact | None = None
distilled_lora: LTXArtifact | None = None
# Official uses block streaming; Omni uses layerwise offload. Both preserve
# bf16 arithmetic while keeping full-resolution two-stage cases on one H100.
enable_layerwise_offload: bool = False
CASES = (
LTX2_REVISION = "47da56e2ad66ce4125a9922b4a8826bf407f9d0a"
LTX23_REVISION = "4229404625088d21c4f112eb640fb04a0900ee25"
LTX2_CHECKPOINT = LTXArtifact(
repo_id="Lightricks/LTX-2",
filename="ltx-2-19b-dev.safetensors",
revision=LTX2_REVISION,
env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT",
)
LTX2_DISTILLED_CHECKPOINT = replace(
LTX2_CHECKPOINT,
filename="ltx-2-19b-distilled.safetensors",
env="VLLM_TEST_LTX2_DISTILLED_OFFICIAL_CHECKPOINT",
)
LTX2_UPSAMPLER = replace(
LTX2_CHECKPOINT,
filename="ltx-2-spatial-upscaler-x2-1.0.safetensors",
env="VLLM_TEST_LTX2_UPSAMPLER",
)
LTX2_DISTILLED_LORA = replace(
LTX2_CHECKPOINT,
filename="ltx-2-19b-distilled-lora-384.safetensors",
env="VLLM_TEST_LTX2_DISTILLED_LORA",
)
LTX23_CHECKPOINT = LTXArtifact(
repo_id="Lightricks/LTX-2.3",
filename="ltx-2.3-22b-dev.safetensors",
revision=LTX23_REVISION,
env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
)
LTX23_DISTILLED_CHECKPOINT = replace(
LTX23_CHECKPOINT,
# The pinned Diffusers checkpoint contains the original merged distilled
# Transformer. Keep the reference Transformer on the same version; the
# spatial upsampler is versioned independently and remains on 1.1 below.
filename="ltx-2.3-22b-distilled.safetensors",
env="VLLM_TEST_LTX23_DISTILLED_OFFICIAL_CHECKPOINT",
)
LTX23_UPSAMPLER = replace(
LTX23_CHECKPOINT,
filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
env="VLLM_TEST_LTX23_UPSAMPLER",
)
LTX23_DISTILLED_LORA = replace(
LTX23_CHECKPOINT,
filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors",
env="VLLM_TEST_LTX23_DISTILLED_LORA",
)
I2V_IMAGE = LTXArtifact(
repo_id="huggingface/documentation-images",
filename="diffusers/svd/rocket.png",
revision="645d8364f0c7a101180b364811b5a11a362e4010",
env="VLLM_TEST_LTX_I2V_IMAGE",
repo_type="dataset",
)
LEGACY_CASES = (
LTXAccuracyCase(
name="ltx2",
pipeline_kind="one_stage",
model_id="Lightricks/LTX-2",
model_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a",
model_revision=LTX2_REVISION,
model_env="VLLM_TEST_LTX2_MODEL",
model_class_name="LTX2Pipeline",
checkpoint_repo="Lightricks/LTX-2",
checkpoint_filename="ltx-2-19b-dev.safetensors",
checkpoint_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a",
checkpoint_env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT",
checkpoint=LTX2_CHECKPOINT,
width=512,
height=384,
num_frames=25,
num_inference_steps=20,
seed=42,
stg_block=29,
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx2_3",
pipeline_kind="one_stage",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2Pipeline",
checkpoint_repo="Lightricks/LTX-2.3",
checkpoint_filename="ltx-2.3-22b-dev.safetensors",
checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25",
checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
checkpoint=LTX23_CHECKPOINT,
width=512,
height=384,
num_frames=25,
num_inference_steps=20,
seed=42,
stg_block=28,
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx2_3_i2v",
pipeline_kind="one_stage",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2Pipeline",
checkpoint_repo="Lightricks/LTX-2.3",
checkpoint_filename="ltx-2.3-22b-dev.safetensors",
checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25",
checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
checkpoint=LTX23_CHECKPOINT,
width=512,
height=384,
num_frames=25,
num_inference_steps=20,
seed=42,
stg_block=28,
image_repo="huggingface/documentation-images",
image_filename="diffusers/svd/rocket.png",
image_revision="645d8364f0c7a101180b364811b5a11a362e4010",
thresholds=STRICT_THRESHOLDS,
image=I2V_IMAGE,
),
)
TWO_STAGE_CASES = (
LTXAccuracyCase(
name="ltx2_distilled_t2v",
pipeline_kind="distilled",
model_id="rootonchair/LTX-2-19b-distilled",
model_revision="388e2846f54aae51687498ffb6b27c7c2c9ce9e5",
model_env="VLLM_TEST_LTX2_DISTILLED_MODEL",
model_class_name="LTX2DistilledPipeline",
checkpoint=LTX2_DISTILLED_CHECKPOINT,
spatial_upsampler=LTX2_UPSAMPLER,
width=1536,
height=1024,
num_frames=121,
num_inference_steps=8,
seed=10,
stg_block=None,
gemma_model_id="Lightricks/LTX-2",
gemma_model_revision=LTX2_REVISION,
gemma_model_env="VLLM_TEST_LTX2_MODEL",
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx23_distilled_i2v",
pipeline_kind="distilled",
model_id="diffusers/LTX-2.3-Distilled-Diffusers",
model_revision="432e0d3c2d1769aaa4d295f9243f7062bf6b47ee",
model_env="VLLM_TEST_LTX23_DISTILLED_MODEL",
model_class_name="LTX2DistilledPipeline",
checkpoint=LTX23_DISTILLED_CHECKPOINT,
spatial_upsampler=LTX23_UPSAMPLER,
width=1536,
height=1024,
num_frames=121,
num_inference_steps=8,
seed=10,
stg_block=None,
gemma_model_id="diffusers/LTX-2.3-Diffusers",
gemma_model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
gemma_model_env="VLLM_TEST_LTX23_MODEL",
thresholds=STRICT_THRESHOLDS,
image=I2V_IMAGE,
),
LTXAccuracyCase(
name="ltx2_two_stage_layer_fused_t2v",
pipeline_kind="two_stage",
model_id="Lightricks/LTX-2",
model_revision=LTX2_REVISION,
model_env="VLLM_TEST_LTX2_MODEL",
model_class_name="LTX2TwoStagePipeline",
checkpoint=LTX2_CHECKPOINT,
spatial_upsampler=LTX2_UPSAMPLER,
distilled_lora=LTX2_DISTILLED_LORA,
width=1536,
height=1024,
num_frames=121,
num_inference_steps=40,
seed=10,
stg_block=29,
enable_layerwise_offload=True,
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx23_two_stage_layer_fused_i2v",
pipeline_kind="two_stage",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2TwoStagePipeline",
checkpoint=LTX23_CHECKPOINT,
spatial_upsampler=LTX23_UPSAMPLER,
distilled_lora=LTX23_DISTILLED_LORA,
width=1536,
height=1024,
num_frames=121,
num_inference_steps=30,
seed=10,
stg_block=28,
enable_layerwise_offload=True,
thresholds=STRICT_THRESHOLDS,
image=I2V_IMAGE,
),
)
CASES = (*LEGACY_CASES, *TWO_STAGE_CASES)
def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None:
start = time.perf_counter()
subprocess.run(command, env=env, timeout=timeout, check=True)
@@ -209,6 +395,7 @@ def _resolve_model(case: LTXAccuracyCase) -> Path:
"model_index.json",
"audio_vae/*",
"connectors/*",
"latent_upsampler/*",
"processor/*",
"scheduler/*",
"text_encoder/config.json",
@@ -223,74 +410,82 @@ def _resolve_model(case: LTXAccuracyCase) -> Path:
)
def _resolve_gemma_root(model: Path) -> Path:
def _resolve_gemma_root(case: LTXAccuracyCase, model: Path) -> Path:
configured_root = os.environ.get("VLLM_TEST_LTX_GEMMA_ROOT")
if configured_root:
root = Path(configured_root)
assert root.is_dir(), f"Gemma root not found: {root}"
return root
if case.gemma_model_id is not None:
configured_model = os.environ.get(case.gemma_model_env or "")
if configured_model and Path(configured_model).is_dir():
return Path(configured_model)
return Path(
snapshot_download(
repo_id=case.gemma_model_id,
revision=case.gemma_model_revision,
allow_patterns=[
"processor/*",
"text_encoder/*",
"tokenizer/*",
],
)
)
return model
def _resolve_checkpoint(case: LTXAccuracyCase, model: Path) -> Path:
configured_checkpoint = os.environ.get(case.checkpoint_env)
if configured_checkpoint:
checkpoint = Path(configured_checkpoint)
assert checkpoint.is_file(), f"Official checkpoint not found: {checkpoint}"
return checkpoint
model_checkpoint = model / case.checkpoint_filename
if model_checkpoint.is_file():
return model_checkpoint
def _resolve_artifact(artifact: LTXArtifact, model: Path | None = None) -> Path:
configured_path = os.environ.get(artifact.env)
if configured_path:
path = Path(configured_path)
assert path.is_file(), f"Configured LTX artifact not found: {path}"
return path
if model is not None:
model_path = model / artifact.filename
if model_path.is_file():
return model_path
return Path(
hf_hub_download(
repo_id=case.checkpoint_repo,
filename=case.checkpoint_filename,
revision=case.checkpoint_revision,
repo_id=artifact.repo_id,
repo_type=None if artifact.repo_type == "model" else artifact.repo_type,
filename=artifact.filename,
revision=artifact.revision,
)
)
def _resolve_image(case: LTXAccuracyCase) -> Path | None:
if case.image_filename is None:
if case.image is None:
return None
if case.image_repo is None or case.image_revision is None:
raise ValueError(f"Incomplete image source for LTX accuracy case {case.name!r}.")
configured_image = os.environ.get("VLLM_TEST_LTX_I2V_IMAGE")
if configured_image:
image = Path(configured_image)
assert image.is_file(), f"LTX I2V conditioning image not found: {image}"
return image
return Path(
hf_hub_download(
repo_id=case.image_repo,
repo_type="dataset",
filename=case.image_filename,
revision=case.image_revision,
)
)
return _resolve_artifact(case.image)
def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
request: dict[str, object] = {
"pipeline_kind": case.pipeline_kind,
"prompt": case.prompt,
"negative_prompt": NEGATIVE_PROMPT,
"width": 512,
"height": 384,
"num_frames": 25,
"negative_prompt": "" if case.pipeline_kind == "distilled" else NEGATIVE_PROMPT,
"width": case.width,
"height": case.height,
"num_frames": case.num_frames,
"fps": 24,
"num_inference_steps": 20,
"seed": 42,
"video_cfg_scale": 3.0,
"audio_cfg_scale": 7.0,
"video_stg_scale": 1.0,
"audio_stg_scale": 1.0,
"video_modality_scale": 3.0,
"audio_modality_scale": 3.0,
"video_rescale_scale": 0.7,
"audio_rescale_scale": 0.7,
"video_stg_blocks": [case.stg_block],
"audio_stg_blocks": [case.stg_block],
"num_inference_steps": case.num_inference_steps,
"seed": case.seed,
}
if case.pipeline_kind != "distilled":
assert case.stg_block is not None
request.update(
video_cfg_scale=3.0,
audio_cfg_scale=7.0,
video_stg_scale=1.0,
audio_stg_scale=1.0,
video_modality_scale=3.0,
audio_modality_scale=3.0,
video_rescale_scale=0.7,
audio_rescale_scale=0.7,
video_stg_blocks=[case.stg_block],
audio_stg_blocks=[case.stg_block],
)
if image is not None:
request["image"] = str(image.resolve())
return request
@@ -299,20 +494,30 @@ def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
assert reference.shape == prediction.shape
assert reference.ndim == 4 and reference.shape[-1] == 3
device = torch.device("cpu")
ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to(device)
psnr = PeakSignalNoiseRatio(data_range=1.0).to(device)
ssim_scores: list[float] = []
psnr_scores: list[float] = []
for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0)
prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0)
ssim_scores.append(float(StructuralSimilarityIndexMeasure(data_range=1.0)(prediction_tensor, reference_tensor)))
psnr_scores.append(float(PeakSignalNoiseRatio(data_range=1.0)(prediction_tensor, reference_tensor)))
difference = np.abs(reference.astype(np.float64) - prediction.astype(np.float64))
max_abs = 0.0
absolute_error_sum = 0.0
with torch.inference_mode():
for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0).to(device)
prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0).to(device)
ssim_scores.append(float(ssim(prediction_tensor, reference_tensor)))
psnr_scores.append(float(psnr(prediction_tensor, reference_tensor)))
ssim.reset()
psnr.reset()
difference = np.abs(reference_frame.astype(np.float64) - prediction_frame.astype(np.float64))
max_abs = max(max_abs, float(difference.max()))
absolute_error_sum += float(difference.sum())
return {
"ssim_mean": float(np.mean(ssim_scores)),
"ssim_min": float(np.min(ssim_scores)),
"psnr_mean_db": float(np.mean(psnr_scores)),
"max_abs": float(difference.max()),
"mean_abs": float(difference.mean()),
"max_abs": max_abs,
"mean_abs": absolute_error_sum / reference.size,
}
@@ -343,13 +548,19 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
@pytest.mark.benchmark
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None:
"""Compare official and Omni raw AV outputs from the same E2E request."""
def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None:
"""Compare official and Omni raw AV outputs from the same single-H100 request."""
configured_artifact_root = os.environ.get("VLLM_TEST_LTX_ARTIFACT_ROOT")
if configured_artifact_root:
accuracy_artifact_root = Path(configured_artifact_root)
accuracy_artifact_root.mkdir(parents=True, exist_ok=True)
output_root = reset_artifact_dir(accuracy_artifact_root / "ltx_official" / case.name)
official_root, official_revision = _official_source(accuracy_artifact_root / "ltx_official")
model = _resolve_model(case)
gemma_root = _resolve_gemma_root(model)
checkpoint = _resolve_checkpoint(case, model)
gemma_root = _resolve_gemma_root(case, model)
checkpoint = _resolve_artifact(case.checkpoint, model)
spatial_upsampler = _resolve_artifact(case.spatial_upsampler, model) if case.spatial_upsampler is not None else None
distilled_lora = _resolve_artifact(case.distilled_lora, model) if case.distilled_lora is not None else None
image = _resolve_image(case)
request_path = output_root / "request.json"
request_path.write_text(json.dumps(_request(case, image), indent=2) + "\n")
@@ -359,8 +570,13 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
str(runner),
"--request",
str(request_path),
"--pipeline-kind",
case.pipeline_kind,
]
if os.environ.get("VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", "").lower() in {"1", "true", "yes", "on"}:
enable_layerwise_offload = case.enable_layerwise_offload or os.environ.get(
"VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", ""
).lower() in {"1", "true", "yes", "on"}
if enable_layerwise_offload:
runner_args.append("--enable-layerwise-offload")
env = os.environ.copy()
env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision
@@ -372,10 +588,9 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
)
official_output = output_root / "official"
_run(
_official_runner_prefix()
+ runner_args
+ [
official_args = _official_runner_prefix() + runner_args
official_args.extend(
[
"--backend",
"official",
"--output-dir",
@@ -386,12 +601,16 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
str(checkpoint),
"--gemma-root",
str(gemma_root),
],
env=env,
]
)
if spatial_upsampler is not None:
official_args.extend(["--spatial-upsampler", str(spatial_upsampler)])
if distilled_lora is not None:
official_args.extend(["--distilled-lora", str(distilled_lora)])
_run(official_args, env=env)
omni_output = output_root / "omni"
_run(
omni_args = (
[sys.executable]
+ runner_args
+ [
@@ -403,9 +622,9 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
str(model),
"--model-class-name",
case.model_class_name,
],
env=env,
]
)
_run(omni_args, env=env)
official_metadata = json.loads((official_output / "metadata.json").read_text())
omni_metadata = json.loads((omni_output / "metadata.json").read_text())
@@ -423,18 +642,27 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
result = {
"case": case.name,
"task": "i2v" if image is not None else "t2v",
"pipeline_kind": case.pipeline_kind,
"model_class_name": case.model_class_name,
"attention_backend": ATTENTION_BACKEND,
"official_revision": official_revision,
"model_revision": case.model_revision,
"checkpoint_revision": case.checkpoint_revision,
"checkpoint_revision": case.checkpoint.revision,
"spatial_upsampler_revision": (case.spatial_upsampler.revision if case.spatial_upsampler is not None else None),
"distilled_lora_revision": case.distilled_lora.revision if case.distilled_lora is not None else None,
"enable_layerwise_offload": enable_layerwise_offload,
"thresholds": asdict(case.thresholds),
"request": _request(case, image),
"video": video_metrics,
"audio": audio_metrics,
}
(output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result, indent=2))
assert video_metrics["ssim_mean"] >= VIDEO_SSIM_MEAN_THRESHOLD
assert video_metrics["ssim_min"] >= VIDEO_SSIM_MIN_THRESHOLD
assert video_metrics["psnr_mean_db"] >= VIDEO_PSNR_MEAN_THRESHOLD
assert audio_metrics["relative_l2"] <= AUDIO_RELATIVE_L2_THRESHOLD
assert audio_metrics["cosine_similarity"] >= AUDIO_COSINE_THRESHOLD
assert video_metrics["ssim_mean"] >= case.thresholds.video_ssim_mean
assert video_metrics["ssim_min"] >= case.thresholds.video_ssim_min
assert video_metrics["psnr_mean_db"] >= case.thresholds.video_psnr_mean_db
if case.thresholds.audio_relative_l2 > 0:
assert audio_metrics["relative_l2"] <= case.thresholds.audio_relative_l2
if case.thresholds.audio_cosine_similarity > 0:
assert audio_metrics["cosine_similarity"] >= case.thresholds.audio_cosine_similarity
@@ -17,7 +17,7 @@ from vllm_omni.platforms import current_omni_platform
# Override in CI if needed: AUDIOX_TEST_MODEL=<model-or-local-path>
models = [os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random")]
# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
_OMNI_RUNNER_PARAMS = [(m, None, {"model_class_name": "AudioXPipeline"}) for m in models]
pytestmark = [
@@ -43,7 +43,7 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
BAGEL_STAGE_CONFIG = get_deploy_config_path("ci/bagel.yaml")
DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
# (model, stage_configs_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
# (model, deploy_config_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
_OMNI_RUNNER_PARAM = (MODEL, BAGEL_STAGE_CONFIG)
@@ -126,7 +126,7 @@ def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[s
]
# (model, stage_configs_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
# (model, deploy_config_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
_cosy_deployment = get_deploy_config_path("cosyvoice3.yaml")
_cosy_model_path = str(_resolve_model_dir())
_OMNI_RUNNER_PARAMS = [
@@ -17,7 +17,7 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
MODEL = "baidu/ERNIE-Image"
# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
_OMNI_RUNNER_PARAM = (
MODEL,
None,
@@ -96,7 +96,7 @@ def test_offline_voice_clone_zh(async_chunk: bool) -> None:
with OmniRunner(
MODEL,
stage_configs_path=_get_deploy_config(async_chunk=async_chunk),
deploy_config=_get_deploy_config(async_chunk=async_chunk),
stage_init_timeout=600,
trust_remote_code=True,
) as omni_runner:
@@ -52,7 +52,7 @@ pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
def omni() -> Generator[Omni, None, None]:
with OmniRunner(
MODEL_NAME,
stage_configs_path=str(STAGE_CONFIG_PATH),
deploy_config=str(STAGE_CONFIG_PATH),
trust_remote_code=True,
) as runner:
yield runner.omni
@@ -84,7 +84,6 @@ def longcat_omni_runner() -> Generator[OmniRunner, None, None]:
with OmniRunner(
model,
seed=42,
stage_configs_path=None,
model_class_name="LongCatVideoAvatarPipeline",
additional_config={
"model_type": "avatar-v1.5",
@@ -38,7 +38,7 @@ _AR_PATCH_SIZE = 16
MODEL_PATH = "bytedance-research/MammothModa2-Preview"
T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG})
_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG)
# Golden pixel reference file. Set UPDATE_GOLDEN=1 to regenerate.
_GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"
@@ -17,7 +17,7 @@ models = ["openbmb/MiniCPM-o-4_5"]
_CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
test_params = [(model, None, {"deploy_config": _CI_DEPLOY, "trust_remote_code": True}) for model in models]
test_params = [(model, _CI_DEPLOY, {"trust_remote_code": True}) for model in models]
def get_question(prompt_type: str = "text") -> str:
@@ -25,7 +25,7 @@ from vllm_omni import Omni
MODEL_NAME = "OpenMOSS-Team/MOSS-TTS-Nano"
STAGE_CONFIG = get_deploy_config_path("moss_tts_nano.yaml")
# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
_OMNI_RUNNER_PARAM = (
MODEL_NAME,
STAGE_CONFIG,
@@ -20,13 +20,11 @@ from tests.helpers.stage_config import get_deploy_config_path
MODEL = "k2-fsa/OmniVoice"
DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml")
# OmniRunner tuple: model, legacy stage config path, extra Omni kwargs.
# The migrated test passes deploy_config through extra Omni kwargs.
# OmniRunner tuple: model, deploy config path, extra Omni kwargs.
_OMNI_RUNNER_PARAM = (
MODEL,
None,
DEPLOY_CONFIG,
{
"deploy_config": DEPLOY_CONFIG,
"trust_remote_code": True,
"log_stats": True,
},
@@ -131,7 +131,7 @@ def test_sensenova_u1_lora_scale_and_deactivation(tmp_path) -> None:
"""Validate LoRA effect, scale sensitivity, bounded perturbation, and clean deactivation."""
lora_request = _make_file_lora_request(tmp_path / "sensenova_lora")
with OmniRunner(MODEL, stage_configs_path=None) as runner:
with OmniRunner(MODEL) as runner:
omni = runner.omni
baseline = _generate(omni)
img_1x = _generate(omni, lora_request, lora_scale=1.0)
@@ -164,10 +164,7 @@ def _generate_sensenova_u1_img2img(
def test_sensenova_u1_img2img(run_level):
"""Test SenseNova-U1 img2img (single-stage diffusion, no deploy YAML)."""
input_image = _load_input_image()
with OmniRunner(
"SenseNova/SenseNova-U1-8B-MoT",
stage_configs_path=None,
) as runner:
with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner:
generated_image = _generate_sensenova_u1_img2img(runner.omni, input_image)
if run_level == "full_model" or run_level == "advanced_model":
_validate_pixels(generated_image)
@@ -182,7 +179,6 @@ def test_sensenova_u1_img2img_cache_dit():
input_image = _load_input_image()
with OmniRunner(
"SenseNova/SenseNova-U1-8B-MoT",
stage_configs_path=None,
cache_backend="cache_dit",
) as runner:
_generate_sensenova_u1_img2img(runner.omni, input_image)
@@ -151,10 +151,7 @@ def _generate_sensenova_u1_image(
@hardware_test(res={"cuda": "H100"})
def test_sensenova_u1_text2img(run_level):
"""Test SenseNova-U1 text2img (single-stage diffusion, no deploy YAML)."""
with OmniRunner(
"SenseNova/SenseNova-U1-8B-MoT",
stage_configs_path=None,
) as runner:
with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner:
generated_image = _generate_sensenova_u1_image(runner.omni)
if run_level == "full_model" or run_level == "advanced_model":
_validate_pixels(generated_image)
@@ -168,7 +165,6 @@ def test_sensenova_u1_text2img_cache_dit():
"""Test SenseNova-U1 text2img with Cache-DiT enabled."""
with OmniRunner(
"SenseNova/SenseNova-U1-8B-MoT",
stage_configs_path=None,
cache_backend="cache_dit",
) as runner:
_generate_sensenova_u1_image(runner.omni)
@@ -159,7 +159,7 @@ def test_soulxsinger_multistage_from_audio(
with OmniRunner(
model,
stage_configs_path=get_deploy_config_path(deploy_yaml),
deploy_config=get_deploy_config_path(deploy_yaml),
async_chunk=False,
) as runner:
sampling = OmniDiffusionSamplingParams(
@@ -212,7 +212,7 @@ def test_soulxsinger_svs_precomputed(soulx_weights: tuple[Path, Path, Path]) ->
with OmniRunner(
str(base_dir),
stage_configs_path=get_deploy_config_path("soulxsinger_svs.yaml"),
deploy_config=get_deploy_config_path("soulxsinger_svs.yaml"),
async_chunk=False,
) as runner:
sampling = OmniDiffusionSamplingParams(
@@ -19,7 +19,7 @@ from vllm_omni.outputs import OmniRequestOutput
MODEL = "stepfun-ai/Step-Audio-2-mini"
DEPLOY_CONFIG = get_deploy_config_path("step_audio2_ci.yaml")
TEST_PARAMS = [(MODEL, None, {"deploy_config": DEPLOY_CONFIG, "trust_remote_code": True})]
TEST_PARAMS = [(MODEL, DEPLOY_CONFIG, {"trust_remote_code": True})]
pytestmark = [pytest.mark.slow, pytest.mark.tts]
@@ -102,7 +102,7 @@ def test_voxtral_tts_offline_streaming():
async def _run():
async_omni = AsyncOmni(
model=MODEL,
stage_configs_path=STAGE_CONFIG,
deploy_config=STAGE_CONFIG,
stage_init_timeout=300,
enforce_eager=True,
)
+109
View File
@@ -0,0 +1,109 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Online-serving smoke tests for the public LTX-2 pipeline entries."""
import os
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.media import generate_synthetic_image
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
DISTILLED_MODEL = os.getenv("VLLM_TEST_LTX23_DISTILLED_MODEL", "diffusers/LTX-2.3-Distilled-Diffusers")
BASE_MODEL = os.getenv("VLLM_TEST_LTX23_MODEL", "diffusers/LTX-2.3-Diffusers")
PROMPT = "A serene lake at sunset with mountains in the background."
SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
HSDP_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
def _server(
model: str,
model_class_name: str,
*extra_args: str,
) -> OmniServerParams:
return OmniServerParams(
model=model,
server_args=[
"--enforce-eager",
*extra_args,
"--model-class-name",
model_class_name,
],
)
def _cases():
return [
pytest.param(
_server(BASE_MODEL, "LTX2Pipeline"),
3,
id="one_stage",
marks=SINGLE_CARD_MARKS,
),
pytest.param(
_server(DISTILLED_MODEL, "LTX2DistilledPipeline"),
8,
id="distilled",
marks=SINGLE_CARD_MARKS,
),
pytest.param(
_server(
DISTILLED_MODEL,
"LTX2DistilledPipeline",
"--use-hsdp",
"--hsdp-shard-size",
"2",
),
8,
id="distilled_hsdp2",
marks=HSDP_MARKS,
),
pytest.param(
_server(
BASE_MODEL,
"LTX2TwoStagePipeline",
"--enable-layerwise-offload",
),
3,
id="two_stage",
marks=SINGLE_CARD_MARKS,
),
]
@pytest.mark.parametrize(
("omni_server", "num_inference_steps"),
_cases(),
indirect=["omni_server"],
)
def test_ltx_pipeline_entries(
omni_server: OmniServer,
num_inference_steps: int,
openai_client: OpenAIClientHandler,
subtests,
):
for task in ("t2v", "i2v"):
with subtests.test(task=task):
form_data = {
"prompt": PROMPT,
"model": omni_server.model,
"height": 128,
"width": 128,
"num_frames": 9,
"fps": 8,
"num_inference_steps": num_inference_steps,
"seed": 42,
}
request_config = {
"model": omni_server.model,
"form_data": form_data,
}
if task == "i2v":
request_config["image_reference"] = (
f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
)
openai_client.send_video_diffusion_request(request_config)
@@ -1,91 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Expansion tests for LTX-2 two-stage pipelines.
Coverage:
- HSDP with LTX2DistilledPipeline (text-to-video and image-to-video)
"""
import os
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.media import generate_synthetic_image
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
MODEL = os.getenv("VLLM_TEST_LTX2_MODEL", "rootonchair/LTX-2-19b-distilled")
PROMPT = "A serene lake at sunset with mountains in the background."
PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
HSDP_ARGS = ["--use-hsdp", "--hsdp-shard-size", "2"]
def _cases():
cases = []
cases.append(
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
*HSDP_ARGS,
"--model-class-name",
"LTX2DistilledPipeline",
],
),
False,
id="t2v_hsdp",
marks=PARALLEL_MARKS,
)
)
cases.append(
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
*HSDP_ARGS,
"--model-class-name",
"LTX2DistilledPipeline",
],
),
True,
id="i2v_hsdp",
marks=PARALLEL_MARKS,
)
)
return cases
@pytest.mark.parametrize(("omni_server", "is_i2v"), _cases(), indirect=["omni_server"])
def test_ltx2_two_stage_hsdp(
omni_server: OmniServer,
is_i2v: bool,
openai_client: OpenAIClientHandler,
):
# Keep CI small while exercising the fixed 8+3-step distilled recipe.
# Height and width describe the final output; Stage 1 runs at half size.
form_data = {
"prompt": PROMPT,
"model": omni_server.model,
"height": 128,
"width": 128,
"num_frames": 9,
"fps": 8,
"num_inference_steps": 8,
"seed": 42,
}
request_config = {
"model": omni_server.model,
"form_data": form_data,
}
if is_i2v:
request_config["image_reference"] = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
openai_client.send_video_diffusion_request(request_config)
-139
View File
@@ -4,7 +4,6 @@ invariant to the specific attributes of vLLM config except in cases where we
explicitly patch values that differ from vLLM.
"""
import argparse
import inspect
import json
import os
@@ -20,7 +19,6 @@ from vllm.engine.arg_utils import EngineArgs
from vllm_omni.config.model import OmniModelConfig
from vllm_omni.engine.arg_utils import OmniEngineArgs
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.stage_init_utils import build_engine_args_dict
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -149,19 +147,6 @@ def test_qwen3_tts_startup_task_type_is_validated():
)
def test_from_cli_args_picks_up_stage_configs_path():
"""from_cli_args should pick up stage_configs_path from namespace."""
ns = argparse.Namespace(
model="facebook/opt-125m",
stage_configs_path="/some/path.yaml",
custom_pipeline_args=None,
)
args = OmniEngineArgs.from_cli_args(ns)
assert args.stage_configs_path == "/some/path.yaml"
assert args.custom_pipeline_args is None
def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch):
"""Ensure Code2Wav mirrors stage max_model_len into nested HF overrides.
@@ -418,130 +403,6 @@ def test_non_override_ar_stage_inputs_embeds_size_matches_hidden_size():
assert omni_config.get_inputs_embeds_size() == omni_config.hf_text_config.hidden_size
def test_stage_configs_path_field():
"""OmniEngineArgs with stage_configs_path should construct without error."""
args = OmniEngineArgs(stage_configs_path="/some/path.yaml")
assert args.stage_configs_path == "/some/path.yaml"
def test_strip_single_engine_args():
"""_strip_single_engine_args should remove EngineArgs fields but keep omni fields."""
kwargs = {
# Parent EngineArgs fields — stripped unless explicitly allowlisted
"compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}',
"tensor_parallel_size": 4,
"gpu_memory_utilization": 0.9,
"model": "some/model",
# Parent field that should be kept (allowlisted)
"worker_extension_cls": "some.Extension",
# OmniEngineArgs-only / non-engine fields — should pass through
"stage_configs_path": "/path/to/yaml",
"custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
"mode": "text-to-image",
"lora_path": "/some/lora",
}
filtered = AsyncOmniEngine._strip_single_engine_args(kwargs)
# Stripped — parent EngineArgs fields
assert "compilation_config" not in filtered
assert filtered["tensor_parallel_size"] == 4
assert "gpu_memory_utilization" not in filtered
assert "model" not in filtered
# Stripped — orchestrator-level OmniEngineArgs field
assert "stage_configs_path" not in filtered
# Kept
assert filtered["worker_extension_cls"] == "some.Extension"
assert filtered["custom_pipeline_args"] == {"pipeline_class": "my.Pipeline"}
assert filtered["mode"] == "text-to-image"
assert filtered["lora_path"] == "/some/lora"
def test_strip_single_engine_args_model_does_not_trigger_warning(mocker):
"""model is always in kwargs (callers set it via from_cli_args/asdict),
so it should not cause the override warning by itself or appear in it."""
mock_warn = mocker.patch("vllm_omni.engine.async_omni_engine.logger.warning")
# Typical caller kwargs: model is always present, no other parent
# EngineArgs fields are explicitly overridden.
AsyncOmniEngine._strip_single_engine_args(
{
"model": "some/model",
"custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
}
)
mock_warn.assert_not_called()
# When there *are* genuinely surprising overrides alongside model,
# the warning should mention them but not model. Keep-listed fields such as
# tensor_parallel_size are intentionally passed through and should not warn.
AsyncOmniEngine._strip_single_engine_args(
{
"model": "some/model",
"compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}',
"tensor_parallel_size": 4,
"custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
}
)
mock_warn.assert_called_once()
warned_args = mock_warn.call_args[0][-1] # the formatted arg list
assert "compilation_config" in warned_args
assert "tensor_parallel_size" not in warned_args
assert "model" not in warned_args
def test_new_format_deploy_yaml_does_not_strip_interleave_cli_overrides(tmp_path):
"""New-format deploy YAMLs must not strip interleave_mm_strings / media_io_kwargs.
Pytest historically passed ``--stage-configs-path`` for deploy YAMLs; stripping
those CLI flags dropped MiniCPM interleaved AV packing and cut Daily-Omni
accuracy from ~78% to ~70%.
"""
from vllm_omni.entrypoints.utils import is_new_format_deploy_config
yaml_path = tmp_path / "minicpm_deploy.yaml"
yaml_path.write_text(
"""
pipeline: minicpmo_4_5
async_chunk: true
stages:
- stage_id: 0
devices: "0"
""".strip(),
encoding="utf-8",
)
path = str(yaml_path)
assert is_new_format_deploy_config(path)
kwargs = {
"stage_configs_path": path,
"trust_remote_code": True,
"interleave_mm_strings": True,
"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}},
"gpu_memory_utilization": 0.55,
}
# Mirror AsyncOmniEngine._resolve_stage_configs strip decision.
if kwargs.get("stage_configs_path") is not None and not is_new_format_deploy_config(kwargs["stage_configs_path"]):
base = AsyncOmniEngine._strip_single_engine_args(dict(kwargs))
else:
base = dict(kwargs)
assert base.get("interleave_mm_strings") is True
assert base.get("media_io_kwargs") == {"video": {"fps": 1, "num_frames": 128}}
# Parent EngineArgs fields also remain available for deploy merge.
assert base.get("gpu_memory_utilization") == 0.55
legacy = tmp_path / "legacy.yaml"
legacy.write_text("stage_args:\n - stage_id: 0\n", encoding="utf-8")
assert not is_new_format_deploy_config(str(legacy))
legacy_kwargs = dict(kwargs, stage_configs_path=str(legacy))
stripped = AsyncOmniEngine._strip_single_engine_args(legacy_kwargs)
assert "interleave_mm_strings" not in stripped
assert "media_io_kwargs" not in stripped
assert "gpu_memory_utilization" not in stripped
# For https://github.com/vllm-project/vllm-omni/issues/3293
def test_tensor_parallel_size_none_is_handled():
"""Ensure the tensor parallel size of None isn't forwarded."""
+1 -1
View File
@@ -68,7 +68,7 @@ async def test_abort():
with ExitStack() as after:
engine = AsyncOmni(
model=model,
stage_configs_path=stage_config,
deploy_config=stage_config,
)
after.callback(engine.shutdown)
+2 -2
View File
@@ -428,7 +428,7 @@ class TestSingleStageModeDetection:
assert engine.single_stage_mode is False
assert engine._single_stage_id_filter is None
def test_stage_configs_path_loads_duplex_runtime_config(self, mocker: MockerFixture):
def test_deploy_config_loads_duplex_runtime_config(self, mocker: MockerFixture):
duplex_session = DuplexSessionRuntimeConfig(max_sessions=2)
get_pipeline_config = mocker.patch(
"vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config",
@@ -441,7 +441,7 @@ class TestSingleStageModeDetection:
engine = self._make_engine_no_thread(
mocker,
stage_configs_path="/fake/duplex.yaml",
deploy_config="/fake/duplex.yaml",
)
get_pipeline_config.assert_called_once_with(
+2
View File
@@ -139,6 +139,7 @@ def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, st
engine_extras={
"model": str(stage_model),
"attention_backend": "FLASHINFER",
"attention_config": {"use_trtllm_attention": False},
"hf_overrides": {"rope_scaling": {"factor": 2}},
"limit_mm_per_prompt": {"audio": 1},
"logits_processors": ["test.custom.LogitsProcessor"],
@@ -229,6 +230,7 @@ def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
assert thinker_args["omni_kv_config"] == {"need_send_cache": True}
assert thinker_args["worker_cls"] == "test.custom.Worker"
assert thinker_args["attention_backend"] == "FLASHINFER"
assert thinker_args["attention_config"] == {"use_trtllm_attention": False}
assert thinker_args["trust_remote_code"] is True
assert thinker_args["dtype"] == "float16"
assert thinker_args["distributed_executor_backend"] == "mp"
@@ -680,7 +680,7 @@ class TestSpeechAPI:
@pytest.mark.asyncio
async def test_create_diffusion_speech_extra_params(self, mocker: MockerFixture):
"""Test that extra_params are correctly applied to sampling_params_list in diffusion mode."""
"""Test public diffusion speech success and extra_params propagation."""
# Mock the engine client
mock_engine = mocker.MagicMock()
@@ -704,7 +704,11 @@ class TestSpeechAPI:
req = OpenAICreateSpeechRequest(input="Hello", extra_params={"new_arg": 123, "existing_arg": "new_value"})
await server._create_diffusion_speech(req)
response = await server.create_speech(req)
assert response.status_code == 200
assert response.media_type == "audio/wav"
assert response.body == b"dummy"
# Verify generate was called
mock_engine.generate.assert_called_once()
@@ -737,6 +741,19 @@ class TestTTSMethods:
yield server
server.shutdown()
def test_diffusion_audio_encode_speed_skips_ar_adapter(self, mocker: MockerFixture):
"""Pure-diffusion speech keeps generic speed handling outside AR adapters."""
server = OmniOpenAIServingSpeech.for_diffusion(
diffusion_engine=mocker.MagicMock(),
model_name="test-model",
)
resolve_adapter = mocker.patch.object(serving_speech_module, "resolve_adapter")
speed = server._audio_encode_speed(OpenAICreateSpeechRequest(input="Hello", speed=1.25))
assert speed == 1.25
resolve_adapter.assert_not_called()
def test_is_tts_detection_no_stage(self, speech_server):
"""Test TTS model detection when no TTS stage exists."""
# Fixture creates server with stage_configs = [] -> _is_tts should be False
+1 -1
View File
@@ -364,7 +364,7 @@ async def test_omni_generate_request_id():
Text modality only; asserts caller-visible ids are preserved across the
multi-stage orchestrator path on H100.
"""
engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG)
engine = AsyncOmni(model=OMNI_MODEL, deploy_config=OMNI_STAGE_CONFIG)
try:
for req_id in _OMNI_REQ_IDS:
async for output in engine.generate(
@@ -478,12 +478,11 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
)
engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
engine._strip_single_engine_args = lambda kwargs: kwargs
_, stage_configs = engine._resolve_stage_configs(
"dummy-model",
{
"stage_configs_path": "dummy.yaml",
"deploy_config": "dummy.yaml",
"additional_config": {"torchair_graph_config": {"enabled": True}},
},
trust_remote_code=False,
@@ -516,12 +515,11 @@ def test_resolve_stage_configs_injects_quantization_config_into_diffusion_stage(
)
engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
engine._strip_single_engine_args = lambda kwargs: kwargs
_, stage_configs = engine._resolve_stage_configs(
"dummy-model",
{
"stage_configs_path": "dummy.yaml",
"deploy_config": "dummy.yaml",
"quantization_config": {"method": "bitsandbytes"},
},
trust_remote_code=False,
+1 -1
View File
@@ -193,7 +193,7 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture)
with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
run_headless(args)
assert captured["args"][1] is None
assert len(captured["args"]) == 2
assert captured["deploy_config_path"] == "/tmp/deploy.yaml"
assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
assert captured["strategy_config_path"] == "/tmp/strategy.yaml"
+9 -7
View File
@@ -381,7 +381,6 @@ class TestLoadAndResolveStageConfigs:
config_path, stage_configs, _ = load_and_resolve_stage_configs(
model="black-forest-labs/FLUX.2-klein-4B",
stage_configs_path=None,
kwargs=kwargs,
trust_remote_code=False,
default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs),
@@ -390,9 +389,7 @@ class TestLoadAndResolveStageConfigs:
assert len(stage_configs) == 1
assert "dtype" in stage_configs[0]["engine_args"]
def test_stage_configs_path_promotes_new_deploy_yaml_without_expanding_replicas(
self, tmp_path, mocker: MockerFixture
):
def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture):
deploy_path = tmp_path / "qwen3_multi.yaml"
deploy_path.write_text(
'stages:\n - stage_id: 0\n devices: "0"\n - stage_id: 1\n devices: "1,2,3"\n num_replicas: 3\n',
@@ -413,18 +410,23 @@ class TestLoadAndResolveStageConfigs:
"vllm_omni.entrypoints.utils.load_stage_configs_from_model",
return_value=(returned_stage_configs, None),
)
cli_overrides = {
"interleave_mm_strings": True,
"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}},
"gpu_memory_utilization": 0.55,
}
config_path, stage_configs, _ = load_and_resolve_stage_configs(
model="dummy-model",
stage_configs_path=str(deploy_path),
kwargs={},
kwargs=cli_overrides,
trust_remote_code=True,
deploy_config_path=str(deploy_path),
)
load_stage_configs.assert_called_once_with(
"dummy-model",
trust_remote_code=True,
base_engine_args={},
base_engine_args=cli_overrides,
deploy_config_path=str(deploy_path),
stage_overrides=None,
strategy_config_path=None,
@@ -0,0 +1,91 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import pytest
import torch
from examples.offline_inference.text_to_video.text_to_video import (
_MODEL_PRESETS,
_detect_preset,
_normalize_float_tensor,
build_text_to_video_prompt,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.mark.parametrize(
("negative_prompt", "expected_negative_prompt"),
[(None, None), ("", ""), ("blurry", "blurry")],
)
def test_text_to_video_builds_canonical_prompt(
negative_prompt: str | None,
expected_negative_prompt: str | None,
) -> None:
result = build_text_to_video_prompt("a robot waves", negative_prompt)
assert result["prompt"] == "a robot waves"
assert result["modalities"] == ["video"]
if expected_negative_prompt is None:
assert "negative_prompt" not in result
else:
assert result["negative_prompt"] == expected_negative_prompt
@pytest.mark.parametrize(
("model", "model_class_name", "preset_name"),
[
("robbyant/lingbot-video-dense-1.3b", None, "lingbot"),
("/models/custom-checkpoint", "LingBotVideoPipeline", "lingbot"),
("Lightricks/LTX-2", "LTX2DistilledPipeline", "ltx2_distilled"),
("BestWishYsh/Helios-Distilled", None, "helios"),
("/models/cosmos-edge", "Cosmos3EdgeVFMTransformer", "cosmos3_edge"),
("/models/vace", "WanVACEPipeline", "vace"),
],
)
def test_detect_preset_is_scoped_to_model_family(
model: str,
model_class_name: str | None,
preset_name: str,
) -> None:
assert _detect_preset(model, model_class_name) is _MODEL_PRESETS[preset_name]
def test_lingbot_preset_matches_lingbot_defaults() -> None:
assert _MODEL_PRESETS["lingbot"] == {
"model_class_name": "LingBotVideoPipeline",
"height": 192,
"width": 320,
"num_frames": 9,
"num_inference_steps": 2,
"guidance_scale": 3.0,
"fps": 24,
"flow_shift": 3.0,
"output": "lingbot_video_output.mp4",
}
def test_normalize_float_tensor_preserves_zero_to_one_frames() -> None:
frames = torch.tensor([0.0, 0.25, 1.0, 1.5])
assert torch.equal(_normalize_float_tensor(frames, "zero_to_one"), torch.tensor([0.0, 0.25, 1.0, 1.0]))
def test_normalize_float_tensor_maps_negative_one_to_one_frames() -> None:
frames = torch.tensor([-1.5, -1.0, 0.0, 1.0, 1.5])
assert torch.equal(
_normalize_float_tensor(frames, "negative_one_to_one"),
torch.tensor([0.0, 0.0, 0.5, 1.0, 1.0]),
)
def test_normalize_float_tensor_uses_one_contract_for_ambiguous_frames() -> None:
frames = [torch.tensor([0.0, 0.5, 1.0]), torch.tensor([-1.0, 0.0, 1.0])]
normalized = [_normalize_float_tensor(frame, "negative_one_to_one") for frame in frames]
assert torch.equal(normalized[0], torch.tensor([0.5, 0.75, 1.0]))
assert torch.equal(normalized[1], torch.tensor([0.0, 0.5, 1.0]))
+4 -9
View File
@@ -2635,7 +2635,7 @@ class OmniRunner:
# affects the test runner wrapper.
init_timeout: int = 1800,
log_stats: bool = False,
stage_configs_path: str | None = None,
deploy_config: str | None = None,
**kwargs,
) -> None:
startup_t0 = time.perf_counter()
@@ -2653,7 +2653,7 @@ class OmniRunner:
stage_init_timeout=stage_init_timeout,
batch_timeout=batch_timeout,
init_timeout=init_timeout,
stage_configs_path=stage_configs_path,
deploy_config=deploy_config,
**kwargs,
)
startup_s = time.perf_counter() - startup_t0
@@ -3179,12 +3179,7 @@ def iter_omni_server(
print("OmniServer stopping...")
else:
if stage_config_path is not None:
from vllm_omni.entrypoints.utils import is_new_format_deploy_config
if is_new_format_deploy_config(stage_config_path):
server_args += ["--deploy-config", stage_config_path]
else:
server_args += ["--stage-configs-path", stage_config_path]
server_args += ["--deploy-config", stage_config_path]
with (
OmniServer(
@@ -3237,7 +3232,7 @@ def iter_omni_runner(
model = model_prefix + model
if run_level == "core_model" and request.node.get_closest_marker("diffusion"):
model = resolve_tiny_model_path(model)
with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner:
with OmniRunner(model, seed=42, deploy_config=stage_config_path, **extra_omni_kwargs) as runner:
print("OmniRunner started successfully")
yield runner
print("OmniRunner stopping...")
+147 -3
View File
@@ -42,6 +42,10 @@ _EXPECTED_FAMILIES = [
defs.AUDIO_UNDERRUN_S,
defs.AUDIO_CONTINUITY_OK_METRIC,
defs.AUDIO_SKIPPED_REQUESTS_METRIC,
defs.DIFFUSION_EXEC_S,
defs.DIFFUSION_EXEC_PER_STEP_S,
defs.DIFFUSION_PREPROCESS_S,
defs.DIFFUSION_POSTPROCESS_S,
]
@@ -55,6 +59,10 @@ class TestRegistration:
mod.observe_audio_underrun("s", "r", 0.01)
mod.inc_audio_continuity_ok("s", "r", 100)
mod.inc_audio_skipped("s", "r", "malformed_codec")
mod.observe_diffusion_exec("s", "r", 0.5)
mod.observe_diffusion_exec_per_step("s", "r", 0.01)
mod.observe_diffusion_preprocess("s", "r", 0.01)
mod.observe_diffusion_postprocess("s", "r", 0.2)
out = generate_latest(REGISTRY).decode()
for name in _EXPECTED_FAMILIES:
@@ -156,6 +164,50 @@ class TestAudio:
assert unknown == 1.0
# ---------------------------------------------------------------------------
# Diffusion observe API
# ---------------------------------------------------------------------------
class TestDiffusion:
def test_diffusion_exec_observed(self, mod: OmniModalityMetrics) -> None:
stage, replica = "dit_exec", "0"
mod.observe_diffusion_exec(stage, replica, 1.25)
out = generate_latest(REGISTRY).decode()
prefix = f'{defs.DIFFUSION_EXEC_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
assert _sample_value(out, prefix) == 1.25
def test_diffusion_preprocess_observed(self, mod: OmniModalityMetrics) -> None:
stage, replica = "dit_pre", "0"
mod.observe_diffusion_preprocess(stage, replica, 0.03)
out = generate_latest(REGISTRY).decode()
prefix = f'{defs.DIFFUSION_PREPROCESS_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
assert _sample_value(out, prefix) == pytest.approx(0.03)
def test_diffusion_postprocess_observed(self, mod: OmniModalityMetrics) -> None:
stage, replica = "dit_post", "0"
mod.observe_diffusion_postprocess(stage, replica, 0.8)
out = generate_latest(REGISTRY).decode()
prefix = f'{defs.DIFFUSION_POSTPROCESS_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
assert _sample_value(out, prefix) == 0.8
def test_diffusion_exec_per_step_observed(self, mod: OmniModalityMetrics) -> None:
stage, replica = "dit_per_step", "0"
mod.observe_diffusion_exec_per_step(stage, replica, 0.078)
out = generate_latest(REGISTRY).decode()
prefix = f'{defs.DIFFUSION_EXEC_PER_STEP_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
assert _sample_value(out, prefix) == pytest.approx(0.078)
def test_log_stats_false_skips(self) -> None:
silent = OmniModalityMetrics(model_name="silent-model", log_stats=False)
silent.observe_diffusion_exec("s", "r", 1.0)
silent.observe_diffusion_exec_per_step("s", "r", 0.02)
silent.observe_diffusion_preprocess("s", "r", 0.01)
silent.observe_diffusion_postprocess("s", "r", 0.5)
out = generate_latest(REGISTRY).decode()
assert 'model_name="silent-model"' not in out
# ---------------------------------------------------------------------------
# Stub-driven routing tests for the finalize / streaming helpers.
# ---------------------------------------------------------------------------
@@ -183,6 +235,18 @@ class _StubModMetrics:
def inc_audio_skipped(self, s, r, reason):
self.calls.append(("inc_audio_skipped", s, r, reason))
def observe_diffusion_exec(self, s, r, v):
self.calls.append(("observe_diffusion_exec", s, r, v))
def observe_diffusion_exec_per_step(self, s, r, v):
self.calls.append(("observe_diffusion_exec_per_step", s, r, v))
def observe_diffusion_preprocess(self, s, r, v):
self.calls.append(("observe_diffusion_preprocess", s, r, v))
def observe_diffusion_postprocess(self, s, r, v):
self.calls.append(("observe_diffusion_postprocess", s, r, v))
class _Bag:
def __init__(self, **kw):
@@ -239,8 +303,88 @@ class TestObserveModalityAtFinalize:
)
assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls
def test_non_audio_output_type_skipped(self):
# Image / video / text output types are out of scope for this module.
def test_diffusion_path_full(self):
stub = _StubModMetrics()
stage_metrics = _Bag(
stage_gen_time_ms=1000.0,
diffusion_metrics={
"preprocess_time_s": 0.05,
"diffusion_engine_exec_time_s": 1.2,
"postprocess_time_s": 0.3,
"image_num": 1,
"resolution": 1024,
"num_inference_steps": 50,
},
)
observe_modality_at_finalize(
stub,
output_type="image",
stage_id=2,
replica_id=0,
stage_metrics=stage_metrics,
engine_outputs=_Bag(),
)
assert ("observe_diffusion_preprocess", "2", "0", 0.05) in stub.calls
assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls
assert ("observe_diffusion_postprocess", "2", "0", 0.3) in stub.calls
assert ("observe_diffusion_exec_per_step", "2", "0", pytest.approx(0.024)) in stub.calls
def test_diffusion_per_step_skipped_without_num_inference_steps(self):
stub = _StubModMetrics()
stage_metrics = _Bag(
stage_gen_time_ms=1000.0,
diffusion_metrics={
"diffusion_engine_exec_time_s": 1.2,
"preprocess_time_s": 0.05,
"postprocess_time_s": 0.3,
},
)
observe_modality_at_finalize(
stub,
output_type="image",
stage_id=2,
replica_id=0,
stage_metrics=stage_metrics,
engine_outputs=_Bag(),
)
assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls
assert not any(c[0] == "observe_diffusion_exec_per_step" for c in stub.calls)
def test_diffusion_metrics_none_skipped(self):
stub = _StubModMetrics()
observe_modality_at_finalize(
stub,
output_type="image",
stage_id=2,
replica_id=0,
stage_metrics=_Bag(stage_gen_time_ms=500.0, diffusion_metrics=None),
engine_outputs=_Bag(),
)
assert not any(c[0].startswith("observe_diffusion") for c in stub.calls)
def test_diffusion_zero_values_observed(self):
stub = _StubModMetrics()
stage_metrics = _Bag(
stage_gen_time_ms=500.0,
diffusion_metrics={
"preprocess_time_s": 0.0,
"diffusion_engine_exec_time_s": 0.0,
"postprocess_time_s": 0.0,
},
)
observe_modality_at_finalize(
stub,
output_type="image",
stage_id=2,
replica_id=0,
stage_metrics=stage_metrics,
engine_outputs=_Bag(),
)
assert ("observe_diffusion_preprocess", "2", "0", 0.0) in stub.calls
assert ("observe_diffusion_exec", "2", "0", 0.0) in stub.calls
assert ("observe_diffusion_postprocess", "2", "0", 0.0) in stub.calls
def test_non_audio_output_type_skips_audio_but_still_checks_diffusion(self):
stub = _StubModMetrics()
for output_type in ("text", "image", "video"):
observe_modality_at_finalize(
@@ -251,7 +395,7 @@ class TestObserveModalityAtFinalize:
stage_metrics=_Bag(stage_gen_time_ms=100.0),
engine_outputs=_Bag(),
)
assert stub.calls == []
assert not any(c[0].startswith("observe_audio") or c[0].startswith("inc_audio") for c in stub.calls)
def test_replica_id_none_skipped(self):
stub = _StubModMetrics()
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -0,0 +1,137 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Regression tests for the dots.tts talker's stop-signal pairing and
per-request state lifecycle.
Pure-CPU tests on a bare talker (``__new__`` + the handful of attributes
the methods under test read) — no weights, no GPU, no engine. Mirrors
tests/model_executor/models/voxcpm2/test_talker_state_eviction.py.
"""
from __future__ import annotations
import functools
from types import SimpleNamespace
import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
torch = pytest.importorskip("torch")
_VOCAB_SIZE = 151672
@functools.lru_cache(maxsize=1)
def _dots_tts_talker_mod():
"""Defer talker import (pulls vLLM model_executor) until first use."""
from vllm_omni.model_executor.models.dots_tts.dots_tts_talker import (
DotsTTSForConditionalGeneration,
_RequestState,
)
return DotsTTSForConditionalGeneration, _RequestState
def _make_bare_talker():
DotsTTSForConditionalGeneration, _ = _dots_tts_talker_mod()
talker = DotsTTSForConditionalGeneration.__new__(DotsTTSForConditionalGeneration)
talker.config = SimpleNamespace(vocab_size=_VOCAB_SIZE)
talker._active_states = {}
talker._pending_requests = []
talker._results_queue = []
talker._audio_queue = []
talker._deferred_cleanup_ids = set()
talker._beta_trace = False
return talker
def _add_state(talker, req_id: str, *, is_stopping: bool = False):
_, _RequestState = _dots_tts_talker_mod()
state = _RequestState(request_id=req_id)
state.prefill_completed = True
state.is_stopping = is_stopping
talker._active_states[req_id] = state
return state
def _hidden(bsz: int) -> torch.Tensor:
return torch.zeros(bsz, 1536, dtype=torch.float32)
class TestComputeLogitsPairing:
"""compute_logits pairs _results_queue entries with batch rows
positionally, so every request must push exactly one entry per step —
including prefills, which push a (req_id, None) placeholder."""
def test_mixed_batch_prefill_placeholder_keeps_rows_aligned(self) -> None:
talker = _make_bare_talker()
_add_state(talker, "prefill-req")
decode_state = _add_state(talker, "decode-req", is_stopping=True)
decode_state.precomputed_stop_logits = torch.tensor([[0.1, 0.9]])
# Row 0 is the prefill (None placeholder), row 1 the stopping decode.
talker._results_queue = [
("prefill-req", None),
("decode-req", torch.tensor([[0.1, 0.9]])),
]
logits = talker.compute_logits(_hidden(2))
assert logits.shape == (2, _VOCAB_SIZE)
# Prefill row: forced continue.
assert logits[0, 0].item() == 1.0
assert logits[0, 1].item() == float("-inf")
# Stopping decode row: stop slot wins.
assert logits[1, 1].item() == 1.0
assert logits[1, 0].item() == 0.0
# Everything else stays masked.
assert logits[0, 2].item() == float("-inf")
assert logits[1, 2].item() == float("-inf")
# Queue drained.
assert talker._results_queue == []
def test_below_threshold_stop_logits_force_continue(self) -> None:
"""Raw (continue, stop) softmax must not reach the sampler: a
prob_stop in (0.5, 0.8] would flip greedy argmax even though
is_stopping stayed False (the 0.8-threshold contract)."""
talker = _make_bare_talker()
_add_state(talker, "req", is_stopping=False)
talker._results_queue = [("req", torch.tensor([[0.3, 0.7]]))]
logits = talker.compute_logits(_hidden(1))
assert logits[0, 0].item() == 1.0
assert logits[0, 1].item() == float("-inf")
def test_empty_queue_defaults_to_all_continue(self) -> None:
talker = _make_bare_talker()
logits = talker.compute_logits(_hidden(3))
assert torch.all(logits[:, 0] == 1.0)
assert torch.all(logits[:, 1] == float("-inf"))
class TestStateEviction:
"""on_requests_finished fires BEFORE forward(), so eviction must be
deferred to _flush_deferred_cleanup at the end of the step — the
finishing request's audio still drains through the current forward."""
def test_finished_request_survives_until_flush(self) -> None:
talker = _make_bare_talker()
_add_state(talker, "req-a")
_add_state(talker, "req-b")
talker.on_requests_finished(["req-a"])
assert "req-a" in talker._active_states # not evicted yet
talker._flush_deferred_cleanup()
assert "req-a" not in talker._active_states
assert "req-b" in talker._active_states
assert talker._deferred_cleanup_ids == set()
def test_unknown_request_id_is_ignored(self) -> None:
talker = _make_bare_talker()
talker.on_requests_finished(["never-seen"])
talker._flush_deferred_cleanup()
assert talker._active_states == {}
@@ -1326,7 +1326,9 @@ class TestRegistry:
assert "higgs_multimodal_qwen3" in OMNI_PIPELINES
def test_deploy_yaml_exists(self):
def test_deploy_yaml_uses_native_flashinfer(self):
from vllm_omni.config.stage_config import load_deploy_config
for name in (
"higgs_multimodal_qwen3.yaml",
"higgs_multimodal_qwen3_high_throughput.yaml",
@@ -1334,6 +1336,10 @@ class TestRegistry:
):
yaml_path = get_deploy_config_path(name)
assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}"
deploy_config = load_deploy_config(yaml_path)
stage_0 = next(stage for stage in deploy_config.stages if stage.stage_id == 0)
assert stage_0.engine_extras["attention_backend"] == "FLASHINFER"
assert stage_0.engine_extras["attention_config"] == {"use_trtllm_attention": False}
# ---- AC-3: Prompt Builder ----
@@ -1,21 +1,24 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from functools import lru_cache
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
import torch.nn as nn
from vllm.platforms import current_platform
from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
HiFTGenerator,
)
from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import (
CFMGraphWrapper,
HiFTGraphWrapper,
)
pytestmark = [pytest.mark.core_model]
pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
class _F0Predictor(nn.Module):
@@ -142,3 +145,248 @@ def test_lazy_capture_limit_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch)
wrapper._capture.assert_not_called()
wrapper.decode_fn.assert_called_once_with(speech_feat, cache_source)
assert result is wrapper.decode_fn.return_value
# ---------------------------------------------------------------------------
# CFMGraphWrapper tests
# ---------------------------------------------------------------------------
class _MiniDiT(nn.Module):
"""Minimal DiT-like module with blocks_forward_chunk for CFM graph testing.
Mimics the upstream cosyvoice2 DiT's cache semantics:
- CausalConv1d.forward_chunk: cat([cnn_cache, x], dim=time) when cache is not None
- Attention.forward_chunk: cat([k, k_cache], dim=seq) when att_cache is not None
"""
def __init__(self, in_dim: int = 16, hidden: int = 8, depth: int = 2) -> None:
super().__init__()
self.in_proj = nn.Linear(in_dim, hidden)
self.blocks = nn.ModuleList([nn.LayerNorm(hidden) for _ in range(depth)])
self.final_layer = nn.Linear(hidden, in_dim)
def blocks_forward_chunk(
self,
x: torch.Tensor,
t: torch.Tensor,
mask: torch.Tensor | None,
cnn_cache: torch.Tensor | None = None,
att_cache: torch.Tensor | None = None,
cnn_cache_buffer: torch.Tensor | None = None,
att_cache_buffer: torch.Tensor | None = None,
) -> torch.Tensor:
x = x.transpose(1, 2)
x = self.in_proj(x)
for b_idx in range(len(self.blocks)):
# Simulate CausalConv1d: cnn_cache contributes to the first frames.
# Upstream: if cnn_cache[b_idx] is None, creates zeros internally (same as zeros).
cnn_b = cnn_cache[b_idx] if cnn_cache is not None else None
if cnn_b is not None:
x[:, : cnn_b.shape[2], :] += cnn_b.transpose(1, 2)
# Simulate Attention: att_cache contributes a bias when non-empty.
# Upstream: if att_cache[b_idx] is None, skips cat entirely (different path).
att_b = att_cache[b_idx] if att_cache is not None else None
if att_b is not None and att_b.shape[3] > 0:
x += att_b.sum(dim=(1, 2), keepdim=False).unsqueeze(1)
x = self.blocks[b_idx](x)
x = x + t
cnn_cache_buffer[b_idx] = x[:, -2:, :].transpose(1, 2).contiguous()
dt = x.shape[1]
att_cache_buffer[b_idx][:, :, :dt, :] = x.unsqueeze(1)
x = self.final_layer(x)
x = x.transpose(1, 2)
return x
def _cfm_inputs(
batch_size: int, chunk_size: int, old_att_len: int, *, device: str = "cuda"
) -> tuple[torch.Tensor, ...]:
depth = 2
hidden = 8
estimator_input = torch.randn(batch_size, 16, chunk_size, device=device)
time_emb = torch.randn(batch_size, 1, hidden, device=device)
cnn_cache = torch.randn(depth, batch_size, hidden, 2, device=device)
att_cache = torch.randn(depth, batch_size, 1, old_att_len, hidden, device=device)
cnn_out = torch.empty(depth, batch_size, hidden, 2, device=device)
att_out = torch.empty(depth, batch_size, 1, old_att_len + chunk_size, hidden, device=device)
return estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes(
monkeypatch: pytest.MonkeyPatch,
) -> None:
pool = torch.cuda.graph_pool_handle()
monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
torch.manual_seed(0)
estimator = _MiniDiT().eval().cuda()
wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=32)
with torch.inference_mode():
for _, chunk_size, old_att_len in ((2, 10, 0), (2, 10, 5)):
inputs = _cfm_inputs(2, chunk_size, old_att_len)
eager_inputs = tuple(v.clone() for v in inputs)
with torch.no_grad():
eager_result = estimator.blocks_forward_chunk(
eager_inputs[0],
eager_inputs[1],
None,
eager_inputs[2],
eager_inputs[3],
eager_inputs[4],
eager_inputs[5],
)
wrapper.replay(*inputs)
replay_inputs = tuple(v.clone() for v in inputs)
graph_result, graph_cnn, graph_att = wrapper.replay(*replay_inputs)
torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5)
torch.testing.assert_close(graph_cnn, eager_inputs[4], rtol=1e-4, atol=1e-5)
torch.testing.assert_close(graph_att, eager_inputs[5], rtol=1e-4, atol=1e-5)
wrapper._capture_graph.cache_clear()
def _cfm_mock_wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 1) -> CFMGraphWrapper:
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
wrapper = object.__new__(CFMGraphWrapper)
wrapper.max_graphs = max_graphs
wrapper.graph_fn = Mock(return_value=torch.tensor([42.0]))
wrapper.device = torch.device("cuda")
@lru_cache(maxsize=max_graphs)
def _capture_graph(key: tuple):
return wrapper._do_capture(key)
wrapper._capture_graph = _capture_graph
wrapper._do_capture = Mock(return_value=None)
return wrapper
def test_cfm_unseen_shape_is_lazily_captured(monkeypatch: pytest.MonkeyPatch) -> None:
wrapper = _cfm_mock_wrapper(monkeypatch)
wrapper._capture_graph = Mock(return_value=None)
inputs = _cfm_inputs(2, 10, 0)
wrapper.replay(*inputs)
wrapper._capture_graph.assert_called_once()
wrapper.graph_fn.assert_called_once()
def test_cfm_capture_failure_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None:
wrapper = _cfm_mock_wrapper(monkeypatch)
wrapper._capture_graph = lru_cache(maxsize=1)(lambda key: None)
inputs = _cfm_inputs(2, 10, 0)
result = wrapper.replay(*inputs)
wrapper.graph_fn.assert_called_once()
assert result[0] is wrapper.graph_fn.return_value
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None:
"""Verify that LRU cache evicts oldest entries when full."""
pool = torch.cuda.graph_pool_handle()
monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
torch.manual_seed(0)
estimator = _MiniDiT().eval().cuda()
wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=2)
with torch.inference_mode():
inputs_a = _cfm_inputs(2, 10, 0)
wrapper.replay(*inputs_a)
assert wrapper._capture_graph.cache_info().currsize == 1
inputs_b = _cfm_inputs(2, 12, 0)
wrapper.replay(*inputs_b)
assert wrapper._capture_graph.cache_info().currsize == 2
wrapper.replay(*inputs_a)
assert wrapper._capture_graph.cache_info().currsize == 2
inputs_c = _cfm_inputs(2, 14, 0)
wrapper.replay(*inputs_c)
assert wrapper._capture_graph.cache_info().currsize == 2
wrapper.replay(*inputs_b)
hits_before = wrapper._capture_graph.cache_info().hits
wrapper.replay(*inputs_b)
hits_after = wrapper._capture_graph.cache_info().hits
assert hits_after == hits_before + 1
wrapper._capture_graph.cache_clear()
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_cfm_none_cache_parity_between_graph_and_eager(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Verify that None→zeros (graph path) matches [None]*depth (eager path).
setup_batch calls _decode_cfm with cnn_cache=None on every request's
prompt pass. The graph path replaces None with zeros; the eager path
passes [None]*depth. Both must produce identical outputs.
"""
torch.accelerator.synchronize()
pool = torch.cuda.graph_pool_handle()
monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
torch.manual_seed(0)
estimator = _MiniDiT().eval().cuda()
wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=32)
depth = len(estimator.blocks)
hidden = 8
batch_size = 2
chunk_size = 20
estimator_input = torch.randn(batch_size, 16, chunk_size, device="cuda")
time_emb = torch.randn(batch_size, 1, hidden, device="cuda")
cnn_out = torch.empty(depth, batch_size, hidden, 2, device="cuda")
att_out = torch.empty(depth, batch_size, 1, chunk_size, hidden, device="cuda")
# Eager path: cnn_cache=[None]*depth, att_cache=[None]*depth
eager_input = estimator_input.clone()
eager_time = time_emb.clone()
eager_cnn_out = torch.empty_like(cnn_out)
eager_att_out = torch.empty_like(att_out)
with torch.no_grad():
eager_result = estimator.blocks_forward_chunk(
eager_input,
eager_time,
None,
[None] * depth,
[None] * depth,
eager_cnn_out,
eager_att_out,
)
# Graph path: cnn_cache=zeros, att_cache=zero-length (simulating _estimator_step's None→zeros)
zero_cnn = torch.zeros_like(cnn_out)
zero_att = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:])
wrapper.replay(estimator_input, time_emb, zero_cnn, zero_att, cnn_out, att_out)
graph_input = estimator_input.clone()
graph_time = time_emb.clone()
graph_cnn_out = torch.empty_like(cnn_out)
graph_att_out = torch.empty_like(att_out)
graph_zero_cnn = torch.zeros_like(cnn_out)
graph_zero_att = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:])
graph_result, graph_cnn, graph_att = wrapper.replay(
graph_input,
graph_time,
graph_zero_cnn,
graph_zero_att,
graph_cnn_out,
graph_att_out,
)
torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5)
torch.testing.assert_close(graph_cnn, eager_cnn_out, rtol=1e-4, atol=1e-5)
torch.testing.assert_close(graph_att, eager_att_out, rtol=1e-4, atol=1e-5)
+10 -4
View File
@@ -17,6 +17,7 @@ from vllm_omni.model_extras import (
build_x_to_text_prompt,
get_extra_body_params,
get_extra_output_params,
get_output_tensor_range,
get_x_to_text_model_family,
should_init_extra_args_for_non_diffusion_stages,
)
@@ -134,6 +135,8 @@ def test_lingbot_extra_registry_declares_request_params() -> None:
}
)
assert get_extra_output_params("LingBotVideoPipeline") == frozenset()
assert get_output_tensor_range("LingBotVideoPipeline") == "zero_to_one"
assert get_output_tensor_range("WanPipeline") == "negative_one_to_one"
assert should_init_extra_args_for_non_diffusion_stages("LingBotVideoPipeline") is False
@@ -230,10 +233,13 @@ def test_ltx_extra_registry_declares_official_guidance_params() -> None:
}
)
assert get_extra_body_params("LTX2Pipeline") == expected
assert get_extra_output_params("LTX2Pipeline") == frozenset()
assert get_extra_body_params("LTX2DistilledPipeline") == expected
assert get_extra_output_params("LTX2DistilledPipeline") == frozenset()
for pipeline_name in (
"LTX2Pipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
):
assert get_extra_body_params(pipeline_name) == expected
assert get_extra_output_params(pipeline_name) == frozenset()
@pytest.mark.core_model
@@ -13,10 +13,13 @@ def _shrink_dit_rope_config(
target_head_dim: int = 32,
target_heads: int = 4,
default_axes_dims_rope: list[int] | None = None,
joint_attention_dim: int | None = None,
) -> dict:
config["num_layers"] = num_layers
if num_single_layers is not None:
config["num_single_layers"] = num_single_layers
if joint_attention_dim is not None:
config["joint_attention_dim"] = joint_attention_dim
axes_dims_rope = config.get("axes_dims_rope", default_axes_dims_rope)
factor = config["attention_head_dim"] / target_head_dim
config["attention_head_dim"] = target_head_dim
@@ -50,22 +53,55 @@ def _shrink_flux_t5_text_encoder(config: dict) -> dict:
return config
def tiny_qwen_image_builder() -> str:
def shrink_text_encoder(config: dict) -> dict:
config["num_hidden_layers"] = 2
config["intermediate_size"] = 64
config["text_config"]["num_hidden_layers"] = 2
config["text_config"]["intermediate_size"] = 64
config["text_config"]["layer_types"] = config["text_config"]["layer_types"][:2]
config["vision_config"]["depth"] = 2
config["vision_config"]["intermediate_size"] = 64
config["vision_config"]["fullatt_block_indexes"] = [0, 1]
return config
def _shrink_qwen_text_encoder_config(config: dict, hidden_size: int = 64) -> dict:
text_config = config["text_config"]
old_head_dim = text_config["hidden_size"] / text_config["num_attention_heads"]
text_config["num_hidden_layers"] = 2
text_config["intermediate_size"] = 64
text_config["hidden_size"] = hidden_size
text_config["num_attention_heads"] = 2
text_config["num_key_value_heads"] = 2
text_config["layer_types"] = text_config["layer_types"][:2]
factor = old_head_dim / (hidden_size / 2)
mrope_section = text_config["rope_scaling"]["mrope_section"]
text_config["rope_scaling"]["mrope_section"] = [round(d / factor) for d in mrope_section]
config["num_hidden_layers"] = 2
config["intermediate_size"] = 64
config["hidden_size"] = hidden_size
config["num_attention_heads"] = 2
config["num_key_value_heads"] = 2
config["vision_config"]["depth"] = 2
config["vision_config"]["intermediate_size"] = 64
config["vision_config"]["fullatt_block_indexes"] = [0, 1]
config["vision_config"]["out_hidden_size"] = hidden_size
return config
def _shrink_qwen_transformer_config(config: dict, joint_attention_dim: int = 64) -> dict:
return _shrink_dit_rope_config(config, joint_attention_dim=joint_attention_dim)
def tiny_qwen_image_builder() -> str:
return build_tiny_from_configs(
"QwenImagePipeline",
"Qwen/Qwen-Image",
transform={"text_encoder": shrink_text_encoder, "transformer": _shrink_dit_rope_config},
transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
)
def tiny_qwen_image_edit_builder() -> str:
return build_tiny_from_configs(
"QwenImageEditPipeline",
"Qwen/Qwen-Image-Edit",
transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
)
def tiny_qwen_image_edit_plus_builder() -> str:
return build_tiny_from_configs(
"QwenImageEditPlusPipeline",
"Qwen/Qwen-Image-Edit-2511",
transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
)
@@ -74,4 +74,14 @@ DIFFUSION_TEST_SETTINGS = {
builder=diff_model_builders.tiny_flux2_builder,
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
),
"QwenImageEditPipeline": DiffusionModelTestOpts(
model="Qwen/Qwen-Image-Edit",
builder=diff_model_builders.tiny_qwen_image_edit_builder,
supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
),
"QwenImageEditPlusPipeline": DiffusionModelTestOpts(
model="Qwen/Qwen-Image-Edit-2511",
builder=diff_model_builders.tiny_qwen_image_edit_plus_builder,
supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
),
}
@@ -21,14 +21,13 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
# if you are adding a new model and see the tests below fail, please follow the pattern
# for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS.
EXCLUDED_MODELS = [
"QwenImageEditPipeline",
"QwenImageEditPlusPipeline",
"QwenImageLayeredPipeline",
"GlmImagePipeline",
"ZImagePipeline",
"OvisImagePipeline",
"WanPipeline",
"WanVACEPipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
+1 -1
View File
@@ -147,7 +147,7 @@ def _run_once(model: str, yaml_path: str, *, async_mode: bool) -> list[dict]:
# (see tests/e2e/accuracy/qwen3_omni). Confirm the exact knob at bench time.
omni = Omni(
model,
stage_configs_path=yaml_path,
deploy_config=yaml_path,
stage_init_timeout=1800,
init_timeout=3600,
async_chunk=async_mode,
+2 -1
View File
@@ -53,7 +53,6 @@ _NON_STAGE_ENGINE_CLI_FIELDS = frozenset(
"model",
"omni",
"output_modalities",
"stage_configs_path",
"stage_id",
"tokenizer",
}
@@ -91,6 +90,7 @@ class _ModelEngineOverrides(TypedDict, total=False):
trust_remote_code: bool
dtype: Any
attention_backend: Any
attention_config: Any
moe_backend: str
hf_overrides: Any
limit_mm_per_prompt: dict[str, Any]
@@ -322,6 +322,7 @@ class OmniStageModelConfig:
trust_remote_code: bool = False
dtype: Any = "auto"
attention_backend: Any = None
attention_config: Any = None
moe_backend: str = "auto"
hf_overrides: Any = None
limit_mm_per_prompt: dict[str, Any] | None = None
+2
View File
@@ -47,6 +47,7 @@ from vllm_omni.model_executor.models.bagel.pipeline import (
)
from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE
from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE
from vllm_omni.model_executor.models.dots_tts.pipeline import DOTS_TTS_PIPELINE
from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE
from vllm_omni.model_executor.models.dynin_omni.pipeline import DYNIN_OMNI_PIPELINE
from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE
@@ -144,6 +145,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
"hunyuan_video_15": HUNYUAN_VIDEO_15_PIPELINE,
"wan2_2_ti2v": WAN2_2_TI2V_PIPELINE,
"voxcpm2": VOXCPM2_PIPELINE,
"dots_tts": DOTS_TTS_PIPELINE,
"cosyvoice3": COSYVOICE3_PIPELINE,
"audex_tts": AUDEX_TTS_PIPELINE,
"audex_tta": AUDEX_TTA_PIPELINE,
+1 -42
View File
@@ -4,7 +4,6 @@
from __future__ import annotations
import dataclasses
import functools
import re
import warnings
@@ -56,7 +55,7 @@ def build_stage_runtime_overrides(
``internal_keys`` defaults to the union of
``arg_utils.internal_blacklist_keys()`` and ``arg_utils.SHARED_FIELDS``
so that neither orchestrator-only fields nor shared-pipeline fields
(``model`` / ``stage_configs_path`` / ``log_stats`` / ``stage_id``) leak
(``model`` / ``log_stats`` / ``stage_id``) leak
into a stage's per-stage runtime overrides — the orchestrator sets those
uniformly for every stage, they are not per-stage knobs. Callers can
pass an explicit set for tests or specialized flows.
@@ -90,46 +89,6 @@ def build_stage_runtime_overrides(
return result
def strip_parent_engine_args(
kwargs: dict[str, Any],
*,
parent_fields: dict[str, dataclasses.Field],
keep_keys: set[str] | frozenset[str] = frozenset(),
strip_keys: set[str] | frozenset[str] = frozenset(),
no_warn_keys: set[str] | frozenset[str] = frozenset(),
) -> tuple[dict[str, Any], list[str]]:
"""Strip parent ``EngineArgs`` fields before merging into stage YAML."""
overridden: list[str] = []
result: dict[str, Any] = {}
for key, value in kwargs.items():
if key in strip_keys:
continue
if key not in parent_fields or key in keep_keys:
result[key] = value
continue
field_def = parent_fields[key]
if field_def.default is not dataclasses.MISSING:
default = field_def.default
elif field_def.default_factory is not dataclasses.MISSING:
default = field_def.default_factory()
else:
default = dataclasses.MISSING
if default is dataclasses.MISSING or value is None:
continue
if dataclasses.is_dataclass(default) and not isinstance(default, type):
default = asdict(default)
if value != default and key not in no_warn_keys:
overridden.append(key)
return result, sorted(overridden)
def _apply_diffusion_parallel_runtime_overrides(
engine_args: dict[str, Any],
runtime_overrides: dict[str, Any],
Loaded 100 of 162 files, more files were not shown because too many files have changed in this diff. Show more