Compare commits

...
Author SHA1 Message Date
hsliuustc0106 bfdf9534e8 Update maintainers and code owners
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-07-01 21:40:03 +08:00
Zeng ChuangandHongsheng Liu 67bd02f767 [Feat] Enable streaming CoT display for HunyuanImage-3.0 AR generation (#4148)
Signed-off-by: zengchuang <zengchuang3@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-06-24 22:08:30 +08:00
Zhisen Zhang 0814d845dc [Refactor] Migrate VACE example to standard task examples + model_extras (#4648)
Signed-off-by: ForestWisdom <zszhang01@gmail.com>
2026-06-24 13:54:36 +00:00
rein yang bb973cb10c [Bugfix] [AURA] Correct aura results when use qwen3-tts in custom voice mode (#4650)
Signed-off-by: R2-Y <ruiruyang2@gmail.com>
2026-06-24 17:29:48 +08:00
Yuanheng Zhao fb3abb38fb [refactor] Flatten Ming-flash-omni-2.0 imagegen args (#4587)
Signed-off-by: Yuanheng <jonathan.zhaoyh@gmail.com>
2026-06-24 17:01:58 +08:00
Chendi.Xue eb81914c64 [S2V] Add sequence parallelism support to WanS2VTransformer3DModel (#4276)
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
2026-06-24 14:04:04 +08:00
1951dca786 [Model] add soulx-singer support (#3862)
Signed-off-by: Mr.Dong <1042542469@qq.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Mr.Dong <1042542469@qq.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
2026-06-24 05:36:45 +00:00
d71cacc74c [perf] avoid HunyuanImage3 timestep scalar sync (#4363)
Signed-off-by: TaffyOfficial <2324465096@qq.com>
Co-authored-by: TaffyOfficial <2324465096@qq.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-06-24 11:05:35 +08:00
Alex Brooks 465f1464d2 [Feat] Support Pipeline Resolution Based on HF Config & Make Pipeline Registry Eager (#3760)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-06-24 11:02:29 +08:00
amy-why-3459 f6894f4308 [BugFix] Delete unnecessary log printing (#4566)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-06-24 09:40:07 +08:00
rongfu.lengandGao Han a2d1dafc65 [Feature] add a new response_format is file type when image generate (#1673)
Signed-off-by: rongfu.leng <lenronfu@gmail.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-06-24 09:04:13 +08:00
MaciejBalaNVandHongsheng Liu 6dc4638a7e Fixed Cosmos3 regression (#4627)
Signed-off-by: Maciej Bala <mbala@nvidia.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-06-24 09:03:34 +08:00
SYLAR d35cdd2dab [Fix] JoyVL serving: align with reference engine (bounded long-term memory, timestamps, max_pixels) + recipe updates (#4623)
Signed-off-by: lishunyang12 <125541396+lishunyang12@users.noreply.github.com>
2026-06-24 07:04:17 +08:00
wangyu 0958e62ecc Add field IDs to bug report issue template (#4605)
Signed-off-by: wangyu <410167048@qq.com>
2026-06-23 20:26:25 +08:00
3901ad8165 [feat] Streaming diffusion video generation output (#3737)
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: Zeyu Huang | 黃澤宇 <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-23 09:57:53 +00:00
QianCyrusandHongsheng Liu b84c99c122 [Bugfix] Fix make_request_output TypeError on multimodal generation-stage outputs (#4579)
Signed-off-by: QianCyrus <happyloser233@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-06-23 14:42:30 +08:00
Guanzhe (Leo) Huangand汪志鹏 733f4f6a36 [Refactor] Migrate cosmos3 example to standard task example + model_extra (#4581)
Signed-off-by: leohuang257 <masugchds@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
2026-06-23 14:42:05 +08:00
Nick Cao 67d14b5ac4 [Refactor] Replace hardcoded CUDA device selection with platform-agnostic API (#4496)
Signed-off-by: Nick Cao <ncao@redhat.com>
2026-06-23 14:38:32 +08:00
Y. Fisher 64efa66869 [NPU] Initialize Ascend config for diffusion workers (#4386)
Signed-off-by: KexiongYu <yukexiong1@huawei.com>
2026-06-23 03:12:44 +00:00
Nick CaoandClaude 6646a66d41 [BugFix] Use ModuleDiscovery for HSDP DiT sharding to support nested pipelines (#3456)
Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-06-23 02:54:14 +00:00
CrazyCoder 378359fc11 [Doc] Stability AI/Stable-Audio-Open (#3664)
Signed-off-by: CrazyCoder <18235787078@163.com>
2026-06-23 10:03:02 +08:00
9c260b3774 Qwen3-TTS full-payload: emit empty-finished payload on degenerate take instead of dropping (fixes Stage-1 300s hang) (#4465)
Signed-off-by: jjw <m@jjw.in>
Signed-off-by: henryj <349043+henryj@users.noreply.github.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-23 00:21:11 +00:00
Yuanheng Zhao 9177b37e2f [NFC] Update misleading seed-tts-eval optional dependency guidance (#4619)
Signed-off-by: Yuanheng <jonathan.zhaoyh@gmail.com>
2026-06-22 16:35:42 -07:00
50bc51ffdf fix: return proper HTTP status codes for audio voice endpoints (#3969)
Signed-off-by: Ian Eaves <ian.k.eaves@gmail.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-06-22 22:57:07 +00:00
BeatSeat 044240d8eb [Model] Add IndexTTS2 text-to-speech support (#3838)
Signed-off-by: BeatSeat <wendavid552@gmail.com>
Signed-off-by: David <david@podcast.ai>
Signed-off-by: wendavid552 <wendavid552@gmail.com>
Signed-off-by: WenDavid <davidzhangcf@foxmail.com>
2026-06-22 19:24:01 +00:00
TaffyOfficialandTaffyOfficial f4d922d8f6 [Bugfix][HunyuanImage3] Align AR RGB conversion with official semantics (#4502)
Signed-off-by: TaffyOfficial <2324465096@qq.com>
Co-authored-by: TaffyOfficial <2324465096@qq.com>
2026-06-22 16:09:03 +00:00
Chenguang Zhengandamy-why-3459 eaa16e1941 [Feature] Support Async Omni output materialization (#4476)
Signed-off-by: Chenguang ZHENG <645327136@qq.com>
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-06-22 16:07:07 +00:00
Zhou TaichangandClaude Opus 4.8 8a6786731e [CI][Rebase] Fix Ready/Merge pipeline upload by stripping source_file_dependencies (#4532)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-22 15:22:11 +00:00
279 changed files with 27202 additions and 2207 deletions

No files matched your search

+3 -3
View File
@@ -14,7 +14,7 @@ steps:
depends_on: image-build
key: upload-ready-pipeline
commands:
- buildkite-agent pipeline upload .buildkite/test-ready.yml
- python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-ready.yml
agents:
queue: "cpu_queue_premerge"
@@ -23,7 +23,7 @@ steps:
depends_on: image-build
key: upload-merge-pipeline
commands:
- buildkite-agent pipeline upload .buildkite/test-merge.yml
- python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-merge.yml
agents:
queue: "cpu_queue_premerge"
@@ -33,6 +33,6 @@ steps:
key: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
- python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-nightly.yml
agents:
queue: "cpu_queue_premerge"
+17 -3
View File
@@ -13,13 +13,16 @@ Test pipeline mode (e.g. test-merge.yml):
recognize this uploader-only key).
Usage:
python3 upload_pipeline.py [--upload] <pipeline.yml>
python3 upload_pipeline.py [--upload] [--all] <pipeline.yml>
# Bootstrap (replaces upload_pipeline_with_skip_ci.sh):
python3 upload_pipeline.py --upload .buildkite/pipeline.yml
# Test pipeline (replaces upload_test_pipeline_with_diff_skip.py):
python3 upload_pipeline.py --upload .buildkite/test-merge.yml
# Full suite (rebase pipeline): keep all steps, still strip the uploader-only key:
python3 upload_pipeline.py --upload --all .buildkite/test-merge.yml
"""
from __future__ import annotations
@@ -437,11 +440,17 @@ def resolve_pipeline_path(arg: str) -> Path:
return ROOT / path
def render_pipeline(path: Path) -> str:
def render_pipeline(path: Path, *, force_all: bool = False) -> str:
text = path.read_text(encoding="utf-8")
diff_range = resolve_diff_range()
changed_files = resolve_changed_files()
# ``--all`` forces the keep-all-steps path (no diff-aware skipping) while still
# stripping ``source_file_dependencies``. Used by the rebase pipeline so main builds
# run the full e2e suite (see .buildkite/rebase-pipeline.yaml).
if force_all:
changed_files = None
if BOOTSTRAP_MARKER in text:
rendered = render_bootstrap_pipeline(
text,
@@ -480,6 +489,11 @@ def main() -> int:
action="store_true",
help="Pipe rendered YAML to buildkite-agent pipeline upload",
)
parser.add_argument(
"--all",
action="store_true",
help="Keep all steps (disable diff-aware skipping); still strips source_file_dependencies",
)
args = parser.parse_args()
path = resolve_pipeline_path(args.pipeline)
@@ -487,7 +501,7 @@ def main() -> int:
_log(f"missing pipeline file: {path}")
return 1
rendered = render_pipeline(path)
rendered = render_pipeline(path, force_all=args.all)
if args.upload:
upload_to_buildkite(rendered)
else:
+64
View File
@@ -363,6 +363,42 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "TTS · IndexTTS2 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_indextts2.py
- tests/e2e/offline_inference/test_indextts2.py
- tests/model_executor/models/indextts2/
- tests/model_executor/stage_input_processors/test_indextts2.py
- tests/assets/indextts2/
- vllm_omni/model_executor/models/indextts2/
- vllm_omni/model_executor/stage_input_processors/indextts2.py
- vllm_omni/deploy/indextts2.yaml
- vllm_omni/deploy/indextts2_low_latency.yaml
- vllm_omni/entrypoints/openai/serving_speech.py
- vllm_omni/entrypoints/openai/tts_adapters/indextts2.py
- vllm_omni/config/pipeline_registry.py
- vllm_omni/model_executor/models/registry.py
timeout_in_minutes: 20
commands:
- |
timeout 20m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
pytest -s -v tests/e2e/online_serving/test_indextts2.py tests/e2e/offline_inference/test_indextts2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
"
agents:
queue: "gpu_4_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
gpus: all
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni · Qwen3-Omni Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_qwen3_omni.py
@@ -509,6 +545,34 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "TTS · SoulX-Singer Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_soulxsinger.py
- vllm_omni/diffusion/models/soulx_singer/
- vllm_omni/deploy/soulxsinger_svc.yaml
- vllm_omni/deploy/soulxsinger_svs.yaml
timeout_in_minutes: 30
commands:
- |
timeout 30m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/offline_inference/test_soulxsinger.py -k 'svc or precomputed' -m 'advanced_model and cuda' --run-level 'advanced_model'
"
agents:
queue: "gpu_1_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion · Qwen Image Edit Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image_edit.py
+2 -1
View File
@@ -1279,7 +1279,7 @@ steps:
- label: ":full_moon: Diffusion X2V · Function Test"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py 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_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py tests/e2e/online_serving/test_ltx2_expansion.py -m "full_model and cuda" --run-level "full_model"
agents:
queue: "mithril-h100-pool"
plugins:
@@ -1322,6 +1322,7 @@ steps:
- 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/test_ltx2_3_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -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
agents:
queue: "mithril-h100-pool"
plugins:
+10 -10
View File
@@ -22,30 +22,30 @@
/.buildkite/ @yenuo26 @congw729
# Engine and model_executor
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan
/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks
/vllm_omni/model_executor/ @tzhouam @gcanlin
/vllm_omni/worker/ @tzhouam @gcanlin
# Diffusion
/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666
/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
/vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
# Diffusion CPU offloading
/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao @xuechendi
# Omni and TTS support
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin
# Benchmarks
/benchmarks/accuracy/ @david6666666
/benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa
/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi
/benchmarks/distributed/ @princepride
/benchmarks/glm_image/ @princepride @JaredforReal
/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride
/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin
# Usability and Observability
/vllm_omni/deploy/ @tzhouam @lishunyang12
@@ -53,13 +53,13 @@
/vllm_omni/profiler/ @gcanlin
# ComfyUI
/apps/ @fhfuih
/apps/ @fhfuih @alex-jw-brooks
# Large Scale Serving
/vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
# Quantization and configuration
/vllm_omni/config/ @lishunyang12
/vllm_omni/config/ @lishunyang12 @alex-jw-brooks
/vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
# Hardware plugins and accelerators
@@ -19,6 +19,7 @@ body:
Consider redacting or replacing sensitive values with placeholders like `<YOUR_TOKEN_HERE>` when sharing configuration or code examples.
- type: textarea
id: current-environment
attributes:
label: Your current environment
description: |
@@ -42,6 +43,7 @@ body:
required: true
- type: textarea
id: code-version
attributes:
label: Your code version
description: |
@@ -66,6 +68,7 @@ body:
required: true
- type: textarea
id: bug-description
attributes:
label: 🐛 Describe the bug
description: |
+1
View File
@@ -241,6 +241,7 @@ datasets/
*.json
!tests/dfx/perf/tests/*.json
!apps/ComfyUI-vLLM-Omni/example_workflows/*.json
!tests/assets/soulxsinger/*.json
*.jsonl
*.parquet
+1 -1
View File
@@ -101,7 +101,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
ENV LD_LIBRARY_PATH="/opt/intel/oneapi/ccl/2021.15/lib:$LD_LIBRARY_PATH:/usr/local/lib/"
ENV VLLM_TARGET_DEVICE=xpu
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+7 -2
View File
@@ -33,16 +33,20 @@ Lead maintainers are responsible for the overall direction and strategy of the p
Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
- [@alex-jw-brooks](https://github.com/alex-jw-brooks): Frontend, Configuration, and Entrypoint Support
- [@amy-why-3459](https://github.com/amy-why-3459): Omni Modality Model and World Model Support
- [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
- [@gcanlin](https://github.com/gcanlin): Hardware plugin, NPU integration, and TTS Model Support
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
- [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs
- [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache
- [@Sy0307](https://github.com/Sy0307): TTS Model and AR Engine Support
- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
- [@wtomin](https://github.com/wtomin): Diffusion models, parallel, and docs
- [@xuechendi](https://github.com/xuechendi): Diffusion and Intel GPU Support
- [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
- [@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
@@ -53,7 +57,7 @@ This section breaks down diffusion-related responsibilities. If you have PRs tou
### Runtime-related
- Diffusion models: @RuixiangMa, @wtomin
- Diffusion models: @RuixiangMa, @wtomin, @xuechendi
### Optimization
@@ -61,6 +65,7 @@ This section breaks down diffusion-related responsibilities. If you have PRs tou
- Attention backend: @ZJY0516
- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
- Offload: @yuanheng-zhao
- Intel GPU support: @xuechendi
- Quantization: @lishunyang12, @david6666666
- Kernel fusion / communication-computation: @ZJY0516
+4 -1
View File
@@ -37,7 +37,7 @@ th {
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
@@ -63,6 +63,8 @@ th {
| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ |
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | |
| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | |
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
@@ -74,6 +76,7 @@ th {
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | |
| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | |
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
+44
View File
@@ -32,6 +32,20 @@ curl -X POST http://localhost:8000/v1/images/generations \
}' | jq -r '.data[0].b64_json' | base64 -d > dragon.png
```
**Using curl save to file:**
```bash
curl -o dragon.png -X POST http://localhost:8000/v1/images/generations \
-H "Content-Type: application/json" \
-d '{
"prompt": "a dragon laying over the spine of the Green Mountains of Vermont",
"size": "1024x1024",
"seed": 42,
"response_format":"file"
}'
```
**Using Python:**
```python
@@ -57,6 +71,36 @@ img = Image.open(io.BytesIO(img_bytes))
img.save("cat.png")
```
**Using Python save to file:**
```python
import requests
import base64
from PIL import Image
import io
import re
response = requests.post(
"http://localhost:8000/v1/images/generations",
json={
"prompt": "a black and white cat wearing a princess tiara",
"size": "1024x1024",
"num_inference_steps": 50,
"seed": 42,
"response_format":"file"
}
)
# save to file
content_disposition = response.headers.get("Content-Disposition", "")
match = re.search(r'filename="?(.+)"?', content_disposition)
filename = match.group(1) if match else "save.png"
with open(filename, "wb") as f:
for chunk in response.iter_content(8192):
f.write(chunk)
print("saved:", filename)
```
**Using OpenAI SDK:**
```python
@@ -197,6 +197,7 @@ Factory function `get_offload_backend()` selects the appropriate backend based o
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` |
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | `"layers"` |
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | `"layers"`, `"customized modules"` |
**Notes:**
+1 -1
View File
@@ -146,7 +146,7 @@ The following tables show which models support each feature:
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
| **Wan2.2-S2V** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
| **Wan2.2-S2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
@@ -1,14 +0,0 @@
{
"prompt": "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. The robotic arm is in the process of pouring a transparent liquid into the glass. To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. The lighting is bright, casting soft shadows on the table. The robotic arm's movements are smooth and controlled, demonstrating precision in its task. As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. The other robotic arm on the right side also stays stationary throughout the video. The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper.",
"negative_prompt": "",
"vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4",
"height": 720,
"width": 1280,
"num_frames": 189,
"num_inference_steps": 35,
"guidance_scale": 6.0,
"flow_shift": 10.0,
"fps": 24,
"condition_frame_indexes_vision": [0, 1],
"condition_video_keep": "first"
}
@@ -74,6 +74,12 @@ 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",
default=False,
help="Enable streaming CoT display; print AR text token-by-token in real time.",
)
parser.add_argument("--log-stats", action="store_true", default=False)
parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.")
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
@@ -222,6 +228,8 @@ def main():
prompt_dict["modalities"] = ["text"]
formatted_prompts.append(prompt_dict)
from vllm.sampling_params import RequestOutputKind
params_list = list(omni.default_sampling_params_list)
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -253,6 +261,10 @@ def main():
sp.width = args.width
elif hasattr(sp, "stop_token_ids"):
sp.stop_token_ids = ar_stop_token_ids
# When --stream is set, request DELTA output from the AR stage
# so we can display CoT text token-by-token in real time.
if args.stream and hasattr(sp, "output_kind"):
sp.output_kind = RequestOutputKind.DELTA
print(f"\n{'=' * 60}")
print("HunyuanImage-3.0 Generation Configuration:")
@@ -281,22 +293,32 @@ def main():
print(f" Prompts: {prompts}")
print(f"{'=' * 60}\n")
omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
# When --stream is set, print AR CoT text token-by-token in real time.
# Otherwise, collect and print the full AR text once when stage 0 finishes.
omni_outputs = omni.generate(
prompts=formatted_prompts,
sampling_params_list=params_list,
py_generator=True,
use_tqdm=False,
)
img_idx = 0
for req_output in omni_outputs:
ro = getattr(req_output, "request_output", None)
txt = ""
if ro and getattr(ro, "outputs", None):
txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
if not txt:
ar_text = getattr(req_output, "custom_output", {}).get("ar_generated_text")
if isinstance(ar_text, list):
txt = "\n".join(text for text in ar_text if text)
else:
txt = ar_text or ""
if txt:
print(f"[Output] Text:\n{txt}")
stage_id = getattr(req_output, "stage_id", None)
# AR stage text — each CompletionOutput.text is already a delta when
# output_kind=DELTA, so we can print it directly (matching the pattern
# in serving_chat.py).
if stage_id == 0 and ro and getattr(ro, "outputs", None):
for o in ro.outputs:
text = getattr(o, "text", "") or ""
if text:
print(text, end="", flush=True)
# Non-streaming: one shot with full text — emit a trailing newline.
if not args.stream:
print(flush=True)
# Collect images from diffusion stage
images = getattr(req_output, "images", None)
if not images and ro and hasattr(ro, "images"):
images = ro.images
@@ -304,7 +326,7 @@ def main():
for j, img in enumerate(images):
save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
img.save(save_path)
print(f"[Output] Saved image to {save_path}")
print(f"\n[Output] Saved image to {save_path}")
img_idx += 1
@@ -1,6 +1,7 @@
# Image-To-Video
This example demonstrates how to generate videos from images using Wan2.2 Image-to-Video models with vLLM-Omni's offline inference API.
This shared example generates videos from images with VACE, Wan2.2, LTX2,
HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
## Local CLI Usage
@@ -48,19 +49,38 @@ python image_to_video.py \
--output i2v_output.mp4
```
### Cosmos3
```bash
# Cosmos3 bundles example frames under assets/ (any RGB image works too):
python image_to_video.py \
--model nvidia/Cosmos3-Nano \
--image /path/to/Cosmos3-Nano/assets/example_i2v_input.jpg \
--prompt "The scene comes to life with smooth, natural motion." \
--negative-prompt "blurry, distorted, low quality" \
--height 720 --width 1280 --num-frames 189 --fps 24 \
--num-inference-steps 35 --guidance-scale 6.0 \
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}' \
--output cosmos3_i2v.mp4
```
Key arguments:
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V).
- `--image`: Path to input image (required).
- `--image`: Path to the first-frame or source image.
- `--last-image`: Optional last-frame condition for models such as VACE.
- `--mask-image`: Optional inpainting mask. White pixels are regenerated and black pixels are preserved.
- `--reference-image`: Optional reference image; repeat it to provide multiple references.
- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3.
- `--prompt`: Text description of desired motion/animation.
- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16.
- `--num-frames`: Number of frames (default 81).
- `--num-frames`: Number of frames (default: model-specific — Wan 81, LTX2 121, Cosmos3 189).
- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
- `--negative-prompt`: Optional list of artifacts to suppress.
- `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
- `--flow-shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
- `--flow-shift`: Scheduler flow shift (default: model-specific — Wan/LTX2 5.0, Cosmos3 10.0).
- `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints.
- `--num-inference-steps`: Number of denoising steps (default 50).
- `--num-inference-steps`: Number of denoising steps (default: model-specific — Wan 50, LTX2 40, Cosmos3 35).
- `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
- `--output`: Path to save the generated video.
- `--vae-use-slicing`: Enable VAE slicing for memory optimization.
@@ -76,5 +96,98 @@ Key arguments:
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
## Wan2.1 VACE Conditional Tasks
The shared script selects the VACE conditioning structure from the media inputs.
No explicit mode parameter is required: the script constructs the source video,
mask, or reference images consumed by the VACE pipeline.
Download the same Hugging Face assets used by the original VACE example:
```bash
wget -O astronaut.jpg https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg
wget -O vace_first_frame.png https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_first_frame.png
wget -O vace_last_frame.png https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_last_frame.png
```
### Image-to-Video (I2V)
```bash
python image_to_video.py \
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
--image astronaut.jpg \
--prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \
--seed 42 --height 480 --width 832 --num-frames 81 \
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
--vae-use-tiling --output vace_i2v_output.mp4
```
### Video-to-Last-Frame (V2LF)
```bash
python image_to_video.py \
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
--last-image astronaut.jpg \
--prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \
--seed 42 --height 480 --width 832 --num-frames 81 \
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
--vae-use-tiling --output vace_v2lf_output.mp4
```
### First-Last-Frame-to-Video (FLF2V)
```bash
python image_to_video.py \
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
--image vace_first_frame.png \
--last-image vace_last_frame.png \
--prompt "CG animation style, a small blue bird takes off from a branch and lands on another branch" \
--seed 42 --height 512 --width 512 --num-frames 81 \
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
--vae-use-tiling --output vace_flf2v_output.mp4
```
### Inpainting
Create a mask matching the original VACE example: a 160-pixel-wide white
vertical stripe marks the region to regenerate.
```bash
python - <<'PY'
from PIL import Image
mask = Image.new("L", (832, 480), 0)
mask.paste(255, (336, 0, 496, 480))
mask.save("vace_center_mask.png")
PY
```
```bash
python image_to_video.py \
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
--image astronaut.jpg \
--mask-image vace_center_mask.png \
--prompt "Shrek, the ogre, walks out of a building in a happy mood" \
--seed 42 --height 480 --width 832 --num-frames 81 \
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
--vae-use-tiling --output vace_inpaint_output.mp4
```
### Reference-to-Video (R2V)
Repeat `--reference-image` to provide more than one reference image.
```bash
python image_to_video.py \
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
--reference-image astronaut.jpg \
--prompt "Camera slowly zooms out from the character walking in a garden" \
--seed 42 --height 480 --width 832 --num-frames 81 \
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
--vae-use-tiling --output vace_r2v_output.mp4
```
The VACE T2V command is documented in the shared
[`text_to_video.py`](../text_to_video/text_to_video.md#wan21-vace-t2v) example.
For Wan2.2 LightX2V-converted local Diffusers directories and related LoRA
assets, see the [LoRA guide](../../../docs/user_guide/diffusion/lora.md#wan22-lightx2v-offline-assembly).
@@ -2,13 +2,14 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, or HunyuanVideo-1.5.
Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, HunyuanVideo-1.5, or Wan2.1 VACE.
Supports:
- Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
- Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B)
- LTX2 image-to-video pipeline
- HunyuanVideo-1.5 I2V: SigLIP + VAE dual image conditioning
- Wan2.1 VACE: first/last-frame, inpainting, and reference conditioning
Usage:
# Wan I2V-A14B (MoE)
@@ -30,9 +31,16 @@ Usage:
python image_to_video.py --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \
--image input.jpg --prompt "A cat playing with yarn" \
--flow-shift 5.0 --guidance-scale 6.0
# Cosmos3 I2V (image conditioning)
python image_to_video.py --model nvidia/Cosmos3-Nano \
--image input.jpg --prompt "The scene comes to life with smooth, natural motion." \
--num-frames 189 --num-inference-steps 35 --guidance-scale 6.0 --fps 24 \
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}'
"""
import argparse
import functools
import json
import time
from pathlib import Path
@@ -43,39 +51,59 @@ import PIL.Image
import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_extras import (
build_image_to_video_prompt,
get_extra_body_params,
get_model_class_name,
)
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
def parse_profiler_config(value: str) -> dict[str, Any]:
def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]:
"""Parse a CLI value as a JSON object, attributing errors to ``flag_name``."""
try:
config = json.loads(value)
except json.JSONDecodeError as e:
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e
if not isinstance(config, dict):
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object")
return config
parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate a video from an image (Wan2.2, LTX2, HunyuanVideo-1.5).")
parser = argparse.ArgumentParser(
description="Generate a video from one or more images (Wan2.2, LTX2, HunyuanVideo-1.5, or Wan2.1 VACE)."
)
parser.add_argument(
"--model",
default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
help="Diffusers I2V model ID or local path (Wan2.2 or HunyuanVideo-1.5).",
help="Diffusers I2V model ID or local path (Wan2.2, HunyuanVideo-1.5, or Wan2.1 VACE).",
)
parser.add_argument(
"--model-class-name",
default=None,
help="Override model class name (e.g., LTX2ImageToVideoPipeline).",
)
parser.add_argument("--image", required=True, help="Path to input image.")
parser.add_argument("--image", help="Path to the first-frame or source image.")
parser.add_argument("--last-image", help="Path to a last-frame condition (used by models such as VACE).")
parser.add_argument("--mask-image", help="Path to an inpainting mask (used by models such as VACE).")
parser.add_argument(
"--reference-image",
action="append",
default=None,
help="Path to a reference image. Repeat to provide multiple references.",
)
parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.")
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.")
parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
parser.add_argument(
"--guidance-scale-high", type=float, default=None, help="Optional separate CFG for high-noise (MoE only)."
)
@@ -83,8 +111,10 @@ def parse_args() -> argparse.Namespace:
"--height", type=int, default=None, help="Video height (auto-calculated from image if not set)."
)
parser.add_argument("--width", type=int, default=None, help="Video width (auto-calculated from image if not set).")
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps.")
parser.add_argument("--num-frames", type=int, default=None, help="Number of frames. Default: model-specific.")
parser.add_argument(
"--num-inference-steps", type=int, default=None, help="Sampling steps. Default: model-specific."
)
parser.add_argument("--boundary-ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.")
parser.add_argument(
"--frame-rate",
@@ -93,7 +123,10 @@ def parse_args() -> argparse.Namespace:
help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
)
parser.add_argument(
"--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
"--flow-shift",
type=float,
default=None,
help="Scheduler flow_shift. Default: model-specific (Wan 5.0, Cosmos3 10.0).",
)
parser.add_argument(
"--sample-solver",
@@ -244,6 +277,19 @@ def parse_args() -> argparse.Namespace:
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
parser.add_argument(
"--extra-body",
type=functools.partial(parse_json_object, flag_name="--extra-body"),
default=None,
help=(
"Model-specific generation params as a JSON object. Keys are filtered "
"against the model's declared extra_body_params (see vllm_omni/model_extras), "
"so unknown keys for the chosen model are silently dropped. "
'Cosmos3 V2V example: \'{"condition_frame_indexes_vision": [0, 1], '
'"condition_video_keep": "first", "flow_shift": 10.0, '
'"max_sequence_length": 4096, "guardrails": false}\'.'
),
)
return parser.parse_args()
@@ -267,31 +313,63 @@ def main():
model_name = str(args.model).lower() if args.model is not None else ""
model_class_name = args.model_class_name
is_ltx2 = "ltx2" in model_name or (model_class_name and "ltx2" in model_class_name.lower())
is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower())
if model_class_name is None and is_ltx2:
model_class_name = "LTX2ImageToVideoPipeline"
# Load input image
image = PIL.Image.open(args.image).convert("RGB")
image = PIL.Image.open(args.image).convert("RGB") if args.image else None
last_image = PIL.Image.open(args.last_image).convert("RGB") if args.last_image else None
mask_image = PIL.Image.open(args.mask_image).convert("L") if args.mask_image else None
reference_images = (
[PIL.Image.open(path).convert("RGB") for path in args.reference_image] if args.reference_image else None
)
dimension_image = image or last_image or (reference_images[0] if reference_images else None)
if dimension_image is None:
raise ValueError("Provide --image, --last-image, or at least one --reference-image.")
fps = args.fps if args.fps is not None else (24 if is_ltx2 else 16)
# Per-model generation defaults, applied only when the matching flag is omitted.
# Cosmos3 would otherwise silently inherit the Wan2.2 defaults (wrong size/steps/shift).
if is_cosmos:
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
24,
6.0,
189,
35,
10.0,
1280 * 720,
16,
)
elif is_ltx2:
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 24, 4.0, 121, 40, 5.0, 512 * 768, 32
else: # Wan2.2 / HunyuanVideo-1.5
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 16, 5.0, 81, 50, 5.0, 480 * 832, 16
fps = args.fps if args.fps is not None else d_fps
frame_rate = args.frame_rate if args.frame_rate is not None else float(fps)
guidance_scale = args.guidance_scale if args.guidance_scale is not None else (4.0 if is_ltx2 else 5.0)
num_frames = args.num_frames if args.num_frames is not None else (121 if is_ltx2 else 81)
num_inference_steps = args.num_inference_steps if args.num_inference_steps is not None else (40 if is_ltx2 else 50)
guidance_scale = args.guidance_scale if args.guidance_scale is not None else d_guidance
num_frames = args.num_frames if args.num_frames is not None else d_num_frames
num_inference_steps = args.num_inference_steps if args.num_inference_steps is not None else d_steps
flow_shift = args.flow_shift if args.flow_shift is not None else d_flow_shift
# Calculate dimensions if not provided
# Calculate dimensions if not provided (model-aware max area).
height = args.height
width = args.width
if height is None or width is None:
# Default to 480P area for Wan2.2 I2V, 512x768 area for LTX2
max_area = 512 * 768 if is_ltx2 else 480 * 832
mod_value = 32 if is_ltx2 else 16
calc_height, calc_width = calculate_dimensions(image, max_area=max_area, mod_value=mod_value)
calc_height, calc_width = calculate_dimensions(dimension_image, max_area=d_max_area, mod_value=d_mod)
height = height or calc_height
width = width or calc_width
# Resize image to target dimensions
image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
media_inputs: dict[str, Any] = {}
if image is not None:
media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
if last_image is not None:
media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS)
if mask_image is not None:
media_inputs["mask"] = mask_image.resize((width, height), PIL.Image.Resampling.NEAREST)
if reference_images is not None:
media_inputs["reference_images"] = [
reference.resize((width, height), PIL.Image.Resampling.LANCZOS) for reference in reference_images
]
# Configure cache based on backend type
cache_config = None
@@ -342,7 +420,7 @@ def main():
vae_use_slicing=args.vae_use_slicing,
vae_use_tiling=args.vae_use_tiling,
boundary_ratio=args.boundary_ratio,
flow_shift=args.flow_shift,
flow_shift=flow_shift,
diffusion_kv_cache_dtype=args.diffusion_kv_cache_dtype,
diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps,
diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers,
@@ -357,7 +435,13 @@ def main():
)
if args.quantization is not None:
omni_kwargs["quantization"] = args.quantization
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
# gate is an engine-level config (offline analog of the server's --no-guardrails).
if args.extra_body and "guardrails" in args.extra_body:
omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
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)
if profiler_enabled:
print("[Profiler] Starting profiling...")
@@ -367,8 +451,8 @@ def main():
print(f"\n{'=' * 60}")
print("Generation Configuration:")
print(f" Model: {args.model}")
print(f" Inference steps: {args.num_inference_steps}")
print(f" Frames: {args.num_frames}")
print(f" Inference steps: {num_inference_steps}")
print(f" Frames: {num_frames}")
print(f" Solver: {args.sample_solver}")
print(f" diffusion_kv_cache_dtype(config): {args.diffusion_kv_cache_dtype}")
print(f" diffusion_kv_cache_skip_steps(config): {args.diffusion_kv_cache_skip_steps}")
@@ -381,30 +465,43 @@ def main():
print(f" Video size: {args.width}x{args.height}")
print(f"{'=' * 60}\n")
prompt_dict = build_image_to_video_prompt(
model_class_name=model_class_name,
prompt=args.prompt,
negative_prompt=args.negative_prompt,
media_inputs=media_inputs,
height=height,
width=width,
num_frames=num_frames,
)
sampling_params = OmniDiffusionSamplingParams(
height=height,
width=width,
generator=generator,
guidance_scale=guidance_scale,
guidance_scale_2=args.guidance_scale_high,
boundary_ratio=args.boundary_ratio,
num_inference_steps=num_inference_steps,
num_frames=num_frames,
frame_rate=frame_rate,
extra_args={
"sample_solver": args.sample_solver,
"flow_shift": flow_shift,
},
)
# Route model-specific knobs through extra_body, filtered against the model's
# declared extra_body_params. Models without a declaration only forward explicit
# --extra-body JSON (e.g. Cosmos3 V2V's condition_frame_indexes_vision).
extra_body = dict(args.extra_body or {})
if declared_extra_body_params:
apply_declared_extra_args(sampling_params, declared_extra_body_params, extra_body)
elif extra_body:
sampling_params.extra_args.update({k: v for k, v in extra_body.items() if v is not None})
generation_start = time.perf_counter()
# omni.generate() returns Generator[OmniRequestOutput, None, None]
frames = omni.generate(
{
"prompt": args.prompt,
"negative_prompt": args.negative_prompt,
"multi_modal_data": {"image": image},
},
OmniDiffusionSamplingParams(
height=height,
width=width,
generator=generator,
guidance_scale=guidance_scale,
guidance_scale_2=args.guidance_scale_high,
boundary_ratio=args.boundary_ratio,
num_inference_steps=num_inference_steps,
num_frames=num_frames,
frame_rate=frame_rate,
extra_args={
"sample_solver": args.sample_solver,
"flow_shift": args.flow_shift,
},
),
)
frames = omni.generate(prompt_dict, sampling_params)
generation_end = time.perf_counter()
generation_time = generation_end - generation_start
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import functools
import json
import time
from pathlib import Path
@@ -38,16 +39,20 @@ def is_nextstep_model(model_name: str) -> bool:
return False
def parse_profiler_config(value: str) -> dict[str, Any]:
def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]:
"""Parse a CLI value as a JSON object, attributing errors to ``flag_name``."""
try:
config = json.loads(value)
except json.JSONDecodeError as e:
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e
if not isinstance(config, dict):
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object")
return config
parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
parser.add_argument(
@@ -269,7 +274,7 @@ def parse_args() -> argparse.Namespace:
)
parser.add_argument(
"--extra-body",
type=parse_profiler_config,
type=functools.partial(parse_json_object, flag_name="--extra-body"),
default=None,
help=(
"Model-specific generation params as a JSON object, e.g. "
@@ -425,6 +430,10 @@ def main():
if use_nextstep:
# NextStep-1.1 requires explicit pipeline class
omni_kwargs["model_class_name"] = "NextStep11Pipeline"
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
# gate is an engine-level config (offline analog of the server's --no-guardrails).
if args.extra_body and "guardrails" in args.extra_body:
omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
omni = Omni(**omni_kwargs)
model_class_name = get_model_class_name(omni)
declared_extra_body_params = get_extra_body_params(model_class_name)
@@ -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 |
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
## Common Quick Start
@@ -436,6 +437,50 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [
---
## IndexTTS-2
2-stage TTS pipeline (GPT AR talker + S2Mel CFM DiT + BigVGAN vocoder) at 22.05 kHz. Every request requires reference audio for zero-shot voice cloning. Supports emotion conditioning via audio, text, or 8-dim vector.
### Quick start
```bash
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
--model IndexTeam/IndexTTS-2 \
--text "你好,这是一个语音合成测试。" \
--ref-audio /path/to/reference.wav
```
### Emotion control
```bash
# Emotion from reference audio
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
--model IndexTeam/IndexTTS-2 \
--text "今天天气真好!" \
--ref-audio /path/to/ref.wav \
--emo-audio /path/to/happy.wav
# Emotion from 8-dim vector (happy angry sad afraid disgusted melancholy surprised calm)
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
--model IndexTeam/IndexTTS-2 \
--text "今天天气真好!" \
--ref-audio /path/to/ref.wav \
--emo-vector 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
# Emotion from text description
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
--model IndexTeam/IndexTTS-2 \
--text "今天天气真好!" \
--ref-audio /path/to/ref.wav \
--emo-text "happy and excited"
```
### Notes
- `--ref-audio` is **required** — IndexTTS-2 does not support text-only synthesis.
- Stage 0 (AR Talker): GPT-2 generates mel codes from text + reference audio.
- Stage 1 (S2Mel + BigVGAN): CFM DiT converts mel codes to waveform at 22.05 kHz.
- Deploy config: `vllm_omni/deploy/indextts2.yaml`. Stage 1 runs with `enforce_eager: true` (DiT has dynamic shapes).
---
## Voxtral TTS
Voxtral-4B-TTS (Mistral). Has its own argparse surface; uses voice presets and the `mistral_common` `SpeechRequest` protocol.
@@ -476,3 +521,100 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
- `--num-prompts N` replicates the prompt for performance measurement.
- `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
- Run `--help` for the full argument surface.
---
## SoulX-Singer
Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Script: `soulxsinger/end2end.py`. Deploy: `vllm_omni/deploy/soulxsinger_svs.yaml` or `soulxsinger_svc.yaml`.
### Prerequisites
Download DiT and preprocess weights, then set up separate SVS / SVC view directories. Copy `soulxsinger/utils/phoneme/phone_set.json` from upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) into the model weights dir as `phoneme/phone_set.json` — HuggingFace does not ship it.
```bash
# 1. DiT weights
export BASE=path/to/SoulX-Singer
export PREPROCESS=path/to/SoulX-Singer-Preprocess
export SVC_DIR=path/to/SoulX-Singer-svc
huggingface-cli download Soul-AILab/SoulX-Singer --local-dir "$BASE"
# 2. Preprocess weights (required)
huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess --local-dir "$PREPROCESS"
export SOULX_PREPROCESS_WEIGHTS_DIR="$PREPROCESS"
# 3. SVS / SVC view directories
mkdir -p "$SVC_DIR"
cp $BASE/{config.yaml,README.md,assets} $SVC_DIR
mv $BASE/model-svc.pt $SVC_DIR/model-svc.pt
cat > "$BASE/config.json" <<'EOF'
{
"model_type": "soulxsinger",
"architectures": ["SoulXSingerPipeline"],
"max_num_seqs": 1
}
EOF
cat > "$SVC_DIR/config.json" <<'EOF'
{
"model_type": "soulxsinger",
"architectures": ["SoulXSingerSVCPipeline"],
"max_num_seqs": 1
}
EOF
```
`config.yaml` hyper-parameters live under `$BASE`; each view's `config.json` `architectures` field is the single source of truth for SVS vs SVC. Point `--model` at the matching directory (`$BASE` for SVS, `$SVC_DIR` for SVC). Deploy YAML is chosen automatically from `config.json`; optional `--svs` / `--svc` only assert the mode matches.
**Online preprocess** is the default: pass `--prompt-audio` and `--target-audio`, and the worker runs vocal separation, F0, and (for SVS) lyrics/MIDI before DiT. Install only what your run needs:
```bash
pip install "BS-RoFormer" # vocal sep + F0 on GPU — SVS and SVC
```
Mandarin SVS also needs FunASR and Chinese G2P; `ffmpeg` must be on `PATH`:
```bash
# install optional dependencies:
pip install -e ".[soulx-svs]"
```
English SVS adds NeMo ASR and NLTK data; pass `--language English`:
```bash
pip install "nemo_toolkit[asr]==2.6.1" lhotse==1.32.2
python -c "import nltk; nltk.download('cmudict'); nltk.download('averaged_perceptron_tagger_eng')"
```
**Precomputed metadata** is the alternative: pass both `--prompt-metadata-path` and `--target-metadata-path` and skip online ASR/ROSVOT — none of the packages above are required. JSON can be produced by integrated preprocess on a prior run, or by upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) `preprocess/` scripts if you prefer to run that outside vLLM-Omni.
### Quick start
```bash
# SVS — default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3
python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \
--model "$BASE" \
--preprocess-weights-dir "$PREPROCESS" \
--control score \
--num-inference-steps 32 \
-o output.wav
python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \
--model "$SVC_DIR" \
--preprocess-weights-dir "$PREPROCESS" \
--svc \
--num-inference-steps 32 \
-o output_svc.wav
```
`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. Long SVS targets are handled in one request. See `end2end.py --help` for `--pitch-shift`, `--vocal-sep`, `--auto-shift`, and language/control options.
### Notes
- Output: 24 kHz mono WAV; batch only.
- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on.
- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only.
---
@@ -0,0 +1,187 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline inference example for IndexTTS2 via vLLM-Omni.
Two-stage pipeline: GPT AR (Stage 0) → S2Mel + BigVGAN (Stage 1).
Output is 22050 Hz mono WAV.
Usage:
python end2end.py \
--model /path/to/IndexTeam/IndexTTS-2 \
--text "你好,这是一个语音合成测试。" \
--ref-audio /path/to/ref.wav
# With emotion audio:
python end2end.py \
--model /path/to/IndexTeam/IndexTTS-2 \
--text "今天天气真好!" \
--ref-audio /path/to/ref.wav \
--emo-audio /path/to/happy.wav
"""
from __future__ import annotations
import os
from pathlib import Path
import soundfile as sf
import torch
from vllm import SamplingParams
from vllm.utils.argparse_utils import FlexibleArgumentParser
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
from vllm_omni import Omni # noqa: E402
from vllm_omni.model_executor.models.indextts2.prompt_utils import ( # noqa: E402
build_indextts2_prefill_prompt_ids,
)
def build_request(
model: str,
text: str,
ref_audio_path: str | None = None,
emo_audio_path: str | None = None,
emo_text: str | None = None,
emo_vector: list[float] | None = None,
emo_alpha: float | None = None,
use_emo_text: bool = False,
use_random: bool = False,
) -> dict:
additional: dict = {"text": [text]}
if ref_audio_path:
additional["voice"] = [str(ref_audio_path)]
if emo_audio_path:
additional["emo_audio"] = [str(emo_audio_path)]
if emo_text:
additional["emo_text"] = [emo_text]
if emo_vector is not None:
additional["emo_vector"] = [emo_vector]
if emo_alpha is not None:
additional["emo_alpha"] = [emo_alpha]
if use_emo_text:
additional["use_emo_text"] = [True]
if use_random:
additional["use_random"] = [True]
return {
"prompt_token_ids": build_indextts2_prefill_prompt_ids(model, text),
"additional_information": additional,
}
def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = 22050) -> None:
audio_np = waveform.float().numpy()
sf.write(path, audio_np, sample_rate)
print(f" Saved {path} ({audio_np.shape}, {sample_rate} Hz)")
def extract_audio(mm: dict) -> tuple[torch.Tensor | None, int]:
audio = mm.get("audio")
if audio is None:
audio = mm.get("model_outputs")
if isinstance(audio, list):
chunks = [chunk.reshape(-1) for chunk in audio if isinstance(chunk, torch.Tensor) and chunk.numel() > 0]
audio = torch.cat(chunks, dim=0) if chunks else None
sr_val = mm.get("sr")
if isinstance(sr_val, list):
sr_val = sr_val[-1] if sr_val else None
if hasattr(sr_val, "item"):
sample_rate = int(sr_val.item())
else:
sample_rate = int(sr_val) if sr_val is not None else 22050
return audio if isinstance(audio, torch.Tensor) else None, sample_rate
def main(args) -> None:
omni = Omni(
model=args.model,
deploy_config=args.deploy_config,
stage_init_timeout=args.stage_init_timeout,
)
# Stage 0: GPT AR (text → mel codes)
gpt_sampling = SamplingParams(
temperature=0.8,
top_p=0.8,
top_k=30,
max_tokens=1500,
repetition_penalty=10.0,
stop_token_ids=[8193],
seed=args.seed if args.seed is not None else 42,
detokenize=False,
)
# Stage 1: S2Mel + BigVGAN (non-AR, params mostly ignored)
s2mel_sampling = SamplingParams(
temperature=0.0,
max_tokens=65536,
detokenize=True,
)
sampling_params = [gpt_sampling, s2mel_sampling]
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
print(f"Synthesizing: {args.text!r}")
if args.ref_audio:
print(f" ref_audio: {args.ref_audio}")
inputs = build_request(
model=args.model,
text=args.text,
ref_audio_path=args.ref_audio,
emo_audio_path=args.emo_audio,
emo_text=args.emo_text,
emo_vector=args.emo_vector,
emo_alpha=args.emo_alpha,
use_emo_text=args.use_emo_text,
use_random=args.use_random,
)
for i, omni_out in enumerate(omni.generate(inputs, sampling_params_list=sampling_params)):
mm = omni_out.multimodal_output
if not mm:
print(f" [req {i}] No multimodal output.")
continue
audio, sr = extract_audio(mm)
if audio is None:
print(f" [req {i}] No waveform in multimodal_output.")
continue
out_path = str(output_dir / f"output_{i}.wav")
save_audio(audio.cpu(), out_path, sr)
print("Done.")
def parse_args():
parser = FlexibleArgumentParser(description="IndexTTS2 offline inference")
parser.add_argument("--model", required=True, help="HF model path for IndexTTS2.")
parser.add_argument("--text", default="你好,这是IndexTTS2语音合成测试。")
parser.add_argument("--ref-audio", required=True, help="Reference audio for voice cloning.")
parser.add_argument("--emo-audio", default=None, help="Emotion reference audio.")
parser.add_argument("--emo-text", default=None, help="Emotion description text.")
parser.add_argument(
"--emo-vector",
type=float,
nargs=8,
default=None,
help="8-dim emotion vector: happy angry sad afraid disgusted melancholic surprised calm.",
)
parser.add_argument("--emo-alpha", type=float, default=None, help="Emotion weight in [0, 1].")
parser.add_argument("--use-emo-text", action="store_true", help="Infer emotion vector from emo-text or text.")
parser.add_argument("--use-random", action="store_true", help="Use random emotion prototypes.")
parser.add_argument("--seed", type=int, default=None)
parser.add_argument(
"--output-dir",
default=os.path.join(
os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")),
"indextts2_output",
),
)
parser.add_argument("--deploy-config", default=None)
parser.add_argument("--stage-init-timeout", type=int, default=600)
return parser.parse_args()
if __name__ == "__main__":
main(parse_args())
@@ -0,0 +1,162 @@
"""Repeatable SoulX-Singer offline benchmark (preprocess + DiT, RTF + stage timings).
Usage:
python benchmark.py --model /path/to/SoulX-Singer --svs \\
--prompt-audio ../SoulX-Singer/example/audio/zh_prompt.mp3 \\
--target-audio ../SoulX-Singer/example/audio/music.mp3 \\
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
-o benchmark.wav
python benchmark.py ... --enable-diffusion-pipeline-profiler
"""
import argparse
import statistics
import time
from pathlib import Path
import soundfile as sf
from end2end import (
SVC_DEPLOY_CONFIG,
SVS_DEPLOY_CONFIG,
add_inference_args,
build_sampling,
extract_audio,
resolve_preprocess_weights_dir,
)
from vllm_omni.entrypoints.omni import Omni
def _audio_duration_sec(outputs) -> float:
audio_np, sr = extract_audio(outputs)
return audio_np.size / sr
def _stage_durations_from_output(outputs) -> dict[str, float]:
omni_out = outputs[0]
durations = getattr(omni_out, "stage_durations", None) or {}
if durations:
return dict(durations)
ro = getattr(omni_out, "request_output", None)
if ro is not None:
return dict(getattr(ro, "stage_durations", None) or {})
return {}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="SoulX-Singer offline benchmark")
add_inference_args(parser)
parser.add_argument("--warmup", type=int, default=1, help="Warmup runs (excluded from stats)")
parser.add_argument("--runs", type=int, default=3, help="Measured runs")
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable low-overhead stage timing in DiffusionOutput.stage_durations",
)
parser.add_argument(
"--enforce-eager",
action="store_true",
help="Disable torch.compile on diff_estimator (eager baseline for A/B)",
)
parser.add_argument(
"--dtype",
type=str,
default=None,
choices=["bfloat16", "bf16", "float16", "fp16", "half"],
help="DiT trunk dtype (default from deploy YAML). Use float16 to match upstream --fp16.",
)
parser.add_argument(
"--output",
"-o",
type=str,
default=None,
help="Save WAV from the last measured run (after warmup)",
)
parser.add_argument(
"--save-each-run",
action="store_true",
help="With --output, also write measured runs as <stem>_run1.wav, _run2.wav, ...",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir)
deploy_config = args.deploy_config or str(SVS_DEPLOY_CONFIG if args.svs else SVC_DEPLOY_CONFIG)
mode = "SVS" if args.svs else "SVC"
omni_kwargs: dict = {
"model": args.model,
"deploy_config": deploy_config,
"async_chunk": False, # SoulX-Singer currently supports only batch mode (pseudo-streaming was stashed)
}
if args.enable_diffusion_pipeline_profiler:
omni_kwargs["enable_diffusion_pipeline_profiler"] = True
if args.enforce_eager:
omni_kwargs["enforce_eager"] = True
if args.dtype is not None:
omni_kwargs["dtype"] = args.dtype
compile_mode = "eager" if args.enforce_eager else "torch.compile"
dtype_label = args.dtype or "from deploy YAML"
print(f"Loading SoulX-Singer {mode} from {args.model} [{compile_mode}, dtype={dtype_label}]")
omni = Omni(**omni_kwargs)
kind = "svs" if args.svs else "svc"
sampling = build_sampling(args, preprocess_weights_dir=preprocess_dir, kind=kind)
prompt = {"prompt_token_ids": [0]}
latencies_ms: list[float] = []
rtfs: list[float] = []
last_stages: dict[str, float] = {}
last_measured_outputs = None
total_runs = args.warmup + args.runs
for run_idx in range(total_runs):
is_warmup = run_idx < args.warmup
label = "warmup" if is_warmup else f"run {run_idx - args.warmup + 1}/{args.runs}"
t0 = time.perf_counter()
outputs = list(omni.generate([prompt], sampling))
elapsed_ms = (time.perf_counter() - t0) * 1000.0
audio_sec = _audio_duration_sec(outputs)
rtf = (elapsed_ms / 1000.0) / audio_sec if audio_sec > 0 else float("inf")
print(f"[{label}] client={elapsed_ms:.1f} ms, audio={audio_sec:.2f}s, RTF={rtf:.3f}")
if not is_warmup:
latencies_ms.append(elapsed_ms)
rtfs.append(rtf)
last_stages = _stage_durations_from_output(outputs)
last_measured_outputs = outputs
if args.output and args.save_each_run:
measured_idx = run_idx - args.warmup + 1
out_path = Path(args.output)
run_path = out_path.with_name(f"{out_path.stem}_run{measured_idx}{out_path.suffix}")
audio_np, sr = extract_audio(outputs)
sf.write(str(run_path), audio_np, sr)
print(f" saved {run_path} ({sr} Hz, {audio_np.size / sr:.2f}s)")
omni.close()
if args.output and last_measured_outputs is not None:
audio_np, sr = extract_audio(last_measured_outputs)
out_path = Path(args.output)
out_path.parent.mkdir(parents=True, exist_ok=True)
sf.write(str(out_path), audio_np, sr)
print(f"\nSaved last measured run → {out_path} ({sr} Hz, {audio_np.size / sr:.2f}s)")
if latencies_ms:
print("\n=== Summary (measured runs) ===")
print(f"client_ms: mean={statistics.mean(latencies_ms):.1f}, stdev={statistics.pstdev(latencies_ms):.1f}")
print(f"RTF: mean={statistics.mean(rtfs):.3f}, stdev={statistics.pstdev(rtfs):.3f}")
if last_stages:
print("\nStage durations (last measured run):")
for name, value in sorted(last_stages.items()):
if name.endswith("_ms"):
print(f" {name}: {value:.1f} ms")
else:
print(f" {name}: {value * 1000.0:.1f} ms")
if __name__ == "__main__":
main()
@@ -0,0 +1,285 @@
"""Offline SoulX-Singer SVS / SVC: single-stage DiT (preprocess inline).
Pass raw prompt/target audio; vLLM-Omni runs integrated preprocess then inference.
Usage:
# SVS (default)
python end2end.py --model /path/to/SoulX-Singer \\
--prompt-audio path/to/prompt/audio.mp3 \\
--target-audio path/to/target/audio.mp3 \\
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
-o output.wav
# SVC (config.json must declare SoulXSingerSVCPipeline; --model points at SVC view)
python end2end.py --model /path/to/SoulX-Singer-svc \\
--prompt-audio path/to/prompt/audio.mp3 \\
--target-audio path/to/target/audio.mp3 \\
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
-o output_svc.wav
"""
import argparse
import os
from pathlib import Path
import numpy as np
import soundfile as sf
from vllm_omni.diffusion.models.soulx_singer.utils import resolve_soulx_kind, validate_soulx_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
REPO_ROOT = Path(__file__).resolve().parents[4]
SVS_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svs.yaml"
SVC_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svc.yaml"
DEFAULT_PROMPT_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "zh_prompt.mp3"
DEFAULT_TARGET_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "music.mp3"
_SAMPLE_RATE = 24000
def _require_paths(paths: dict[str, str | None]) -> None:
missing = [name for name, p in paths.items() if not p or not os.path.isfile(p)]
if missing:
raise FileNotFoundError("Missing file(s): " + ", ".join(f"{k}={paths[k]!r}" for k in missing))
def resolve_preprocess_weights_dir(explicit: str | None) -> str:
if explicit:
path = Path(explicit).expanduser().resolve()
if not (path / "rmvpe" / "rmvpe.pt").is_file():
raise FileNotFoundError(f"--preprocess-weights-dir must contain rmvpe/rmvpe.pt: {path}")
return str(path)
env = os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR")
if env:
path = Path(env).expanduser().resolve()
if (path / "rmvpe" / "rmvpe.pt").is_file():
return str(path)
raise FileNotFoundError(
"SoulX preprocess weights not found. Pass --preprocess-weights-dir or set "
"SOULX_PREPROCESS_WEIGHTS_DIR to a directory containing rmvpe/rmvpe.pt "
"(huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess)."
)
def build_sampling(
args: argparse.Namespace,
*,
preprocess_weights_dir: str,
kind: str,
) -> OmniDiffusionSamplingParams:
extra: dict = {
"prompt_audio": os.path.abspath(args.prompt_audio),
"target_audio": os.path.abspath(args.target_audio),
"preprocess_weights_dir": preprocess_weights_dir,
"vocal_sep": args.vocal_sep,
"auto_shift": args.auto_shift,
"pitch_shift": args.pitch_shift,
}
if kind == "svs":
extra["language"] = args.language
extra["control"] = args.control
if args.prompt_metadata_path:
extra["prompt_metadata_path"] = os.path.abspath(args.prompt_metadata_path)
if args.target_metadata_path:
extra["target_metadata_path"] = os.path.abspath(args.target_metadata_path)
if args.prompt_metadata_path or args.target_metadata_path:
extra["audio_path"] = os.path.abspath(args.prompt_wav_path or args.prompt_audio)
extra = validate_soulx_extra_args(kind, extra)
return OmniDiffusionSamplingParams(
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
extra_args=extra,
)
def extract_audio(outputs) -> tuple[np.ndarray, int]:
out = outputs[0]
mm = getattr(out, "multimodal_output", None)
if mm is None:
ro = getattr(out, "request_output", None)
if ro is not None:
mm = getattr(ro, "multimodal_output", None)
if mm is None and ro.outputs:
mm = getattr(ro.outputs[0], "multimodal_output", None)
if not mm or "audio" not in mm:
raise RuntimeError("No audio in multimodal_output")
audio_val = mm["audio"]
if isinstance(audio_val, list):
chunks = [
np.asarray(chunk.detach().cpu().float().numpy() if hasattr(chunk, "detach") else chunk).reshape(-1)
for chunk in audio_val
if chunk is not None
]
audio_np = np.concatenate(chunks, axis=0) if chunks else np.array([], dtype=np.float32)
elif hasattr(audio_val, "cpu"):
audio_np = audio_val.detach().cpu().numpy().squeeze()
else:
audio_np = np.asarray(audio_val, dtype=np.float32).squeeze()
sr_val = mm.get("sr", mm.get("audio_sample_rate", _SAMPLE_RATE))
if isinstance(sr_val, list) and sr_val:
sr_val = sr_val[-1]
if hasattr(sr_val, "item"):
sr_val = sr_val.item()
return audio_np.astype(np.float32, copy=False), int(sr_val)
def add_inference_args(parser: argparse.ArgumentParser) -> None:
mode = parser.add_mutually_exclusive_group()
mode.add_argument(
"--svc",
action="store_true",
help="Assert SVC mode (must match config.json architectures)",
)
parser.add_argument(
"--model",
type=str,
required=True,
help="SoulX-Singer weight directory ($BASE for SVS, $SVC_DIR for SVC)",
)
parser.add_argument(
"--deploy-config",
type=str,
default=None,
help="Override deploy YAML (default: soulxsinger_svs.yaml or soulxsinger_svc.yaml)",
)
parser.add_argument(
"--prompt-audio",
type=str,
default=str(DEFAULT_PROMPT_AUDIO),
help="Prompt / reference audio (mp3 or wav)",
)
parser.add_argument(
"--target-audio",
type=str,
default=str(DEFAULT_TARGET_AUDIO),
help="Target content audio (mp3 or wav)",
)
parser.add_argument(
"--prompt-metadata-path",
type=str,
default=None,
help="SVS precomputed prompt metadata.json (skips online ASR/ROSVOT for prompt)",
)
parser.add_argument(
"--target-metadata-path",
type=str,
default=None,
help="SVS precomputed target metadata.json (skips online ASR/ROSVOT for target)",
)
parser.add_argument(
"--prompt-wav-path",
type=str,
default=None,
help="SVS prompt vocal wav for metadata processor when using --prompt-metadata-path (default: --prompt-audio)",
)
parser.add_argument(
"--preprocess-weights-dir",
type=str,
default=None,
help="SoulX-Singer-Preprocess weights (or set SOULX_PREPROCESS_WEIGHTS_DIR)",
)
parser.add_argument(
"--language",
type=str,
default="Mandarin",
help="SVS lyrics language (Mandarin, Cantonese, English)",
)
parser.add_argument(
"--control",
type=str,
default="melody",
choices=["score", "melody"],
help="SVS control mode",
)
parser.add_argument(
"--vocal-sep",
action=argparse.BooleanOptionalAction,
default=False,
help="Run vocal separation during preprocess",
)
parser.add_argument(
"--auto-shift",
action=argparse.BooleanOptionalAction,
default=True,
help="Auto pitch shift to align target pitch with prompt",
)
parser.add_argument("--pitch-shift", type=int, default=0)
parser.add_argument("--num-inference-steps", type=int, default=32)
parser.add_argument("--guidance-scale", type=float, default=3.0)
parser.add_argument(
"--seed",
type=int,
default=42,
help="Optional CFM RNG seed. Omit for non-deterministic sampling.",
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="SoulX-Singer offline SVS / SVC (preprocess + DiT)")
add_inference_args(parser)
parser.add_argument("--output", "-o", type=str, default="output.wav")
args = parser.parse_args()
return args
def main() -> None:
args = parse_args()
kind = resolve_soulx_kind(args.model)
if not (args.svc ^ (kind == "svs")):
raise ValueError(f"mode mismatch: {args.model}/config.json (declares {kind!r}), --svc={args.svc}")
required = {
"prompt_audio": args.prompt_audio,
"target_audio": args.target_audio,
}
if kind == "svs" and bool(args.prompt_metadata_path) ^ bool(args.target_metadata_path):
raise ValueError(
"SVS precomputed metadata requires both --prompt-metadata-path and "
"--target-metadata-path (or neither, for online preprocess)."
)
if args.prompt_metadata_path:
required["prompt_metadata_path"] = args.prompt_metadata_path
required["target_metadata_path"] = args.target_metadata_path
if args.prompt_wav_path:
required["prompt_wav_path"] = args.prompt_wav_path
_require_paths(required)
preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir)
deploy_config = args.deploy_config or str(SVC_DEPLOY_CONFIG if kind == "svc" else SVS_DEPLOY_CONFIG)
mode = kind.upper()
print(f"Loading SoulX-Singer {mode} from {args.model}")
print(f" deploy: {deploy_config}")
print(f" prompt: {args.prompt_audio}")
print(f" target: {args.target_audio}")
omni = Omni(model=args.model, deploy_config=deploy_config, async_chunk=False) # SoulX-Singer is batch-only
sampling = build_sampling(
args,
preprocess_weights_dir=preprocess_dir,
kind=kind,
)
prompt = {"prompt_token_ids": [0], "multi_modal_data": {}}
print("Running generate (inline preprocess + DiT)...")
outputs = list(omni.generate([prompt], sampling))
if not outputs:
raise RuntimeError("No output from omni.generate")
err = getattr(outputs[0], "error", None)
if err is not None:
raise RuntimeError(f"Pipeline error: {err}")
audio_np, sr = extract_audio(outputs)
sf.write(args.output, audio_np, sr)
print(f"Saved {args.output} ({sr} Hz, {len(audio_np) / sr:.2f}s)")
omni.close()
if __name__ == "__main__":
main()
@@ -6,9 +6,11 @@ A unified script for text-to-video generation. Supports multiple models with mod
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|---|---|---|---|---|---|
| `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 |
| `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 |
| `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) |
| `BestWishYsh/Helios-Base` / `Helios-Mid` / `Helios-Distilled` | 384x640 | 99 | 50 | 5.0 / 5.0 / 1.0 | — |
## Local CLI Usage
@@ -30,6 +32,28 @@ python text_to_video.py \
--output t2v_out.mp4
```
### Wan2.1 VACE (T2V)
VACE text-to-video uses this shared entrypoint. Conditional VACE tasks use
the shared [`image_to_video.py`](../image_to_video/README.md#wan21-vace-conditional-tasks)
entrypoint, which constructs the pipeline-native conditioning data from the
provided media inputs. No explicit mode parameter is required.
```bash
python text_to_video.py \
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
--prompt "A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves." \
--seed 0 \
--height 480 \
--width 832 \
--num-frames 81 \
--num-inference-steps 30 \
--guidance-scale 5.0 \
--flow-shift 5.0 \
--vae-use-tiling \
--output vace_t2v_output.mp4
```
LTX2 example:
```bash
@@ -101,6 +125,20 @@ python text_to_video.py \
--output quick_test.mp4
```
### Cosmos3
```bash
python text_to_video.py \
--model nvidia/Cosmos3-Nano \
--prompt "A robot arm is cleaning a plate in the kitchen." \
--negative-prompt "blurry, distorted, low quality, jittery, deformed" \
--height 720 --width 1280 --num-frames 189 --fps 24 \
--num-inference-steps 35 --guidance-scale 6.0 \
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false,
"use_resolution_template": false, "use_duration_template": false}' \
--output cosmos3_t2v.mp4
```
### Helios (T2V)
Helios ships three variants. Model-specific knobs (declared in
@@ -167,6 +205,7 @@ python text_to_video.py \
- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio).
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
- `--flow-shift`: scheduler flow_shift parameter.
- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).
### Wan2.2-specific
@@ -11,12 +11,23 @@ import numpy as np
import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
_MODEL_PRESETS = {
"vace": {
"height": 480,
"width": 832,
"num_frames": 81,
"num_inference_steps": 30,
"guidance_scale": 5.0,
"fps": 16,
"output": "vace_t2v_output.mp4",
},
"wan": {
"height": 720,
"width": 1280,
@@ -35,6 +46,16 @@ _MODEL_PRESETS = {
"fps": 24,
"output": "hunyuan_video_15_output.mp4",
},
"cosmos": {
"height": 720,
"width": 1280,
"num_frames": 189,
"num_inference_steps": 35,
"guidance_scale": 6.0,
"fps": 24,
"flow_shift": 10.0,
"output": "cosmos3_t2v_output.mp4",
},
"helios": {
"height": 384,
"width": 640,
@@ -49,6 +70,10 @@ _MODEL_PRESETS = {
def _detect_preset(model: str) -> dict:
model_lower = model.lower()
if "vace" in model_lower:
return _MODEL_PRESETS["vace"]
if "cosmos" in model_lower:
return _MODEL_PRESETS["cosmos"]
if "hunyuan" in model_lower:
return _MODEL_PRESETS["hunyuan"]
if "helios" in model_lower:
@@ -106,8 +131,12 @@ def parse_args() -> argparse.Namespace:
type=parse_extra_body,
default=None,
help="JSON dict of model-specific extra_body params (declared in vllm_omni/model_extras/), "
"merged into sampling extra_args. Example (Helios-Distilled): "
'\'{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}\'.',
"merged into sampling extra_args. Unknown keys for the chosen model are dropped. "
'Cosmos3 example: \'{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false, '
'"use_resolution_template": false, "use_duration_template": false}\' '
'(add "generate_sound": true and optional "sound_duration" for synchronized audio). '
'Helios-Distilled example: \'{"is_enable_stage2": true, '
'"pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}\'.',
)
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
@@ -357,7 +386,14 @@ def main():
omni_kwargs["cache_config"] = cache_config
omni_kwargs["enable_cache_dit_summary"] = args.enable_cache_dit_summary
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
# gate is an engine-level config (offline analog of the server's --no-guardrails).
if args.extra_body and "guardrails" in args.extra_body:
omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
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)
if profiler_enabled:
print("[Profiler] Starting profiling...")
@@ -392,14 +428,21 @@ def main():
)
if args.guidance_scale_high is not None:
sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
if args.extra_body:
# Model-specific knobs (declared in vllm_omni/model_extras/) routed via extra_args.
sampling_kwargs["extra_args"] = dict(args.extra_body)
sampling_params = OmniDiffusionSamplingParams(**sampling_kwargs)
# Route model-specific knobs through extra_body, filtered against the model's
# declared extra_body_params. Models without a declaration only forward explicit
# --extra-body JSON (preserving the generic flags' legacy behavior).
extra_body = dict(args.extra_body or {})
if declared_extra_body_params:
apply_declared_extra_args(sampling_params, declared_extra_body_params, extra_body)
elif extra_body:
sampling_params.extra_args.update({k: v for k, v in extra_body.items() if v is not None})
generation_start = time.perf_counter()
frames = omni.generate(
prompt_dict,
OmniDiffusionSamplingParams(**sampling_kwargs),
sampling_params,
)
generation_end = time.perf_counter()
generation_time = generation_end - generation_start
@@ -1,148 +0,0 @@
# VACE Video Generation
Generate videos from text prompts, images, or video conditions using vLLM-Omni's VACE diffusion pipeline.
- `vace_video_generation.py`: command-line script for multi-mode video generation with advanced options.
## Table of Contents
- [Overview](#overview)
- [Quick Start](#quick-start)
- [Key Arguments](#key-arguments)
- [More CLI Examples](#more-cli-examples)
## Overview
[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single unified model, including text-to-video, image-to-video, first-last-frame interpolation, inpainting, and reference image-guided generation.
### Supported Models
| Model | Architecture | Peak VRAM (GiB) * | Model Weights (GiB) | HuggingFace |
|-------|-------------|-------------------|---------------------|-------------|
| Wan2.1-VACE (1.3B) | Wan2.1 | TBD | ~10 | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
| Wan2.1-VACE (14B) | Wan2.1 | TBD | ~38 | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
!!! info
*Peak VRAM: based on basic single-card usage, 480x832 resolution, 81 frames, 30 inference steps, without any acceleration/optimization features.
Default model: `Wan-AI/Wan2.1-VACE-14B-diffusers`
## Quick Start
### Python API
Text-to-video generation:
```python
import torch
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
if __name__ == "__main__":
omni = Omni(model="Wan-AI/Wan2.1-VACE-1.3B-diffusers")
prompt = "A sleek robot stands in a vast warehouse filled with boxes"
outputs = omni.generate(
prompt,
OmniDiffusionSamplingParams(
height=480,
width=832,
num_frames=81,
num_inference_steps=30,
guidance_scale=5.0,
),
)
video = outputs[0].images
from diffusers.utils import export_to_video
export_to_video(list(video[0]), "t2v_output.mp4", fps=16)
omni.close()
```
### Local CLI Usage
```bash
python vace_video_generation.py \
--mode t2v \
--prompt "A sleek robot stands in a vast warehouse filled with boxes" \
--height 480 --width 832 --num-frames 81 \
--output t2v_output.mp4
```
## Key Arguments
| Argument | Type | Default | Description |
| -------- | ---- | ------- | ----------- |
| `--mode` | str | `"t2v"` | VACE task mode: `t2v`, `i2v`, `v2lf`, `flf2v`, `inpaint`, `r2v` |
| `--model` | str | `"Wan-AI/Wan2.1-VACE-14B-diffusers"` | Model ID or local path |
| `--prompt` | str | `"A cat walking in a garden"` | Text description of desired video |
| `--negative-prompt` | str | `""` | Negative prompt for classifier-free guidance |
| `--image` | str | `None` | Input image path (for I2V, R2V, FLF2V, inpaint modes) |
| `--last-image` | str | `None` | Last frame image path (for FLF2V mode) |
| `--height` | int | `480` | Output video height in pixels (should be a multiple of 16) |
| `--width` | int | `832` | Output video width in pixels (should be a multiple of 16) |
| `--num-frames` | int | `81` | Number of video frames to generate |
| `--num-inference-steps` | int | `30` | Number of denoising steps (more steps = higher quality, slower) |
| `--guidance-scale` | float | `5.0` | Classifier-free guidance scale |
| `--flow-shift` | float | `5.0` | Scheduler flow shift parameter |
| `--seed` | int | `42` | Random seed for deterministic sampling |
| `--fps` | int | `16` | Frames per second for the saved MP4 |
| `--output` | str | `"vace_output.mp4"` | Path to save the generated video |
| `--vae-use-tiling` | flag | on | Enable VAE tiling for memory optimization |
| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree for multi-GPU inference |
| `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference |
| `--tensor-parallel-size` | int | — | Tensor parallel size |
| `--enforce-eager` | flag | off | Disable torch.compile |
> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options (`--ulysses-degree`, `--cfg-parallel-size`).
## More CLI Examples
### Image-to-Video (I2V)
First frame is kept, remaining frames are generated:
```bash
python vace_video_generation.py \
--mode i2v \
--image astronaut.jpg \
--prompt "An astronaut emerging from a cracked egg on the moon" \
--height 480 --width 832 --num-frames 81 \
--output i2v_output.mp4
```
### First-Last-Frame Interpolation (FLF2V)
```bash
python vace_video_generation.py \
--mode flf2v \
--image first_frame.jpg --last-image last_frame.jpg \
--prompt "A bird takes off from a branch and lands on another" \
--height 512 --width 512 --num-frames 81 \
--output flf2v_output.mp4
```
### Inpainting
Center vertical stripe is masked and regenerated:
```bash
python vace_video_generation.py \
--mode inpaint \
--image scene.jpg \
--prompt "Shrek walks out of a building" \
--height 480 --width 832 --num-frames 81 \
--output inpaint_output.mp4
```
### Reference Image-guided (R2V)
```bash
python vace_video_generation.py \
--mode r2v \
--image reference.jpg \
--prompt "Camera slowly zooms out from the character" \
--height 480 --width 832 --num-frames 81 \
--output r2v_output.mp4
```
@@ -1,209 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""VACE video generation example.
VACE (Video All-in-one Creation Engine) supports multiple video tasks:
- T2V: Text-to-Video
- I2V: Image-to-Video (first frame conditioning)
- V2LF: Video-to-Last-Frame
- FLF2V: First-Last-Frame interpolation
- Inpainting: Masked region generation
- R2V: Reference image-guided generation
Usage examples:
# T2V (text-to-video)
python vace_video_generation.py --mode t2v --prompt "A robot in a warehouse"
# I2V (image-to-video, first frame kept)
python vace_video_generation.py --mode i2v --image input.jpg --prompt "..."
# FLF2V (first-last frame interpolation)
python vace_video_generation.py --mode flf2v --image first.jpg --last-image last.jpg
# R2V (reference image guided)
python vace_video_generation.py --mode r2v --image ref.jpg --prompt "..."
"""
import argparse
import time
from pathlib import Path
import numpy as np
import PIL.Image
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.platforms import current_omni_platform
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="VACE video generation.")
parser.add_argument(
"--model",
default="Wan-AI/Wan2.1-VACE-14B-diffusers",
help="VACE model ID or local path.",
)
parser.add_argument(
"--mode",
default="t2v",
choices=["t2v", "i2v", "v2lf", "flf2v", "inpaint", "r2v"],
help="Generation mode.",
)
parser.add_argument("--prompt", default="A cat walking in a garden", help="Text prompt.")
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
parser.add_argument("--image", type=str, default=None, help="Input image path (for I2V, R2V, FLF2V, inpaint).")
parser.add_argument("--last-image", type=str, default=None, help="Last frame image path (for FLF2V).")
parser.add_argument("--video-dir", type=str, default=None, help="Directory of video frames (for inpaint).")
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG guidance scale.")
parser.add_argument("--height", type=int, default=480, help="Video height.")
parser.add_argument("--width", type=int, default=832, help="Video width.")
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
parser.add_argument("--num-inference-steps", type=int, default=30, help="Sampling steps.")
parser.add_argument("--flow-shift", type=float, default=5.0, help="Scheduler flow_shift.")
parser.add_argument("--output", type=str, default="vace_output.mp4", help="Output video path.")
parser.add_argument("--fps", type=int, default=16, help="Output video FPS.")
parser.add_argument("--vae-use-tiling", action="store_true", default=True, help="Enable VAE tiling.")
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
parser.add_argument("--ring-degree", type=int, default=1, help="Ring attention degree.")
parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
return parser.parse_args()
def build_prompts(args):
"""Build prompt dict with multi_modal_data based on mode."""
h, w, nf = args.height, args.width, args.num_frames
gray = PIL.Image.new("RGB", (w, h), (128, 128, 128))
mask_black = PIL.Image.new("L", (w, h), 0)
mask_white = PIL.Image.new("L", (w, h), 255)
prompt_data = {
"prompt": args.prompt,
"negative_prompt": args.negative_prompt,
}
if args.mode == "t2v":
return prompt_data
if args.mode == "r2v":
assert args.image, "--image required for R2V mode"
ref_img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
prompt_data["multi_modal_data"] = {"reference_images": [ref_img]}
return prompt_data
if args.mode == "i2v":
assert args.image, "--image required for I2V mode"
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
prompt_data["multi_modal_data"] = {
"video": [img] + [gray] * (nf - 1),
"mask": [mask_black] + [mask_white] * (nf - 1),
}
return prompt_data
if args.mode == "v2lf":
assert args.image, "--image required for V2LF mode"
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
prompt_data["multi_modal_data"] = {
"video": [gray] * (nf - 1) + [img],
"mask": [mask_white] * (nf - 1) + [mask_black],
}
return prompt_data
if args.mode == "flf2v":
assert args.image and args.last_image, "--image and --last-image required for FLF2V"
first = PIL.Image.open(args.image).convert("RGB").resize((w, h))
last = PIL.Image.open(args.last_image).convert("RGB").resize((w, h))
prompt_data["multi_modal_data"] = {
"video": [first] + [gray] * (nf - 2) + [last],
"mask": [mask_black] + [mask_white] * (nf - 2) + [mask_black],
}
return prompt_data
if args.mode == "inpaint":
assert args.image, "--image required for inpaint mode"
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
d = 80
frames, masks = [], []
for _ in range(nf):
base = np.array(img).copy()
mask = PIL.Image.new("L", (w, h), 0)
stripe = PIL.Image.new("L", (2 * d, h), 255)
mask.paste(stripe, (w // 2 - d, 0))
base[np.array(mask) > 128] = 128
frames.append(PIL.Image.fromarray(base))
masks.append(mask)
prompt_data["multi_modal_data"] = {"video": frames, "mask": masks}
return prompt_data
raise ValueError(f"Unknown mode: {args.mode}")
def main():
args = parse_args()
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
parallel_config = DiffusionParallelConfig(
ulysses_degree=args.ulysses_degree,
ring_degree=args.ring_degree,
cfg_parallel_size=args.cfg_parallel_size,
)
omni = Omni(
model=args.model,
vae_use_tiling=args.vae_use_tiling,
flow_shift=args.flow_shift,
enforce_eager=args.enforce_eager,
parallel_config=parallel_config,
)
prompt_data = build_prompts(args)
print(f"\n{'=' * 60}")
print(f"VACE {args.mode.upper()} Generation")
print(f" Model: {args.model}")
print(f" Size: {args.width}x{args.height}, {args.num_frames} frames, {args.num_inference_steps} steps")
print(f"{'=' * 60}\n")
start = time.perf_counter()
outputs = omni.generate(
prompt_data,
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,
),
)
elapsed = time.perf_counter() - start
video = outputs[0].images
if isinstance(video, list):
video = video[0]
if isinstance(video, torch.Tensor):
video = video.cpu().numpy()
if video.ndim == 5:
video = video[0]
print(f"Output shape: {video.shape}, Time: {elapsed:.1f}s")
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
from diffusers.utils import export_to_video
if np.issubdtype(video.dtype, np.integer):
video = video.astype(np.float32) / 255.0
export_to_video(list(video), str(output_path), fps=args.fps)
print(f"Saved to {output_path}")
omni.close()
if __name__ == "__main__":
main()
@@ -88,7 +88,7 @@ curl -s http://127.0.0.1:8091/v1/chat/completions \
### Optional knobs
Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_args.image_gen`:
Pass image-gen overrides as flat keys on the diffusion-stage `sampling_params_list[1].extra_args`:
| Key | Default | Description |
| --- | --- | --- |
@@ -97,25 +97,48 @@ Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_
| `cfg` | 2.0 | Classifier-free guidance scale. |
| `seed` | 42 | Per-request RNG seed (deterministic when ByT5 is also seed-stable). |
| `byte5_text` | (auto) | Override the glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "...". ` format. |
| `negative_prompt` | empty | Real CFG negative conditioning (must be set on **stage-0 thinker** `sampling_params` so `expand_cfg_prompts` spawns the companion). |
| `negative_prompt` | empty | Real CFG negative conditioning (set on **stage-0 thinker** `extra_args` so `expand_cfg_prompts` spawns the companion). |
Example with all the knobs:
```bash
curl http://127.0.0.1:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
"modalities": ["image"],
"sampling_params_list": [
{"temperature":0.4,"top_p":0.9,"top_k":1,"max_tokens":1,"seed":42,
"extra_args":{"image_gen":{"negative_prompt":"ugly, blurry, distorted"}}},
{"seed":42,"extra_args":{"image_gen":{
"steps":6,"cfg":1.5,"height":512,"width":512,"seed":123,
"byte5_text":["理解与生成统一"]}}}
],
"messages": [{"role": "user", "content": "Draw a poster."}]
}' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen_knobs.png
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
"modalities": ["image"],
"sampling_params_list": [
{
"temperature": 0.4,
"top_p": 0.9,
"top_k": 1,
"max_tokens": 1,
"seed": 42,
"extra_args": {
"negative_prompt": "ugly, blurry, distorted"
}
},
{
"seed": 42,
"extra_args": {
"steps": 6,
"cfg": 1.5,
"height": 512,
"width": 512,
"seed": 123,
"byte5_text": ["理解与生成统一"]
}
}
],
"messages": [
{
"role": "user",
"content": "Draw a poster."
}
]
}' \
| jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
| base64 -d > ming_imagegen_knobs.png
```
### img2img (reference image)
@@ -0,0 +1,94 @@
# Streaming Video Generation
This example uses the custom WebSocket endpoint `WS /v1/realtime/video` to receive a video byte stream as chunks are produced.
It covers text-only video generation. Image/reference input is intentionally not included for now.
## Start The Server
Start a diffusion video model with streaming output enabled:
```bash
vllm serve BestWishYsh/Helios-Distilled \
--omni \
--diffusion-streaming-output \
--port 8000
```
The `--diffusion-streaming-output` CLI flag is forwarded as `streaming_output=True` in the default diffusion stage `engine_args`, then loaded by `OmniDiffusionConfig.from_kwargs()`.
## WebSocket Protocol
| Direction | Message | Format | Description |
| --- | --- | --- | --- |
| Client to server | `session.start` | JSON text: `{"type":"session.start","model":"...","prompt":"...","format":"m4s"}` | Starts generation. `format` is optional and accepts `m4s` (default). Sampling fields such as `width`, `height`, `fps`, `num_frames`, and `extra_params` may be included. |
| Server to client | `video.start` | JSON text: `{"type":"video.start","request_id":"...","format":"m4s","config":{...}}` | Confirms the session and mirrors the accepted `format`. |
| Server to client | Video chunk | Binary WebSocket frame | Fragmented MP4 (`m4s`) video bytes. |
| Client to server | `session.stop` | JSON text: `{"type":"session.stop"}` | Requests cancellation of the active session. |
| Client to server | `session.ping` | JSON text: `{"type":"session.ping"}` | Optional keepalive; refreshes the server stall clock. |
| Server to client | `session.done` | JSON text: `{"type":"session.done","request_id":"...","chunks":3,"stopped":false}` | Ends a completed or stopped session. |
| Server to client | `session.pong` | JSON text: `{"type":"session.pong"}` | Reply to `session.ping`. |
| Server to client | `error` | JSON text: `{"type":"error","message":"..."}` | Reports invalid input, unsupported formats, generation failures, control-message errors, or stall timeout. |
During generation the client normally sends only `session.start` and then receives binary chunks; silence on the client socket is expected. The server closes the session with a stall error only when there is no engine progress and no `session.ping` for about 60 seconds.
## Install Client Dependency
```bash
pip install websockets
# For the Gradio demo:
pip install 'vllm-omni[demo]' websockets
```
## Run The Client
```bash
python streaming_video_client.py \
--host 127.0.0.1 \
--port 8000 \
--model BestWishYsh/Helios-Distilled \
--prompt "A serene lakeside sunrise with mist over the water." \
--width 640 \
--height 384 \
--fps 16 \
--num-frames 99 \
--guidance-scale 1.0 \
--seed 42 \
--output helios_stream.mp4
```
The client sends one `session.start` message, prints each received binary video chunk with its byte size and elapsed time, and saves the received bytes to `--output` after `session.done`.
The client remuxes the gathered stream to a regular progressive MP4 file so that local playback knows the video duration.
## Run The Gradio Demo
```bash
python gradio_demo.py \
--host 127.0.0.1 \
--port 7860
```
The Gradio demo requests fMP4 (`m4s`) chunks and appends them directly in the browser with a Media Source Extensions player.
## Model Choice
### Helios
The example uses `BestWishYsh/Helios-Distilled` model by default.
To ensure streaming-level generation speed, `pyramid_num_inference_steps_list` is suggested to be as low as `[1, 1, 1]`. Both example clients uses the following Helios-Distilled preset by default:
```json
{
"is_enable_stage2": true,
"pyramid_num_stages": 3,
"pyramid_num_inference_steps_list": [1, 1, 1],
"is_amplify_first_chunk": true
}
```
Disable it in the CLI example with `--no-helios-distilled-preset`, or override/extend it with `--extra-params`:
```bash
python streaming_video_client.py \
--extra-params '{"pyramid_num_inference_steps_list":[2, 2, 2]}'
```
@@ -0,0 +1,301 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Gradio demo for `/v1/realtime/video`.
Usage:
python gradio_demo.py --host 127.0.0.1 --port 7860
Start a vLLM-Omni server with streaming output enabled first, for example:
vllm serve BestWishYsh/Helios-Distilled --omni --diffusion-streaming-output --port 8000
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Any
try:
import gradio as gr
except ImportError:
raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
DEFAULT_PROMPT = "A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement."
VIDEO_STREAM_VIEW_HTML = Path(__file__).with_name("video-stream-view.html")
VIDEO_STREAM_VIEW_JS = Path(__file__).with_name("video-stream-view.js")
def _maybe_set(payload: dict[str, Any], key: str, value: Any) -> None:
if value is None:
return
if isinstance(value, str) and not value.strip():
return
payload[key] = value
def _optional_int(value: int | float | None) -> int | None:
return None if value is None else int(value)
def _optional_float(value: int | float | None) -> float | None:
return None if value is None else float(value)
def _build_session_start(
*,
model: str,
prompt: str,
negative_prompt: str,
width: int | float | None,
height: int | float | None,
fps: int | float | None,
num_frames: int | float | None,
num_inference_steps: int | float | None,
guidance_scale: int | float | None,
guidance_scale_2: int | float | None,
boundary_ratio: int | float | None,
flow_shift: int | float | None,
true_cfg_scale: int | float | None,
seed: int | float | None,
) -> dict[str, Any]:
payload: dict[str, Any] = {
"type": "session.start",
"model": model,
"prompt": prompt,
"format": "m4s",
}
_maybe_set(payload, "negative_prompt", negative_prompt)
_maybe_set(payload, "width", _optional_int(width))
_maybe_set(payload, "height", _optional_int(height))
_maybe_set(payload, "fps", _optional_int(fps))
_maybe_set(payload, "num_frames", _optional_int(num_frames))
_maybe_set(payload, "num_inference_steps", _optional_int(num_inference_steps))
_maybe_set(payload, "guidance_scale", _optional_float(guidance_scale))
_maybe_set(payload, "guidance_scale_2", _optional_float(guidance_scale_2))
_maybe_set(payload, "boundary_ratio", _optional_float(boundary_ratio))
_maybe_set(payload, "flow_shift", _optional_float(flow_shift))
_maybe_set(payload, "true_cfg_scale", _optional_float(true_cfg_scale))
_maybe_set(payload, "seed", _optional_int(seed))
if "Helios" in model:
payload["extra_params"] = {
"is_enable_stage2": True,
"pyramid_num_stages": 3,
"pyramid_num_inference_steps_list": [1, 1, 1],
"is_amplify_first_chunk": True,
}
return payload
def _set_optional_field_enabled(enabled: bool) -> Any:
return gr.update(interactive=enabled)
def build_browser_stream_config(
host: str,
port: int | float,
model: str,
prompt: str,
negative_prompt: str,
width: int | float | None,
height: int | float | None,
fps: int | float,
num_frames: int | float | None,
num_inference_steps: int | float | None,
include_guidance_scale: bool,
guidance_scale: int | float | None,
include_guidance_scale_2: bool,
guidance_scale_2: int | float | None,
include_seed: bool,
seed: int | float | None,
include_boundary_ratio: bool,
boundary_ratio: int | float | None,
include_flow_shift: bool,
flow_shift: int | float | None,
include_true_cfg_scale: bool,
true_cfg_scale: int | float | None,
) -> tuple[str, Any]:
"""Build the browser-side WebSocket request config for the HTML/MSE player."""
if not model.strip():
raise gr.Error("Model is required.")
if not prompt.strip():
raise gr.Error("Prompt is required.")
payload = _build_session_start(
model=model.strip(),
prompt=prompt.strip(),
negative_prompt=negative_prompt,
width=width,
height=height,
fps=fps,
num_frames=num_frames,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale if include_guidance_scale else None,
guidance_scale_2=guidance_scale_2 if include_guidance_scale_2 else None,
boundary_ratio=boundary_ratio if include_boundary_ratio else None,
flow_shift=flow_shift if include_flow_shift else None,
true_cfg_scale=true_cfg_scale if include_true_cfg_scale else None,
seed=seed if include_seed else None,
)
config = {
"url": f"ws://{host}:{int(port)}/v1/realtime/video",
"payload": payload,
}
return json.dumps(config, ensure_ascii=False), gr.update(value="Streaming...", interactive=False)
def create_demo() -> gr.Blocks:
with gr.Blocks(title="Streaming Video Generation") as demo:
gr.Markdown("# Streaming Video Generation")
gr.Markdown(
"Connects to `WS /v1/realtime/video` in the browser and appends fMP4 chunks "
"directly to a Media Source Extensions video player."
)
with gr.Row():
with gr.Column(scale=1):
host = gr.Textbox(label="Server Host", value="127.0.0.1")
port = gr.Number(label="Server Port", value=8000, precision=0)
model = gr.Textbox(label="Model", value="BestWishYsh/Helios-Distilled")
prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=3)
negative_prompt = gr.Textbox(label="Negative Prompt", value="", lines=2)
with gr.Row():
width = gr.Number(label="Width", value=640, precision=0)
height = gr.Number(label="Height", value=384, precision=0)
with gr.Row():
fps = gr.Number(label="FPS", value=16, precision=0)
num_frames = gr.Number(label="Num Frames", value=132, precision=0)
num_inference_steps = gr.Number(label="Num Inference Steps", value=50, precision=0)
with gr.Row():
with gr.Column():
guidance_scale = gr.Number(
label="Guidance Scale",
value=1.0,
interactive=True,
scale=4,
)
include_guidance_scale = gr.Checkbox(label="Send", value=True, scale=1, min_width=70)
with gr.Column():
guidance_scale_2 = gr.Number(
label="Guidance Scale 2",
value=None,
interactive=False,
scale=4,
)
include_guidance_scale_2 = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
with gr.Column():
true_cfg_scale = gr.Number(
label="True CFG Scale",
value=None,
interactive=False,
scale=4,
)
include_true_cfg_scale = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
# I want it to wrap after 3, but it auto wrap after 2. I'll rather fully delegate to auto wrapping.
with gr.Column():
boundary_ratio = gr.Number(
label="Boundary Ratio",
value=None,
interactive=False,
scale=4,
)
include_boundary_ratio = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
with gr.Column():
flow_shift = gr.Number(label="Flow Shift", value=None, interactive=False, scale=4)
include_flow_shift = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
with gr.Column():
seed = gr.Number(label="Seed", value=42, precision=0, interactive=True, scale=4)
include_seed = gr.Checkbox(label="Send", value=True, scale=1, min_width=70)
with gr.Row():
start_button = gr.Button("Start", variant="primary", elem_id="streaming-video-start")
hidden_stream_config = gr.Textbox(value="", visible=False)
with gr.Column(scale=1):
gr.HTML(
VIDEO_STREAM_VIEW_HTML.read_text(encoding="utf-8"),
label="Streaming Preview",
js_on_load=VIDEO_STREAM_VIEW_JS.read_text(encoding="utf-8"),
)
inputs = [
host,
port,
model,
prompt,
negative_prompt,
width,
height,
fps,
num_frames,
num_inference_steps,
include_guidance_scale,
guidance_scale,
include_guidance_scale_2,
guidance_scale_2,
include_seed,
seed,
include_boundary_ratio,
boundary_ratio,
include_flow_shift,
flow_shift,
include_true_cfg_scale,
true_cfg_scale,
]
for toggle, field in [
(include_guidance_scale, guidance_scale),
(include_guidance_scale_2, guidance_scale_2),
(include_seed, seed),
(include_boundary_ratio, boundary_ratio),
(include_flow_shift, flow_shift),
(include_true_cfg_scale, true_cfg_scale),
]:
toggle.change(
fn=_set_optional_field_enabled,
inputs=toggle,
outputs=field,
)
start_button.click(
fn=build_browser_stream_config,
inputs=inputs,
outputs=[hidden_stream_config, start_button],
).then(
fn=lambda config: config,
inputs=hidden_stream_config,
outputs=hidden_stream_config,
js="""
(config) => {
if (window.vllmStreamingVideoStart) {
window.vllmStreamingVideoStart(config);
}
return config;
}
""",
)
return demo
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Streaming video generation Gradio demo")
parser.add_argument("--host", default="127.0.0.1", help="Host/IP for gradio launch")
parser.add_argument("--port", type=int, default=7860, help="Port for gradio launch")
parser.add_argument("--share", action="store_true", help="Create a public Gradio link")
return parser.parse_args()
def main() -> None:
args = parse_args()
demo = create_demo()
demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share)
if __name__ == "__main__":
main()
@@ -0,0 +1,279 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Text-only WebSocket client for `/v1/realtime/video`.
The client prints a line for every binary video chunk it receives and writes
the received bytes to disk when the session finishes or is interrupted.
Image/reference input is intentionally not supported in this example yet.
Requirements:
pip install websockets
"""
import argparse
import asyncio
import io
import json
import os
import time
from pathlib import Path
from typing import Any
import av
try:
from websockets.asyncio.client import connect # pyright: ignore[reportMissingImports]
except ImportError:
try:
from websockets import connect # pyright: ignore[reportMissingImports]
except ImportError:
print("Please install websockets: pip install websockets")
raise SystemExit(1)
HELIOS_DISTILLED_EXTRA_PARAMS = {
"is_enable_stage2": True,
"pyramid_num_stages": 3,
"pyramid_num_inference_steps_list": [1, 1, 1],
"is_amplify_first_chunk": True,
}
def _json_object(value: str, *, argument_name: str) -> dict[str, Any]:
try:
parsed = json.loads(value)
except json.JSONDecodeError as exc:
raise argparse.ArgumentTypeError(f"{argument_name} must be valid JSON: {exc}") from exc
if not isinstance(parsed, dict):
raise argparse.ArgumentTypeError(f"{argument_name} must be a JSON object")
return parsed
def _maybe_set(payload: dict[str, Any], key: str, value: Any) -> None:
if value is not None:
payload[key] = value
def build_session_start(args: argparse.Namespace) -> dict[str, Any]:
payload: dict[str, Any] = {
"type": "session.start",
"model": args.model,
"prompt": args.prompt,
}
_maybe_set(payload, "negative_prompt", args.negative_prompt)
_maybe_set(payload, "width", args.width)
_maybe_set(payload, "height", args.height)
_maybe_set(payload, "size", args.size)
_maybe_set(payload, "seconds", args.seconds)
_maybe_set(payload, "fps", args.fps)
_maybe_set(payload, "num_frames", args.num_frames)
_maybe_set(payload, "num_inference_steps", args.num_inference_steps)
_maybe_set(payload, "guidance_scale", args.guidance_scale)
_maybe_set(payload, "guidance_scale_2", args.guidance_scale_2)
_maybe_set(payload, "boundary_ratio", args.boundary_ratio)
_maybe_set(payload, "flow_shift", args.flow_shift)
_maybe_set(payload, "true_cfg_scale", args.true_cfg_scale)
_maybe_set(payload, "seed", args.seed)
extra_params: dict[str, Any] = {}
if "Helios-Distilled" in args.model and args.helios_distilled_preset:
extra_params.update(HELIOS_DISTILLED_EXTRA_PARAMS)
if args.extra_params:
extra_params.update(args.extra_params)
if extra_params:
payload["extra_params"] = extra_params
return payload
def _count_decoded_frames(video_bytes: bytes, *, stream_format: str) -> int | None:
"""Return the number of decodable video frames in a streamed byte sequence, if known."""
if not video_bytes:
return 0
av_format = "mp4" if stream_format == "m4s" else None
try:
with av.open(io.BytesIO(video_bytes), format=av_format, mode="r") as container:
stream = container.streams.video[0]
return sum(1 for _ in container.decode(stream))
except Exception:
return None
def save_video(
chunks: list[bytes],
output_path: Path,
*,
fps: float,
stream_format: str,
) -> None:
if not chunks:
print("No video chunks received; nothing to save.")
return
from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes
streamed_bytes = b"".join(chunks)
playback_bytes = finalize_streaming_video_bytes(
streamed_bytes,
input_format="m4s",
fps=fps,
)
output_path.parent.mkdir(parents=True, exist_ok=True)
tmp_path = output_path.with_suffix(output_path.suffix + ".tmp")
with tmp_path.open("wb") as f:
f.write(playback_bytes)
os.replace(tmp_path, output_path)
print(
f"Saved {len(chunks)} streamed chunk(s), {len(streamed_bytes)} streamed bytes "
f"-> {len(playback_bytes)} playback bytes at {output_path}"
)
async def stream_video(args: argparse.Namespace) -> None:
payload = build_session_start(args)
output_path = Path(args.output)
url = f"ws://{args.host}:{args.port}/v1/realtime/video"
chunks: list[bytes] = []
total_bytes = 0
total_frames = 0
accumulated_stream = bytearray()
stream_format = "m4s"
done = False
ws = None
try:
async with connect(url, max_size=None) as websocket:
ws = websocket
await websocket.send(json.dumps(payload))
print(f"Connected: {url}")
print(f"Sent session.start: {json.dumps(payload, ensure_ascii=False)}")
started_at = time.perf_counter()
last_chunk_at = started_at
while True:
message = await websocket.recv()
now = time.perf_counter()
if isinstance(message, bytes):
chunks.append(message)
chunk_bytes = len(message)
chunk_elapsed = now - last_chunk_at
last_chunk_at = now
total_bytes += chunk_bytes
accumulated_stream.extend(message)
decoded_total = _count_decoded_frames(bytes(accumulated_stream), stream_format=stream_format)
if decoded_total is not None:
chunk_frames = max(decoded_total - total_frames, 0)
total_frames = decoded_total
else:
chunk_frames = 0
total_elapsed = now - started_at
print(
f"[chunk {len(chunks):03d}] "
f"bytes={chunk_bytes} "
f"frames={chunk_frames} "
f"elapsed={chunk_elapsed:.2f}s "
f"total_bytes={total_bytes} "
f"total_frames={total_frames} "
f"total_elapsed={total_elapsed:.2f}s"
)
continue
msg = json.loads(message)
msg_type = msg.get("type")
if msg_type == "video.start":
stream_format = msg.get("format") or stream_format
print(f"Video session started: request_id={msg.get('request_id')} format={msg.get('format')}")
elif msg_type == "session.done":
print(f"Session complete: {json.dumps(msg, ensure_ascii=False)}")
done = True
break
elif msg_type == "error":
print(f"ERROR: {msg.get('message', msg)}")
else:
print(f"Received control message: {msg}")
except (KeyboardInterrupt, asyncio.CancelledError):
print("\nInterrupted. Asking server to stop the active session...")
if ws is not None:
try:
await ws.send(json.dumps({"type": "session.stop"}))
await ws.close(code=1000, reason="client interrupted")
except Exception:
pass
raise
finally:
print(
"Saving video... (May take a while. Remuxing concatenated chunks into one progressive MP4 file with total duration metadata in file header)"
)
save_video(
chunks,
output_path,
fps=float(payload.get("fps") or args.fps),
stream_format=stream_format,
)
if not done:
print("Saved partial output because the stream did not finish normally.")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Streaming video generation WebSocket client")
parser.add_argument("--host", default="127.0.0.1", help="API server host")
parser.add_argument("--port", type=int, default=8000, help="API server port")
parser.add_argument(
"--model",
default="BestWishYsh/Helios-Distilled",
help="Model name served by the API server",
)
parser.add_argument(
"--prompt",
default="A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement.",
help="Text prompt for video generation",
)
parser.add_argument("--output", default="streaming_video_output.mp4", help="Path for the received video bytes")
parser.add_argument("--width", type=int, default=640, help="Video width")
parser.add_argument("--height", type=int, default=384, help="Video height")
parser.add_argument("--size", default=None, help="Video size as WIDTHxHEIGHT; alternative to width/height")
parser.add_argument("--seconds", default=None, help="Clip duration in seconds")
parser.add_argument("--fps", type=int, default=16, help="Frames per second")
parser.add_argument("--num-frames", type=int, default=99, help="Number of generated frames")
parser.add_argument("--negative-prompt", default=None, help="Negative prompt")
parser.add_argument("--num-inference-steps", type=int, default=None, help="Number of diffusion steps")
parser.add_argument("--guidance-scale", type=float, default=1.0, help="Classifier-free guidance scale")
parser.add_argument("--guidance-scale-2", type=float, default=None, help="High-noise CFG scale for video models")
parser.add_argument("--boundary-ratio", type=float, default=None, help="Boundary split ratio for video models")
parser.add_argument("--flow-shift", type=float, default=None, help="Scheduler flow_shift for video models")
parser.add_argument("--true-cfg-scale", type=float, default=None, help="True CFG scale for models that support it")
parser.add_argument("--seed", type=int, default=42, help="Random seed")
parser.add_argument(
"--helios-distilled-preset",
action=argparse.BooleanOptionalAction,
default=True,
help="Send the Helios-Distilled extra_params used by the bundled example.",
)
parser.add_argument(
"--extra-params",
type=lambda value: _json_object(value, argument_name="--extra-params"),
default=None,
help="JSON object merged into extra_params; overrides preset keys on conflict.",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
try:
asyncio.run(stream_video(args))
except KeyboardInterrupt:
pass
if __name__ == "__main__":
main()
@@ -0,0 +1,11 @@
<!-- This file and the corresponding JS file are needed to display a streamable video view in the Gradio demo.
Gradio's built-in video player does not support modern M4S streaming format. -->
<div id="vllm-streaming-video-view" style="display:flex; flex-direction:column; gap:10px;">
<video id="vllm-streaming-video-player" autoplay muted controls playsinline
style="width:100%; min-height:320px; background:#111; border-radius:8px;"></video>
<div style="display:flex; align-items:center; gap:10px;">
<span id="vllm-streaming-video-status" style="font-weight:600;">Ready</span>
</div>
<pre id="vllm-streaming-video-log" style="height:220px; overflow:auto; margin:0; padding:10px; border-radius:8px;
background:#0f172a; color:#dbeafe; font-size:12px; white-space:pre-wrap;">Ready.</pre>
</div>
@@ -0,0 +1,198 @@
/**
* This file and the corresponding HTML file are needed to display a streamable video view in the Gradio demo.
* Gradio's built-in video player does not support modern M4S streaming format.
*/
(function () {
const playerId = "vllm-streaming-video-player";
const statusId = "vllm-streaming-video-status";
const logId = "vllm-streaming-video-log";
let ws = null;
let mediaSource = null;
let sourceBuffer = null;
let queue = [];
let chunkCount = 0;
let totalBytes = 0;
let done = false;
let objectUrl = null;
function el(id) {
return document.getElementById(id);
}
function setStatus(text) {
const node = el(statusId);
if (node) node.textContent = text;
}
function log(line) {
const node = el(logId);
if (!node) return;
node.textContent += "\n" + line;
node.scrollTop = node.scrollHeight;
}
function setStartButton(disabled, label) {
const btn = document.querySelector("#streaming-video-start button") || document.getElementById("streaming-video-start");
if (!btn) return;
btn.disabled = disabled;
btn.textContent = label;
}
function chooseMime() {
const candidates = [
'video/mp4; codecs="avc1.42E01E"',
'video/mp4; codecs="avc1.4D401F"',
'video/mp4; codecs="avc1.64001F"',
"video/mp4"
];
for (const mime of candidates) {
if (window.MediaSource && MediaSource.isTypeSupported(mime)) return mime;
}
return null;
}
function pump() {
if (!sourceBuffer || sourceBuffer.updating || queue.length === 0) return;
try {
sourceBuffer.appendBuffer(queue.shift());
} catch (err) {
done = true;
setStatus("SourceBuffer error");
setStartButton(false, "Restart");
log("SourceBuffer append error: " + err.message);
finishStream();
}
}
function finishStream() {
if (!mediaSource || mediaSource.readyState !== "open") return;
if (sourceBuffer && sourceBuffer.updating) return;
if (queue.length > 0) return;
try {
mediaSource.endOfStream();
} catch (_) { }
}
function stopCurrent() {
if (ws && ws.readyState === WebSocket.OPEN) {
try { ws.send(JSON.stringify({ type: "session.stop" })); } catch (_) { }
try { ws.close(); } catch (_) { }
}
ws = null;
sourceBuffer = null;
mediaSource = null;
queue = [];
if (objectUrl) URL.revokeObjectURL(objectUrl);
objectUrl = null;
}
window.vllmStreamingVideoStart = function (configJson) {
stopCurrent();
chunkCount = 0;
totalBytes = 0;
done = false;
const logNode = el(logId);
if (logNode) logNode.textContent = "Starting...";
setStatus("Starting...");
let config;
try {
config = JSON.parse(configJson);
} catch (err) {
setStatus("Invalid request");
setStartButton(false, "Restart");
log("Invalid request config: " + err.message);
return configJson;
}
const mime = chooseMime();
if (!mime) {
setStatus("MSE unsupported");
setStartButton(false, "Restart");
log("This browser does not support MP4 Media Source Extensions playback.");
return configJson;
}
const video = el(playerId);
if (!video) return configJson;
mediaSource = new MediaSource();
objectUrl = URL.createObjectURL(mediaSource);
video.src = objectUrl;
mediaSource.addEventListener("sourceopen", () => {
try {
sourceBuffer = mediaSource.addSourceBuffer(mime);
sourceBuffer.mode = "segments";
sourceBuffer.addEventListener("updateend", () => {
pump();
if (done) finishStream();
});
} catch (err) {
setStatus("SourceBuffer error");
setStartButton(false, "Restart");
log("SourceBuffer error: " + err.message);
return;
}
ws = new WebSocket(config.url);
ws.binaryType = "arraybuffer";
ws.onopen = () => {
setStatus("Streaming...");
setStartButton(true, "Streaming...");
log("Connected: " + config.url);
ws.send(JSON.stringify(config.payload));
log("Sent session.start: " + JSON.stringify(config.payload));
};
ws.onmessage = (event) => {
if (typeof event.data === "string") {
const msg = JSON.parse(event.data);
if (msg.type === "video.start") {
log("Video session started: request_id=" + (msg.request_id || "") + " format=" + (msg.format || ""));
} else if (msg.type === "session.done") {
done = true;
setStatus("Done");
setStartButton(false, "Restart");
log("Session complete: " + JSON.stringify(msg));
finishStream();
} else if (msg.type === "error") {
done = true;
setStatus("Error");
setStartButton(false, "Restart");
log("ERROR: " + (msg.message || JSON.stringify(msg)));
finishStream();
} else {
log("Control message: " + JSON.stringify(msg));
}
return;
}
const data = event.data;
queue.push(data);
chunkCount += 1;
totalBytes += data.byteLength;
log(`[chunk ${String(chunkCount).padStart(3, "0")}] bytes=${data.byteLength} total_bytes=${totalBytes}`);
pump();
};
ws.onerror = () => {
done = true;
setStatus("Error");
setStartButton(false, "Restart");
log("ERROR: WebSocket error");
finishStream();
};
ws.onclose = () => {
if (!done) {
done = true;
setStatus("Closed");
setStartButton(false, "Restart");
log("WebSocket closed before session.done");
finishStream();
}
};
}, { once: true });
return configJson;
};
console.log("video-stream-view.html loaded");
})();
@@ -16,6 +16,7 @@ For the full list of supported architectures across all modalities, see
|---|---|---|---|---|---|
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | compat only, non-chunk | uploaded audio voice only; no presets | — |
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | ✓ (`ref_audio` / `speaker_embedding`) | ✓ (PCM stream) | IP labels + structured `instructions` | — |
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | — (caption-controlled) | — | caption fields (`instructions`) | — |
| MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ |
@@ -23,6 +24,7 @@ For the full list of supported architectures across all modalities, see
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
| VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
| SoulX-Singer | `Soul-AILab/SoulX-Singer` | ✓ (prompt audio) | — (batch only) | — (prompt + target audio) | — (chat client) |
CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](../../offline_inference/text_to_speech/README.md#cosyvoice3) instead.
@@ -34,7 +36,7 @@ Launch the server (defaults shown — adjust `--port`, `--gpu-memory-utilization
vllm serve <hf-repo-or-local-path> --omni --port 8091
```
Send a TTS request via curl:
Send a TTS request via curl. These generic snippets assume a model with a preset/default voice; voice-cloning-only models such as IndexTTS-2 require `ref_audio` or an uploaded audio `voice` (see model-specific sections below).
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
@@ -90,7 +92,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
```
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 24 kHz for the others).
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 22.05 kHz for IndexTTS-2, and 24 kHz for many others).
For full request-shape documentation (all parameters, response formats, error codes), see the [Speech API reference](../../../docs/serving/speech_api.md).
@@ -135,6 +137,40 @@ bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
---
## IndexTTS-2
2-stage TTS (GPT AR + S2Mel CFM DiT + BigVGAN) at 22.05 kHz. Requests use `ref_audio` for voice cloning, or an uploaded audio `voice` from `/v1/audio/voices`. Supports emotion conditioning via `emo_audio`, `emo_text`, or `emo_vector` passed in `extra_params`.
### Launch
```bash
vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092
# or, to pass the bundled deploy config explicitly:
bash examples/online_serving/text_to_speech/indextts2/run_server.sh
```
### Sending requests
```bash
# Voice cloning (ref_audio required)
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
--text "你好,世界!" \
--ref-audio /path/to/reference.wav
# With emotion audio
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
--text "今天心情很好!" \
--ref-audio /path/to/ref.wav \
--emo-audio /path/to/happy.wav
```
### Notes
- Output: 22.05 kHz mono WAV.
- Provide `ref_audio` on the documented raw request path, or pass `voice` only when it names an uploaded audio voice; IndexTTS-2 does not provide a built-in text-only preset voice.
- Emotion params (`emo_audio`, `emo_text`, `emo_vector`, `emo_alpha`, `use_emo_text`, `use_random`) are passed via the `extra_params` field. Official precedence is `use_emo_text` > `emo_vector` > `emo_audio` > same emotion as the speaker reference.
- `stream=true` is accepted as an OpenAI-compatible response path, but IndexTTS-2 is not async-chunk streaming; audio is produced after S2Mel receives the full mel-code sequence.
- Deploy config: `vllm_omni/deploy/indextts2.yaml` (auto-loaded).
---
## Fish Speech S2 Pro
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
@@ -626,3 +662,65 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
- Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
- Voice cloning is gated upstream and may require a recent `mistral_common`.
- A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now.
---
## SoulX-Singer
Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Single-stage DiT with inline preprocess. Uses the `/v1/chat/completions` endpoint with multimodal input (`prompt_audio` + `target_audio`).
### Prerequisites
Download DiT and preprocess weights, then set up separate SVS / SVC view directories and install dependencies as described in the [offline README](../../offline_inference/text_to_speech/README.md#soulx-singer). `config.json` `architectures` field is the single source of truth for SVS vs SVC — point `MODEL` at the matching directory.
### Launch
```bash
# SVS (default)
export MODEL=/path/to/SoulX-Singer
export PREPROCESS=/path/to/SoulX-Singer-Preprocess
bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh
# SVC
export MODE=svc
export MODEL=/path/to/SoulX-Singer-svc
bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh
```
Or equivalently, set `SOULX_PREPROCESS_WEIGHTS_DIR` and launch directly:
```bash
export SOULX_PREPROCESS_WEIGHTS_DIR=$PREPROCESS
vllm serve $MODEL --omni \
--deploy-config vllm_omni/deploy/soulxsinger_${MODE}.yaml \
--port 8192 --trust-remote-code --enforce-eager
```
### Sending requests
Audio paths must be reachable from the server host (local filesystem or data URL). The client sends prompt vocal via `input_audio` and target accompaniment via `extra_args['target_audio']`.
```bash
# Default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3
python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \
--prompt-audio /path/on/server/zh_prompt.mp3 \
--target-audio /path/on/server/music.mp3 \
--preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \
-o output.wav
```
Use precomputed metadata to skip online preprocess with following command:
```bash
python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \
--prompt-metadata-path /path/on/server/zh_prompt.json \
--target-metadata-path /path/on/server/music.json \
--audio-path /path/on/server/zh_prompt.mp3 \
-o output.wav
```
`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. See `openai_chat_client.py --help` for `--vocal-sep`, `--language`, `--num-inference-steps`, `--guidance-scale`, and `--seed`.
### Notes
- Output: 24 kHz mono WAV; batch only.
- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on.
- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only.
@@ -0,0 +1,25 @@
#!/bin/bash
# Launch vLLM-Omni server for IndexTTS2
#
# Usage from repository root:
# examples/online_serving/text_to_speech/indextts2/run_server.sh
# CUDA_VISIBLE_DEVICES=0 PORT=8092 MODEL=/path/to/IndexTeam/IndexTTS-2 examples/online_serving/text_to_speech/indextts2/run_server.sh
set -e
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(cd -- "$SCRIPT_DIR/../../../.." && pwd)"
MODEL="${MODEL:-IndexTeam/IndexTTS-2}"
PORT="${PORT:-8092}"
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$ROOT_DIR/vllm_omni/deploy/indextts2.yaml}"
echo "Starting IndexTTS2 server with model: $MODEL"
FLASHINFER_DISABLE_VERSION_CHECK=1 \
vllm serve "$MODEL" \
--host 0.0.0.0 \
--port "$PORT" \
--omni \
--trust-remote-code \
--deploy-config "$DEPLOY_CONFIG"
@@ -0,0 +1,126 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""OpenAI-compatible client for IndexTTS2 TTS via /v1/audio/speech endpoint.
Examples:
# With reference audio for voice cloning
python speech_client.py --text "你好,世界!" \
--ref-audio /path/to/reference.wav
# With emotion audio
python speech_client.py --text "今天心情很好!" \
--ref-audio /path/to/ref.wav \
--emo-audio /path/to/happy.wav
Server setup:
vllm serve IndexTeam/IndexTTS-2 --omni --host 0.0.0.0 --port 8092
"""
from __future__ import annotations
import argparse
import base64
import os
import httpx
DEFAULT_API_BASE = "http://localhost:8092"
DEFAULT_API_KEY = "sk-empty"
def encode_audio_to_base64(audio_path: str) -> str:
if not os.path.exists(audio_path):
raise FileNotFoundError(f"Audio file not found: {audio_path}")
ext = audio_path.lower().rsplit(".", 1)[-1]
mime = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac"}.get(ext, "audio/wav")
with open(audio_path, "rb") as f:
b64 = base64.b64encode(f.read()).decode("utf-8")
return f"data:{mime};base64,{b64}"
def main() -> None:
parser = argparse.ArgumentParser(description="IndexTTS2 OpenAI speech client")
parser.add_argument("--text", type=str, required=True)
parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio for voice cloning")
parser.add_argument("--emo-audio", type=str, default=None, help="Emotion reference audio")
parser.add_argument("--emo-text", type=str, default=None, help="Emotion description text")
parser.add_argument(
"--emo-vector",
type=float,
nargs=8,
default=None,
help="8-dim emotion vector: happy angry sad afraid disgusted melancholic surprised calm",
)
parser.add_argument("--emo-alpha", type=float, default=None, help="Emotion weight in [0, 1]")
parser.add_argument("--use-emo-text", action="store_true", help="Infer emotion vector from emo-text or text")
parser.add_argument("--use-random", action="store_true", help="Use random emotion prototypes")
parser.add_argument("--model", type=str, default="IndexTeam/IndexTTS-2")
parser.add_argument("--voice", type=str, default=None, help="Uploaded voice name to use instead of --ref-audio")
parser.add_argument("--output", type=str, default="output.wav")
parser.add_argument("--api-base", type=str, default=DEFAULT_API_BASE)
parser.add_argument("--api-key", type=str, default=DEFAULT_API_KEY)
parser.add_argument("--response-format", type=str, default="wav")
args = parser.parse_args()
if not args.ref_audio and not args.voice:
parser.error("IndexTTS2 requires --ref-audio or --voice for voice cloning")
payload: dict = {
"model": args.model,
"input": args.text,
"response_format": args.response_format,
}
if args.voice:
payload["voice"] = args.voice
if args.ref_audio:
ref = args.ref_audio
if ref.startswith(("http://", "https://", "data:")):
payload["ref_audio"] = ref
else:
payload["ref_audio"] = encode_audio_to_base64(ref)
extra_params = {}
if args.emo_audio:
emo = args.emo_audio
if emo.startswith(("http://", "https://", "data:")):
extra_params["emo_audio"] = emo
else:
extra_params["emo_audio"] = encode_audio_to_base64(emo)
if args.emo_text:
extra_params["emo_text"] = args.emo_text
if args.emo_vector is not None:
extra_params["emo_vector"] = args.emo_vector
if args.emo_alpha is not None:
extra_params["emo_alpha"] = args.emo_alpha
if args.use_emo_text:
extra_params["use_emo_text"] = True
if args.use_random:
extra_params["use_random"] = True
if extra_params:
payload["extra_params"] = extra_params
url = f"{args.api_base}/v1/audio/speech"
print(f"POST {url}")
print(f" text: {args.text}")
if args.ref_audio:
print(f" ref_audio: {args.ref_audio[:80]}...")
with httpx.Client(timeout=300) as client:
resp = client.post(
url,
json=payload,
headers={"Authorization": f"Bearer {args.api_key}"},
)
if resp.status_code != 200:
print(f"Error {resp.status_code}: {resp.text[:500]}")
return
with open(args.output, "wb") as f:
f.write(resp.content)
print(f"Saved: {args.output} ({len(resp.content):,} bytes)")
if __name__ == "__main__":
main()
@@ -0,0 +1,186 @@
#!/usr/bin/env python3
"""SoulX-Singer OpenAI-compatible chat client (SVS / SVC).
Sends prompt audio via ``input_audio`` and target accompaniment via
``extra_args['target_audio']`` (server-local path). For integrated preprocess,
also pass ``preprocess_weights_dir`` in ``extra_args``.
Usage:
python openai_chat_client.py \\
--prompt-audio /path/on/server/zh_prompt.mp3 \\
--target-audio /path/on/server/music.mp3 \\
--preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \\
-o output.wav
"""
from __future__ import annotations
import argparse
import base64
import io
import sys
from pathlib import Path
import requests
import soundfile
import torch
def _audio_to_data_url(path: Path) -> str:
with path.open("rb") as handle:
data = base64.b64encode(handle.read()).decode("ascii")
return f"data:audio/mpeg;base64,{data}"
def _save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None:
audio = audio.to(torch.float32)
peak = audio.abs().max().clamp(min=1e-8)
audio = audio / peak
path.parent.mkdir(parents=True, exist_ok=True)
soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16")
def _decode_audio_from_response(body: dict) -> tuple[torch.Tensor, int]:
for choice in body.get("choices", []):
audio_obj = choice.get("message", {}).get("audio")
if isinstance(audio_obj, dict) and audio_obj.get("data"):
data, sr = soundfile.read(
io.BytesIO(base64.b64decode(audio_obj["data"])),
dtype="float32",
always_2d=True,
)
return torch.from_numpy(data).transpose(0, 1), sr
brief = {k: v for k, v in body.items() if k != "choices"}
raise RuntimeError(f"no audio in response message.audio: {brief}")
def main() -> int:
repo_root = Path(__file__).resolve().parents[4]
default_assets = repo_root / "tests" / "assets" / "soulxsinger"
parser = argparse.ArgumentParser(description="SoulX-Singer online chat client")
parser.add_argument("--port", type=int, default=8192)
parser.add_argument("--model", default="Soul-AILab/SoulX-Singer")
parser.add_argument(
"--prompt-audio",
default=str(default_assets / "zh_prompt.mp3"),
help="Prompt vocal audio (path on server if using extra_args, or local for input_audio)",
)
parser.add_argument(
"--target-audio",
default=str(default_assets / "music.mp3"),
help="Target accompaniment path on the server (extra_args['target_audio'])",
)
parser.add_argument(
"--prompt-metadata-path",
default=None,
help="SVS precomputed prompt metadata.json",
)
parser.add_argument(
"--target-metadata-path",
default=None,
help="SVS precomputed target metadata.json",
)
parser.add_argument(
"--audio-path",
default=None,
help="SVS prompt vocal wav for precomputed metadata",
)
parser.add_argument("--preprocess-weights-dir", default=None)
parser.add_argument("--output", "-o", default="soulxsinger_out.wav")
parser.add_argument("--svc", action="store_true", help="Use SVC mode knobs")
parser.add_argument("--language", default="Mandarin")
parser.add_argument("--num-inference-steps", type=int, default=32)
parser.add_argument("--guidance-scale", type=float, default=3.0)
parser.add_argument(
"--seed",
type=int,
default=42,
help="Optional CFM seed. Omit for non-deterministic sampling.",
)
parser.add_argument(
"--auto-shift",
action=argparse.BooleanOptionalAction,
default=True,
help="Auto pitch shift (default: on, original upstream infer.sh)",
)
parser.add_argument(
"--control",
default="melody",
choices=["melody", "score"],
help="SVS control mode",
)
parser.add_argument("--vocal-sep", action="store_true")
args = parser.parse_args()
meta_paths = (args.prompt_metadata_path, args.target_metadata_path, args.audio_path)
if any(meta_paths) and not all(meta_paths):
print(
"ERROR: precomputed metadata requires --prompt-metadata-path, "
"--target-metadata-path, and --audio-path together.",
file=sys.stderr,
)
return 2
extra_args: dict = {
"vocal_sep": args.vocal_sep,
"auto_shift": args.auto_shift,
"pitch_shift": 0,
}
if all(meta_paths):
extra_args.update(
{
"prompt_metadata_path": str(Path(args.prompt_metadata_path).expanduser().resolve()),
"target_metadata_path": str(Path(args.target_metadata_path).expanduser().resolve()),
"audio_path": str(Path(args.audio_path).expanduser().resolve()),
}
)
content = [{"type": "text", "text": "soulx-singer"}]
else:
prompt_path = Path(args.prompt_audio).expanduser().resolve()
if not prompt_path.is_file():
print(f"ERROR: prompt audio not found: {prompt_path}", file=sys.stderr)
return 2
extra_args["prompt_audio"] = str(prompt_path)
extra_args["target_audio"] = str(Path(args.target_audio).expanduser().resolve())
if args.preprocess_weights_dir:
extra_args["preprocess_weights_dir"] = str(Path(args.preprocess_weights_dir).expanduser().resolve())
content = [
{"type": "text", "text": "soulx-singer"},
{
"type": "input_audio",
"input_audio": {"data": _audio_to_data_url(prompt_path), "format": "mp3"},
},
]
if not args.svc:
extra_args["language"] = args.language
extra_args["control"] = args.control
payload = {
"model": args.model,
"modalities": ["audio"],
"messages": [{"role": "user", "content": content}],
"num_inference_steps": args.num_inference_steps,
"guidance_scale": args.guidance_scale,
"extra_args": extra_args,
}
if args.seed is not None:
payload["seed"] = args.seed
print(f"POST http://localhost:{args.port}/v1/chat/completions")
response = requests.post(
f"http://localhost:{args.port}/v1/chat/completions",
headers={"Content-Type": "application/json"},
json=payload,
timeout=1800,
)
response.raise_for_status()
audio, sample_rate = _decode_audio_from_response(response.json())
_save_wav(audio, Path(args.output), sample_rate)
duration = audio.shape[-1] / sample_rate
print(f"saved {args.output} sr={sample_rate}Hz duration={duration:.2f}s")
return 0
if __name__ == "__main__":
raise SystemExit(main())
@@ -0,0 +1,39 @@
#!/bin/bash
# Launch vLLM-Omni server for SoulX-Singer (single-stage DiT, preprocess inline).
#
# Usage:
# MODEL=/path/to/SoulX-Singer PREPROCESS=/path/to/Preprocess \
#
# Audio paths in client extra_args must be readable on the server host.
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)"
MODEL="${MODEL:-Soul-AILab/SoulX-Singer}"
MODE="${MODE:-svs}"
PORT="${PORT:-8192}"
GPUS="${GPUS:-0}"
if [[ "$MODE" == "svc" ]]; then
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svc.yaml}"
else
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svs.yaml}"
fi
echo "Starting SoulX-Singer server"
echo " MODEL=$MODEL"
echo " MODE=$MODE"
echo " PORT=$PORT"
echo " DEPLOY_CONFIG=$DEPLOY_CONFIG"
echo " CUDA_VISIBLE_DEVICES=$GPUS"
CUDA_VISIBLE_DEVICES="$GPUS" \
vllm serve "$MODEL" \
--omni \
--deploy-config "$DEPLOY_CONFIG" \
--host 0.0.0.0 \
--port "$PORT" \
--trust-remote-code \
--enforce-eager
+7
View File
@@ -40,6 +40,7 @@ dev = [
"pytest-asyncio>=0.21.0",
"pytest-cov>=4.0.0",
"pytest-mock>=3.10.0",
"websockets>=15.0.1",
"datasets>=2.14.0",
"mypy==1.11.1",
"pre-commit==4.0.1",
@@ -118,6 +119,12 @@ docs = [
"pydantic",
]
soulx-svs = [
"funasr>=1.0.0",
"g2pM>=0.1.2.5",
"g2p-en>=2.1.0",
"ToJyutping>=3.2.0",
]
[project.urls]
Homepage = "https://github.com/vllm-project/vllm-omni"
+139
View File
@@ -0,0 +1,139 @@
# IndexTTS-2 for voice-cloned TTS on 1x GPU
## Summary
- Vendor: IndexTeam
- Model: `IndexTeam/IndexTTS-2`
- Task: Text-to-speech with voice cloning and optional emotion control
- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API
- Maintainer: Community
## When to use this recipe
Use this recipe to serve IndexTTS-2 as a two-stage TTS system for
voice-cloned speech synthesis. Each request provides synthesis text plus either
a reference audio clip or an uploaded audio voice; the model returns 22.05 kHz
mono speech. IndexTTS-2 also supports optional emotion conditioning through
model-specific `extra_params`.
IndexTTS-2 is **not an async-chunk streaming model** in the bundled deployment
configuration: the S2Mel flow-matching stage consumes the full mel-code sequence
from the AR talker before it produces audio. The OpenAI-compatible endpoint can
still be called with `stream=true` as an HTTP compatibility path, but that should
not be interpreted as incremental model streaming.
## References
- Upstream model:
[IndexTeam/IndexTTS-2 on Hugging Face](https://huggingface.co/IndexTeam/IndexTTS-2)
- Online serving example:
[`examples/online_serving/text_to_speech/README.md#indextts-2`](../../examples/online_serving/text_to_speech/README.md#indextts-2)
- Offline inference example:
[`examples/offline_inference/text_to_speech/README.md#indextts-2`](../../examples/offline_inference/text_to_speech/README.md#indextts-2)
- OpenAI-compatible client:
[`examples/online_serving/text_to_speech/indextts2/speech_client.py`](../../examples/online_serving/text_to_speech/indextts2/speech_client.py)
- Deploy config:
[`vllm_omni/deploy/indextts2.yaml`](../../vllm_omni/deploy/indextts2.yaml)
## Hardware Support
### GPU
### 1x L4 24GB or larger CUDA GPU
#### Environment
- OS: Linux
- Python: 3.10+
- Driver / runtime: NVIDIA CUDA environment
- vLLM version: Match the repository requirements for your checkout
- vLLM-Omni version or commit: Use the commit you are deploying from
#### Command
Start the server from the repository root:
```bash
vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092
```
The bundled deployment config is selected automatically for the `indextts2`
model type. To make the config choice explicit, pass it directly:
```bash
vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092 \
--deploy-config vllm_omni/deploy/indextts2.yaml
```
#### Verification
Run the bundled client. Local audio paths passed to this client are converted to
base64 data URLs before they are sent to the HTTP API:
```bash
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
--api-base http://localhost:8092 \
--text "你好,这是 IndexTTS-2 语音合成测试。" \
--ref-audio /path/to/reference.wav \
--output indextts2.wav
```
For a raw HTTP request, pass `ref_audio` as an HTTP(S) URL, a base64 data URL,
or a `file://` URI allowed by the server's `--allowed-local-media-path`. Do not
pass an arbitrary local filesystem path directly in JSON unless your client
converts it first:
```bash
curl http://localhost:8092/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"model": "IndexTeam/IndexTTS-2",
"input": "你好,这是 IndexTTS-2 语音合成测试。",
"response_format": "wav",
"ref_audio": "data:audio/wav;base64,<BASE64_ENCODED_AUDIO>"
}' \
--output indextts2.wav
```
Optional emotion conditioning is passed through `extra_params`. Use one emotion
source mode at a time:
```bash
curl http://localhost:8092/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"model": "IndexTeam/IndexTTS-2",
"input": "今天心情很好!",
"response_format": "wav",
"ref_audio": "data:audio/wav;base64,<BASE64_ENCODED_REFERENCE_AUDIO>",
"extra_params": {
"emo_vector": [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
"emo_alpha": 0.8
}
}' \
--output indextts2_happy.wav
```
#### Notes
- Hardware scope: the default config is a single-GPU CUDA deployment. Stage 0
and Stage 1 share GPU 0 by default.
- Audio output: 22.05 kHz mono WAV.
- Voice cloning: `ref_audio` is required for the documented raw request path.
Alternatively, pass `voice` only when it names an uploaded audio voice from
`/v1/audio/voices`; IndexTTS-2 does not provide a built-in text-only preset
voice in this serving recipe.
- Reference audio format: use an HTTP(S) URL, a base64 data URL, or a `file://`
URI allowed by the server's `--allowed-local-media-path` for raw API calls.
The bundled Python client accepts local paths and converts them to base64 data
URLs.
- Emotion control: `extra_params.emo_audio`, `extra_params.emo_vector`, and
`extra_params.use_emo_text` are alternative conditioning modes. Official
precedence is `use_emo_text` > `emo_vector` > `emo_audio` > same emotion as
the speaker reference. The `emo_vector` order is `[happy, angry, sad, afraid,
disgusted, melancholic, surprised, calm]`.
- Streaming limitation: `async_chunk` is disabled in
`vllm_omni/deploy/indextts2.yaml`; S2Mel requires the full mel-code sequence
before audio generation.
- First request can be slower because reference encoders and auxiliary models
are lazily loaded and cached.
+34 -5
View File
@@ -5,7 +5,7 @@
## Summary
- Vendor: JD (Joy Future Academy)
- Model: [`jdopensource/JoyAI-VL-Interaction-Preview`](https://huggingface.co/jdopensource/JoyAI-VL-Interaction-Preview) (8B, vanilla Qwen3-VL weights)
- Model: [`jdopensource/JoyAI-VL-Interaction-Preview`](https://huggingface.co/jdopensource/JoyAI-VL-Interaction-Preview) (8B, Qwen3-VL architecture with weights retrained for streaming interaction)
- Task: Per-tick proactive interaction over a live video stream — the model decides
on its own each second to speak, stay silent, or delegate a hard question
- Mode: Online serving — an OpenAI-compatible interaction orchestrator in front of
@@ -25,7 +25,8 @@ framework internals and how to add another full-duplex model, see
- OS: Linux
- Python: 3.10+
- Hardware: 1x GPU. The default `T_s=100` frame window + `--max-model-len 131072` wants
≈48GB+; for ~24GB, lower `--chunk-frames`, `--max-model-len`, and the image limit together
≈48GB+; for ~24GB, lower `--chunk-frames`, `--max-model-len`, and the image limit together,
and/or load the weights in fp8 (see "Smaller GPU: fp8 weights" below — ~9.9 GiB vs 16.8 GiB)
- vLLM / vLLM-Omni: versions from your current checkout
## Start server
@@ -33,7 +34,7 @@ framework internals and how to add another full-duplex model, see
From repository root:
```bash
# 1. Serve the model (plain `vllm serve`, NOT --omni; it is vanilla Qwen3-VL).
# 1. Serve the model (plain `vllm serve`, NOT --omni; it uses the Qwen3-VL architecture).
# The image limit must cover the short-term frame window (chunk_frames, default 100 = T_s);
# prefix caching keeps the accumulating window cheap. Lower both for smaller GPUs.
vllm serve jdopensource/JoyAI-VL-Interaction-Preview \
@@ -46,6 +47,26 @@ python -m vllm_omni.experimental.fullduplex.joyvl.serving.server --port 8070 \
--main-backend-url http://127.0.0.1:8061/v1 --main-model JoyAI-VL-Interaction-Preview
```
### Smaller GPU: fp8 weights
Add `--quantization fp8` to the step-1 `vllm serve` to load the weights in fp8 online
from the bf16 checkpoint (no separate quantized checkpoint needed; uses vLLM's generic
fp8 path, since the model is a standard Qwen3-VL VLM):
```bash
vllm serve jdopensource/JoyAI-VL-Interaction-Preview \
--served-model-name JoyAI-VL-Interaction-Preview --port 8061 \
--quantization fp8 --max-model-len 131072 --enable-prefix-caching \
--limit-mm-per-prompt '{"image":256,"video":1}'
```
Measured here: model weights **9.9 GiB** vs **16.8 GiB** bf16 (**−41%**), which lets the 8B
fit a much smaller card. This is a **memory** option, not a speed one: at one request per
tick (the streaming regime) fp8 was **not** faster (≈19% slower in our test) because
single-stream decode is memory-latency-bound and pays fp8's dynamic-scale overhead; its
throughput gain needs batching. Decisions stay coherent, but fp8 can shift the
speak/silence boundary, so for timing-critical alerting validate output against bf16 first.
Optional one-shot launch (model + orchestrator + JD webui + ASR/TTS/background,
all env-configurable). The JD webui frontend is external — set `WEBUI_DIR` to point at it:
@@ -176,14 +197,22 @@ the audio-track caveat.
## Notes
- `--omni` is **not** used: the model is standard Qwen3-VL, so stock `vllm serve`
runs the forward pass; this recipe only adds the interaction/serving layer.
- `--omni` is **not** used: the model keeps the Qwen3-VL architecture (only the
weights are retrained), so stock `vllm serve` runs the forward pass; this recipe
only adds the interaction/serving layer.
- On a host without `nvcc` / `ninja`, `vllm serve` of the 8B can crash engine-core in the
FlashInfer sampler JIT (`FileNotFoundError: 'ninja'`) during `profile_run`. Set
`VLLM_USE_FLASHINFER_SAMPLER=0` (or install `ninja`) to work around it.
- `force_silence_before_query` is on by default — the model stays silent until an
instruction arrives; give a standing task (e.g. "translate the on-screen text")
to arm proactive output.
- **Frame resolution is the main latency knob.** The system/memory prefix is prefix-cached,
so per-tick *new* compute is dominated by the new frame's vision tokens. Measured on the
8B: ~256×192 frames hit Qwen3-VL's min-pixel floor (~72 vision tokens, ~17 ms/tick) vs
~302 tokens / ~38 ms at 640×480 — about 2× cheaper. Downsample frames to ~256×192 for the
tightest latency / highest concurrency; one GPU then sustains ~150–180 concurrent 1 fps
streams with p95 < 200 ms. Per-tick latency is already far inside the 1 fps budget, so
serving is rarely the bottleneck — resolution is the lever if you need more headroom.
- Speech is external and pluggable: point `ASR_URL` / `TTS_URL` at the bridges in
`examples/online_serving/joyvl_interaction/bridges/` or any compatible service.
- The decision prompts, sampling, and 3-tier summary memory (`T_s=100`, mid→long every
+3
View File
@@ -43,6 +43,7 @@ recipes/
| [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
| [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
| [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
| [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
@@ -57,10 +58,12 @@ recipes/
| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB |
| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) |
| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB |
| [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
| [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
| [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
| [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) |
| [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
| [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
+170
View File
@@ -0,0 +1,170 @@
# Stable-Audio-Open Text-To-Audio Generation on 1x GPU
> Text-to-audio recipe for Stable Audio Open with offline inference and
> OpenAI-compatible online serving on 1x RTX 4090 24GB.
## Summary
- Vendor: Stability AI
- Model: `stabilityai/stable-audio-open-1.0`
- Task: Text-to-audio generation
- Mode: Offline inference and online serving
- Maintainer: Community
## When to use this recipe
Use this recipe when you want to run Stable Audio Open on a single RTX 4090
24GB GPU for music or sound-effect generation. The recipe covers a 10-second
offline validation sample with TeaCache and online serving through the
`/v1/audio/generate` endpoint.
## References
- Model: <https://huggingface.co/stabilityai/stable-audio-open-1.0>
- Offline example:
[`examples/offline_inference/text_to_audio`](../../examples/offline_inference/text_to_audio)
- Online example:
[`examples/online_serving/stable_audio`](../../examples/online_serving/stable_audio)
- Audio generation API:
[`docs/serving/audio_generate_api.md`](../../docs/serving/audio_generate_api.md)
- Related issue:
[#2645](https://github.com/vllm-project/vllm-omni/issues/2645)
## Hardware Support
### GPU
#### 1x RTX 4090 24GB
##### Environment
- OS: Ubuntu 22.04.5
- Python: 3.12
- GPU: NVIDIA GeForce RTX 4090, 24564 MiB VRAM
- Driver / runtime: NVIDIA driver 595.80, CUDA-capable runtime matching the
repository build
- vLLM version: 0.22.0
- vLLM-Omni version: source checkout
- PyTorch: 2.11.0+cu130
- Model path used in the commands below: `/path/to/stable-audio-open-1.0`
Stable Audio Open is a gated Hugging Face model. Accept the model license on
the Hugging Face model card before downloading the checkpoint.
```bash
hf auth login
hf download stabilityai/stable-audio-open-1.0 \
--local-dir /path/to/stable-audio-open-1.0
```
##### Commands
Run a 10-second offline validation sample from the repository root:
```bash
python examples/offline_inference/text_to_audio/text_to_audio.py \
--model /path/to/stable-audio-open-1.0 \
--prompt "A gentle piano melody with soft room ambience" \
--negative-prompt "Low quality, distorted, noisy" \
--seed 42 \
--guidance-scale 7.0 \
--audio-length 10.0 \
--num-inference-steps 50 \
--cache-backend tea_cache \
--output examples/offline_inference/text_to_audio/stable_audio_10s.wav
```
Start the online serving endpoint:
```bash
vllm-omni serve /path/to/stable-audio-open-1.0 \
--host 0.0.0.0 \
--port 8091 \
--gpu-memory-utilization 0.9 \
--trust-remote-code \
--enforce-eager \
--omni
```
Generate a 10-second WAV file from the repository root in another terminal:
```bash
curl http://localhost:8091/health
curl -X POST http://localhost:8091/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "A gentle piano melody with soft room ambience",
"audio_length": 10.0,
"num_inference_steps": 50,
"guidance_scale": 7.0,
"negative_prompt": "Low quality, distorted, noisy",
"seed": 42,
"response_format": "wav"
}' \
--output examples/online_serving/stable_audio/piano_10s.wav
```
##### Verification
Check that:
- The offline command writes a valid WAV file.
- The server responds on `http://localhost:8091/health`.
- The online request writes a valid WAV file.
- The generated audio sample rate is 44.1 kHz.
- The generated duration is approximately 10 seconds.
- Peak sampled GPU memory is within the 24GB RTX 4090 budget. In the validated
run, offline and online generation each peaked at about 12.6 GiB.
Validate the offline outputs:
```bash
ls -lh examples/offline_inference/text_to_audio/stable_audio_10s.wav
python - <<'PY'
import soundfile as sf
path = "examples/offline_inference/text_to_audio/stable_audio_10s.wav"
audio, sample_rate = sf.read(path)
print("sample_rate:", sample_rate)
print("shape:", audio.shape)
print("duration:", len(audio) / sample_rate)
PY
```
Validate the online output:
```bash
ls -lh examples/online_serving/stable_audio/piano_10s.wav
python - <<'PY'
import soundfile as sf
path = "examples/online_serving/stable_audio/piano_10s.wav"
audio, sample_rate = sf.read(path)
print("sample_rate:", sample_rate)
print("shape:", audio.shape)
print("duration:", len(audio) / sample_rate)
PY
```
##### Notes
- `stable-audio-open-1.0` can generate up to about 47 seconds of 44.1 kHz
stereo audio. This recipe validates 10-second WAV outputs.
- `--cache-backend tea_cache` is supported and was used for the 10-second
offline validation command.
- The model is gated on Hugging Face and requires license acceptance before
download.
- If online serving fails while importing `torchaudio`, make sure the
`torchaudio` wheel matches the installed PyTorch and CUDA build. The
validated environment used `torch==2.11.0+cu130` and
`torchaudio==2.11.0+cu130`.
- The `NIXL is not available`, `GLOO_SOCKET_IFNAME`, and `torchsde` boundary
warnings observed during validation did not prevent successful generation.
- This recipe was validated on one RTX 4090 24GB GPU only. Other GPU counts,
ROCm, XPU, and NPU setups are not covered here.
- Long generations, higher inference-step counts, and non-WAV response formats
were not benchmarked in this recipe.
+147
View File
@@ -0,0 +1,147 @@
# Wan2.1 VACE Unified Video Generation
## Summary
- Vendor: Wan-AI
- Models: `Wan2.1-VACE-1.3B-diffusers`, `Wan2.1-VACE-14B-diffusers`
- Tasks: T2V, I2V, V2LF, FLF2V, inpainting, and R2V
- Mode: offline inference through the shared task examples
The shared examples construct VACE's pipeline-native conditioning data from the
provided media inputs. They do not require a mode parameter or a model-specific
example script.
## References
- [VACE project](https://github.com/ali-vilab/VACE)
- [Hugging Face 14B checkpoint](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers)
- [ModelScope 14B checkpoint](https://modelscope.cn/models/Wan-AI/Wan2.1-VACE-14B-diffusers)
- [Shared text-to-video example](../../examples/offline_inference/text_to_video/text_to_video.md)
- [Shared image-to-video example](../../examples/offline_inference/image_to_video/README.md)
- [Migration RFC #4539](https://github.com/vllm-project/vllm-omni/issues/4539)
## Hardware Support
The 1.3B checkpoint was validated on an RTX 5090; the 81-frame T2V command below
used approximately 20 GiB of peak reserved GPU memory with VAE tiling. The 14B
checkpoint requires approximately 70 GiB of disk space and was validated on a
single 48 GiB L40S with VAE tiling and layerwise offload. Start with the smoke
configuration before attempting an 81-frame run.
## Setup
The public examples use the Hugging Face model ID by default:
```bash
MODEL="${MODEL:-Wan-AI/Wan2.1-VACE-14B-diffusers}"
ASSET_DIR="${ASSET_DIR:-$HOME/datasets/vace}"
OUT="${OUT:-$PWD/vace-results/14b}"
mkdir -p "$ASSET_DIR" "$OUT"
```
For environments where Hugging Face downloads are unavailable, download the
same checkpoint from ModelScope and point `MODEL` at the local directory:
```bash
MODEL="$HOME/models/Wan2.1-VACE-14B-diffusers"
modelscope download --model Wan-AI/Wan2.1-VACE-14B-diffusers --local_dir "$MODEL"
```
Download the public conditioning images and create a mask matching the output
size. White pixels are regenerated and black pixels are preserved.
```bash
wget -O "$ASSET_DIR/astronaut.jpg" \
https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg
wget -O "$ASSET_DIR/vace_first_frame.png" \
https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_first_frame.png
wget -O "$ASSET_DIR/vace_last_frame.png" \
https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_last_frame.png
python - "$ASSET_DIR/vace_center_mask.png" <<'PY'
import sys
from PIL import Image
mask = Image.new("L", (832, 480), 0)
mask.paste(255, (336, 0, 496, 480))
mask.save(sys.argv[1])
PY
ASTRONAUT="$ASSET_DIR/astronaut.jpg"
FIRST="$ASSET_DIR/vace_first_frame.png"
LAST="$ASSET_DIR/vace_last_frame.png"
MASK="$ASSET_DIR/vace_center_mask.png"
```
## Commands
The commands below reproduce the original VACE examples with `81 frames / 30
steps`. Most tasks use `480x832`; FLF2V uses `512x512`. For a smoke test,
replace the frame and step counts with `5 frames / 2 steps`.
```bash
# T2V
python examples/offline_inference/text_to_video/text_to_video.py \
--model "$MODEL" \
--prompt "A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves." \
--seed 0 \
--height 480 --width 832 --num-frames 81 --num-inference-steps 30 \
--guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \
--output "$OUT/t2v.mp4"
# I2V
python examples/offline_inference/image_to_video/image_to_video.py \
--model "$MODEL" --image "$ASTRONAUT" \
--prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \
--seed 42 \
--height 480 --width 832 --num-frames 81 --num-inference-steps 30 \
--guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \
--output "$OUT/i2v.mp4"
# V2LF
python examples/offline_inference/image_to_video/image_to_video.py \
--model "$MODEL" --last-image "$ASTRONAUT" \
--prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \
--seed 42 \
--height 480 --width 832 --num-frames 81 --num-inference-steps 30 \
--guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \
--output "$OUT/v2lf.mp4"
# FLF2V
python examples/offline_inference/image_to_video/image_to_video.py \
--model "$MODEL" --image "$FIRST" --last-image "$LAST" \
--prompt "CG animation style, a small blue bird takes off from a branch and lands on another branch" \
--seed 42 --height 512 --width 512 --num-frames 81 --num-inference-steps 30 \
--guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \
--output "$OUT/flf2v.mp4"
# Inpaint
python examples/offline_inference/image_to_video/image_to_video.py \
--model "$MODEL" --image "$ASTRONAUT" --mask-image "$MASK" \
--prompt "Shrek, the ogre, walks out of a building in a happy mood" \
--seed 42 \
--height 480 --width 832 --num-frames 81 --num-inference-steps 30 \
--guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \
--output "$OUT/inpaint.mp4"
# R2V; repeat --reference-image to add references
python examples/offline_inference/image_to_video/image_to_video.py \
--model "$MODEL" --reference-image "$ASTRONAUT" \
--prompt "Camera slowly zooms out from the character walking in a garden" \
--seed 42 \
--height 480 --width 832 --num-frames 81 --num-inference-steps 30 \
--guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \
--output "$OUT/r2v.mp4"
```
## Verification
```bash
for video in "$OUT"/*.mp4; do
ffprobe -v error -select_streams v:0 \
-show_entries stream=codec_name,width,height,nb_frames,r_frame_rate,duration \
-of default=noprint_wrappers=1 "$video"
done
```
For the 1.3B checkpoint, `--enable-layerwise-offload` is generally unnecessary.
+42 -44
View File
@@ -367,60 +367,58 @@ vllm serve nvidia/Cosmos3-Nano-Policy-DROID \
#### Command
```python
# cosmos3_offline.py — run with: python cosmos3_offline.py
import torch
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
Cosmos3 runs through the standard task examples; pass model-specific knobs via
`--extra-body`. Guardrails are on by default — pass `"guardrails": false` for a
quick local run (install `cosmos-guardrail` + accept the gated repo to enable them).
```bash
# Text-to-image -> examples/offline_inference/text_to_image
python examples/offline_inference/text_to_image/text_to_image.py \
--model nvidia/Cosmos3-Nano \
--prompt "A photorealistic red sports car at golden hour, cinematic lighting." \
--negative-prompt "blurry, distorted, low quality" \
--height 1024 --width 1024 --num-inference-steps 50 --guidance-scale 7.0 \
--extra-body '{"flow_shift": 3.0, "guardrails": false}' \
--output cosmos3_t2i.png
def main():
omni = Omni(
model="nvidia/Cosmos3-Nano",
model_class_name="Cosmos3OmniDiffusersPipeline",
trust_remote_code=True,
enforce_eager=True,
# Guardrails are disabled here for a quick local run; install
# cosmos-guardrail + gated-repo access and drop this to enable them.
model_config={"guardrails": False},
)
gen = torch.Generator(device="cpu").manual_seed(42)
# Text-to-video -> examples/offline_inference/text_to_video
python examples/offline_inference/text_to_video/text_to_video.py \
--model nvidia/Cosmos3-Nano \
--prompt "A robot arm is cleaning a plate in the kitchen." \
--negative-prompt "blurry, distorted, low quality, jittery, deformed" \
--height 720 --width 1280 --num-frames 189 --fps 24 \
--num-inference-steps 35 --guidance-scale 6.0 \
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false,
"use_resolution_template": false, "use_duration_template": false}' \
--output cosmos3_t2v.mp4
# Text-to-image (modalities=["image"]). For T2V use modalities=["video"]
# plus num_frames/fps; for I2V add multi_modal_data={"image": <PIL.Image>}.
outputs = omni.generate(
{
"prompt": "A photorealistic red sports car at golden hour, cinematic lighting.",
"negative_prompt": "blurry, distorted, low quality",
"modalities": ["image"],
},
OmniDiffusionSamplingParams(
height=1024, width=1024, generator=gen,
guidance_scale=7.0, num_inference_steps=50, num_outputs_per_prompt=1,
),
)
outputs[0].request_output.images[0].save("cosmos3_t2i.png")
omni.close()
if __name__ == "__main__":
main()
# Image-to-video -> examples/offline_inference/image_to_video
# (Cosmos3 bundles example frames under assets/; any RGB image works too.)
python examples/offline_inference/image_to_video/image_to_video.py \
--model nvidia/Cosmos3-Nano \
--image /path/to/Cosmos3-Nano/assets/example_i2v_input.jpg \
--prompt "The scene comes to life with smooth, natural motion." \
--height 720 --width 1280 --num-frames 189 --fps 24 \
--num-inference-steps 35 --guidance-scale 6.0 \
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}' \
--output cosmos3_i2v.mp4
```
#### Verification
```bash
python cosmos3_offline.py
python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image', im.size, im.mode)"
ffprobe -v error -show_entries stream=codec_type,nb_frames,width,height cosmos3_t2v.mp4
```
#### Notes
- Same `Cosmos3OmniDiffusersPipeline` as online; mode is chosen by
`prompt["modalities"]` (`["image"]` → T2I, `["video"]` → T2V) plus
`num_frames`/`fps`, `multi_modal_data={"image": ...}` for I2V, and
`multi_modal_data={"video": [<PIL frames>]}` or a video tensor/array for V2V.
For video, frames are returned in `outputs[0].request_output.images` as an
`(B, F, H, W, 3)` array.
- The offline entry must be guarded by `if __name__ == "__main__":` — the engine
spawns workers with the `spawn` start method.
- A single `Cosmos3OmniDiffusersPipeline` serves every mode; the standard examples
select it automatically from `model_index.json`. T2I is chosen by the
`text_to_image` prompt builder (which marks `modalities=["image"]`); `text_to_video`
defaults to T2V; `image_to_video` adds `multi_modal_data={"image": ...}` (I2V).
V2V is served online (`/v1/videos/sync`).
- Model-specific knobs (`flow_shift`, `max_sequence_length`, `condition_*`,
`generate_sound`/`sound_duration`, `guardrails`, `action_*`, ...) are declared
once in `vllm_omni/model_extras/cosmos3.py` and forwarded through `--extra-body`;
unknown keys for the model are dropped.
Binary file not shown.
+16
View File
@@ -0,0 +1,16 @@
[
{
"index": "vocal_240_51240",
"language": "Mandarin",
"time": [
240,
51240
],
"duration": "0.21 0.18 0.34 0.38 0.26 0.22 0.50 0.20 0.33 0.13 0.60 0.22 0.38 0.52 0.30 0.12 1.82 0.98 0.20 0.26 0.38 0.36 0.58 0.54 0.26 0.50 1.46 3.03 0.24 0.24 0.29 0.25 0.20 0.14 0.80 0.22 0.54 0.30 0.24 0.16 0.58 0.28 0.38 1.73 0.79 0.24 0.30 0.34 0.30 0.30 0.34 0.21 0.21 0.36 0.34 0.23 1.85 1.90 0.23 0.39 0.68 0.50 0.31 0.43 0.76 0.38 2.00 1.87 0.68 0.72 0.56 0.62 0.80 0.40 0.42 1.68 1.79 0.70 0.66 0.54 0.24 0.48 0.68 0.40 2.34 0.14",
"text": "<SP> 只 是 因 为 为 在 人 群 群 中 多 看 了 你 一 眼 <SP> 再 也 没 能 忘 掉 你 容 颜 <SP> 梦 想 着 着 偶 偶 然 然 有 <SP> 一 一 天 再 相 见 <SP> 从 此 我 开 始 始 孤 孤 单 思 念 念 <SP> 想 想 你 时 你 你 在 天 边 <SP> 想 你 时 你 在 眼 前 前 <SP> 想 你 时 你 你 在 脑 海 <SP>",
"phoneme": "<SP> zh_zhi3 zh_shi4 zh_yin1 zh_wei4 zh_wei2 zh_zai4 zh_ren2 zh_qun2 zh_qun2 zh_zhong1 zh_duo1 zh_kan4 zh_le5 zh_ni3 zh_yi1 zh_yan3 <SP> zh_zai4 zh_ye3 zh_mei2 zh_neng2 zh_wang4 zh_diao4 zh_ni3 zh_rong2 zh_yan2 <SP> zh_meng4 zh_xiang3 zh_zhe5 zh_zhe5 zh_ou3 zh_ou3 zh_ran2 zh_ran2 zh_you3 <SP> zh_yi1 zh_yi1 zh_tian1 zh_zai4 zh_xiang1 zh_jian4 <SP> zh_cong2 zh_ci3 zh_wo3 zh_kai1 zh_shi3 zh_shi3 zh_gu1 zh_gu1 zh_dan1 zh_si1 zh_nian4 zh_nian4 <SP> zh_xiang3 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_tian1 zh_bian1 <SP> zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_zai4 zh_yan3 zh_qian2 zh_qian2 <SP> zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_nao3 zh_hai3 <SP>",
"note_pitch": "0 64 64 64 66 68 66 66 66 64 64 64 64 66 64 60 61 0 63 63 63 64 66 63 61 59 56 0 68 66 66 68 68 66 66 64 64 0 64 66 61 61 61 64 0 62 63 63 64 64 66 65 66 61 58 59 56 0 69 71 66 68 68 71 66 64 61 0 66 61 68 66 64 63 61 59 0 71 66 68 68 71 66 64 61 0",
"note_type": "1 2 2 2 2 3 2 2 2 3 2 2 2 2 2 2 2 1 2 2 2 2 2 2 2 2 2 1 2 2 2 3 2 3 2 3 2 1 2 3 2 2 2 2 1 2 2 2 2 2 3 2 3 2 2 2 3 1 2 3 2 2 2 3 2 2 2 1 2 2 2 2 2 2 2 3 1 2 2 2 2 3 2 2 2 1",
"f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 318.9 332.2 331.2 323.9 312.3 0.0 0.0 0.0 0.0 0.0 0.0 340.9 342.3 340.2 337.4 333.5 329.3 327.8 328.3 329.1 330.5 331.9 331.5 329.2 327.6 327.4 327.7 329.6 330.5 329.5 328.1 328.5 330.1 330.4 329.5 328.9 331.0 333.6 332.8 331.7 330.2 330.6 330.9 329.5 327.6 332.4 343.5 363.9 368.7 369.3 368.2 366.7 371.3 382.0 392.9 402.3 409.8 413.8 416.3 416.4 415.4 414.1 413.4 415.3 417.4 418.3 413.6 416.0 0.0 0.0 0.0 0.0 368.2 362.5 361.3 362.0 362.9 364.3 366.4 367.5 368.4 368.9 368.1 367.6 369.7 371.1 371.1 369.5 368.0 367.3 366.6 367.1 369.5 372.0 371.3 368.2 368.1 369.6 369.9 373.4 376.5 360.2 285.6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 384.0 376.7 372.0 371.7 373.9 376.3 375.5 361.2 333.0 317.8 319.0 329.1 333.5 333.2 333.2 333.3 328.0 314.5 0.0 0.0 0.0 320.6 326.5 332.9 335.7 334.1 329.9 327.4 326.3 326.6 329.1 330.5 330.0 328.0 326.7 328.6 331.1 332.4 332.1 332.8 333.2 333.1 331.8 328.6 318.9 293.2 323.6 327.5 330.0 332.8 331.1 319.4 272.2 0.0 0.0 0.0 0.0 278.5 294.5 310.0 318.5 322.7 327.4 334.1 337.2 336.3 328.5 321.0 324.2 341.5 362.2 374.7 375.8 370.5 366.0 365.2 368.7 371.7 374.4 373.0 368.7 370.4 374.8 375.1 372.5 368.7 363.0 357.1 359.0 366.8 377.1 379.5 371.5 359.9 351.0 358.4 371.7 377.9 375.8 367.8 359.7 363.6 375.3 379.9 377.5 372.9 359.8 345.2 337.1 334.5 333.4 332.8 329.6 319.2 292.8 261.7 253.4 262.8 273.0 278.1 279.3 278.6 277.6 277.9 278.3 278.0 277.3 276.7 275.6 275.3 276.1 277.8 278.2 277.7 277.6 277.2 276.4 275.5 273.9 271.8 270.7 274.7 282.7 285.5 281.4 273.1 264.4 262.2 268.9 278.4 284.3 283.6 277.3 268.5 262.2 263.2 270.7 278.8 285.0 287.3 282.0 272.1 265.2 266.0 272.1 278.2 283.8 285.1 281.4 272.6 267.2 269.1 276.1 282.7 289.3 290.5 285.6 273.9 267.5 267.4 271.8 277.7 282.6 283.1 277.8 269.8 265.6 268.6 273.8 281.1 286.2 287.2 282.3 270.4 264.5 263.3 263.0 268.5 268.3 268.7 270.9 277.2 277.4 284.3 285.7 282.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 311.5 306.4 306.0 306.8 305.8 305.0 307.2 310.8 312.4 313.5 314.4 314.6 313.7 312.2 311.1 311.1 312.8 314.6 312.7 306.5 299.7 302.7 308.7 315.2 314.2 310.9 309.4 309.8 310.6 310.7 310.8 311.4 312.2 313.0 312.5 312.2 312.2 311.9 312.0 312.2 311.4 306.1 303.3 309.3 315.0 325.3 324.4 323.9 325.0 326.4 327.7 327.7 328.1 328.6 327.8 328.3 328.3 327.6 328.6 336.1 347.2 359.4 371.0 374.9 371.3 366.5 369.5 375.5 376.2 370.2 365.7 367.3 372.9 376.9 375.3 368.4 360.1 358.0 365.4 380.8 383.5 380.4 374.4 367.6 365.9 372.0 376.3 377.9 375.7 372.6 368.6 361.2 353.7 0.0 309.4 302.3 299.2 300.6 305.7 308.5 308.6 309.0 310.9 310.9 309.5 308.6 309.0 311.3 313.0 314.4 314.5 313.0 312.1 311.2 308.0 299.5 295.8 295.0 285.7 278.0 280.1 281.2 279.9 278.3 278.3 279.5 279.7 272.2 259.1 0.0 0.0 0.0 240.7 236.0 233.8 233.5 234.6 235.7 235.6 236.2 239.8 245.4 248.8 250.1 250.4 249.9 249.1 247.2 241.7 231.7 216.0 197.6 192.1 197.2 204.3 206.5 205.4 201.7 200.8 203.5 208.2 210.9 210.3 207.1 203.5 202.3 202.1 202.8 205.1 207.3 208.1 207.3 205.4 202.1 199.3 200.5 206.1 212.9 212.9 206.5 197.1 191.1 191.2 196.8 203.6 207.8 208.7 206.4 201.7 197.3 197.4 200.7 206.5 209.9 209.9 206.9 203.3 201.2 203.4 207.7 210.5 208.9 207.6 206.5 204.2 202.0 203.6 205.7 210.8 213.1 214.3 210.6 204.1 199.7 202.3 211.9 217.7 215.1 215.0 215.3 213.3 0.0 0.0 216.3 0.0 0.0 195.2 209.0 205.3 201.3 196.8 195.7 195.8 195.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 431.9 425.8 421.0 417.7 419.3 422.7 423.4 421.1 417.9 411.9 408.6 411.5 414.9 419.6 426.7 424.2 0.0 0.0 0.0 280.8 279.4 369.7 371.8 370.6 369.8 370.5 370.5 371.4 374.6 374.0 361.0 332.7 328.0 0.0 359.0 365.0 371.1 373.8 371.9 369.7 369.6 371.0 378.3 390.4 403.4 414.1 417.9 417.5 417.0 416.6 415.7 414.9 414.5 413.8 413.5 413.6 413.8 415.2 416.9 417.3 415.4 415.0 414.1 411.4 409.2 403.5 392.7 375.6 367.5 365.8 368.8 370.7 370.8 369.0 367.0 366.0 366.2 367.8 369.2 368.8 367.2 366.7 367.3 367.3 368.8 368.5 366.8 365.3 364.0 363.0 365.0 367.3 367.7 365.5 364.0 364.9 367.4 370.3 371.3 369.9 366.5 364.1 368.3 385.8 406.7 409.4 399.7 376.0 357.2 359.9 372.4 377.5 366.9 345.1 320.6 315.8 321.8 329.0 329.9 326.4 324.5 324.8 324.9 325.3 325.1 324.7 325.9 326.8 326.6 326.6 326.4 325.4 323.2 322.5 323.9 326.2 328.6 329.9 329.5 329.0 328.2 327.7 327.9 329.6 331.2 331.4 334.2 340.1 335.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.Line truncated
}
]
Binary file not shown.
+16
View File
@@ -0,0 +1,16 @@
[
{
"index": "vocal_320_10687",
"language": "Mandarin",
"time": [
320,
10687
],
"duration": "0.23 0.34 0.26 0.70 0.52 0.46 0.36 0.44 0.14 0.24 0.64 0.47 0.51 1.10 0.28 0.38 0.32 0.32 0.38 0.32 0.31 0.19 1.45",
"text": "<SP> 除 了 想 你 你 <SP> 除 了 了 爱 你 你 <SP> 我 什 么 什 么 都 愿 愿 意",
"phoneme": "<SP> zh_chu2 zh_le5 zh_xiang3 zh_ni3 zh_ni3 <SP> zh_chu2 zh_le5 zh_le5 zh_ai4 zh_ni3 zh_ni3 <SP> zh_wo3 zh_shen2 zh_me5 zh_shen2 zh_me5 zh_dou1 zh_yuan4 zh_yuan4 zh_yi4",
"note_pitch": "0 62 65 67 67 69 0 67 69 67 65 67 69 67 67 66 64 64 60 60 65 67 0",
"note_type": "1 2 2 2 2 3 1 2 2 3 2 2 3 1 2 2 2 2 2 2 2 3 2",
"f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 294.1 288.2 290.6 294.6 295.7 292.8 291.5 294.4 295.4 294.7 293.5 292.2 294.4 295.8 293.2 297.7 320.1 338.1 348.3 348.7 344.4 342.6 346.2 354.8 356.9 353.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 401.1 403.8 405.0 402.1 398.4 395.6 393.3 392.0 391.4 390.3 390.2 390.3 390.3 391.3 392.8 393.6 391.7 390.5 391.4 391.6 391.5 393.5 393.8 390.8 387.4 387.8 389.3 390.8 392.2 391.6 390.2 389.8 389.1 388.4 390.0 395.5 397.2 396.7 395.5 395.1 394.6 394.9 395.6 395.2 394.6 395.4 395.9 394.0 391.7 390.7 391.7 392.6 391.6 395.7 405.8 441.7 462.5 463.8 450.2 430.9 414.5 415.2 426.7 439.8 454.4 462.9 447.8 422.5 400.6 403.2 423.5 451.3 482.3 492.8 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 435.9 414.3 406.2 402.8 398.2 396.7 396.7 395.2 398.3 416.9 441.6 446.9 441.5 437.3 435.3 432.9 431.6 432.3 434.8 436.7 433.8 422.8 406.7 390.9 382.0 381.7 384.7 382.7 368.1 357.4 355.6 355.1 352.4 348.1 346.4 348.7 351.6 355.0 354.4 351.7 349.9 349.2 348.1 346.0 345.4 344.2 344.4 345.5 346.6 349.0 349.7 349.1 349.5 349.6 349.7 349.4 349.5 352.2 354.5 355.7 355.6 356.9 359.4 361.5 363.8 360.4 354.3 357.2 363.9 372.4 382.6 399.1 402.9 400.6 395.5 390.5 388.9 390.2 391.1 391.9 391.5 390.4 390.2 391.3 391.6 391.0 388.6 386.2 389.1 403.8 430.2 441.8 449.5 448.1 443.2 438.3 432.9 430.7 434.6 442.0 447.6 446.0 440.3 434.7 431.1 435.7 442.3 445.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 423.8 402.1 398.5 397.6 393.4 390.5 391.4 402.5 427.6 442.5 442.5 435.1 430.1 430.8 439.9 447.4 442.2 426.1 412.3 399.8 391.8 389.5 388.5 387.8 386.4 384.7 384.5 387.9 391.2 391.8 392.9 393.8 392.0 392.0 395.4 398.1 398.2 396.3 393.1 391.1 388.9 386.6 383.1 381.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 357.8 367.6 370.9 365.4 359.3 355.7 358.1 372.7 396.8 404.0 398.4 392.6 389.2 388.8 383.6 362.3 341.1 325.1 326.3 327.9 331.3 333.3 326.0 319.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 356.4 353.5 349.0 346.5 338.4 328.4 323.7 326.2 334.2 338.4 331.6 309.8 280.6 256.4 252.8 256.1 258.3 256.9 257.6 261.1 259.7 258.7 258.6 260.2 262.2 262.4 262.9 264.2 263.4 260.2 256.4 253.2 238.7 223.3 231.6 257.7 258.1 258.4 258.8 258.0 256.8 255.8 254.0 255.1 258.1 261.9 263.7 262.6 256.7 253.1 250.2 246.7 258.7 294.4 327.1 342.6 346.0 344.1 341.2 342.3 345.2 350.1 364.7 382.5 396.1 396.1 389.2 381.8 381.1 387.2 397.0 399.3 390.7 374.3 360.1 350.6 346.6 347.4 350.7 354.4 354.3 351.7 349.8 348.4 346.9 347.0 348.4 349.5 351.0 352.3 353.6 353.3 350.5 348.3 345.5 344.3 344.4 347.0 350.6 352.0 351.0 350.8 350.1 347.6 345.7 347.0 350.3 351.7 350.7 348.5 346.9 347.7 349.0 349.1 348.5 346.8 346.3 348.4 349.0 349.2 351.1 349.6 348.3 350.5 351.1 348.0 347.6 349.1 351.3 356.0 361.3 360.6 354.0 341.0 316.2 302.9 302.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0"
}
]
Binary file not shown.
+53 -131
View File
@@ -1,146 +1,68 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the central pipeline registry (2.5/N)."""
from __future__ import annotations
import logging
"""Tests for out of tree registration to OMNI_PIPELINES."""
import pytest
from transformers import PretrainedConfig
from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES
from vllm_omni.config.stage_config import (
_PIPELINE_REGISTRY,
PipelineConfig,
StageExecutionType,
StagePipelineConfig,
register_pipeline,
)
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline
from vllm_omni.config.stage_config import PipelineConfig, pipeline_cfg_resolver
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestCentralRegistryDeclarations:
"""Every in-tree pipeline must be declared exactly once in the central registry."""
@pytest.fixture
def custom_resolver():
"""Build a reusable custom resolver for PipelineConfigs."""
def test_omni_entries_visible_in_registry(self):
for key in _OMNI_PIPELINES:
assert key in _PIPELINE_REGISTRY
class CustomConfigType(PretrainedConfig):
pass
def test_expected_omni_pipelines_present(self):
# Guard against accidental removal during future refactors.
assert "qwen2_5_omni" in _OMNI_PIPELINES
assert "qwen2_5_omni_thinker_only" in _OMNI_PIPELINES
assert "qwen3_omni_moe" in _OMNI_PIPELINES
assert "aura_omni" in _OMNI_PIPELINES
assert "qwen3_tts" in _OMNI_PIPELINES
@pipeline_cfg_resolver(config_type=CustomConfigType)
def custom_resolver(
hf_config: CustomConfigType,
) -> PipelineConfig:
return PipelineConfig(model_type="resolved_type")
return custom_resolver
class TestLazyLoading:
"""Pipelines are imported only on first access."""
def test_contains_without_import(self):
# ``in`` hits the lazy map, not the loaded cache.
assert "qwen3_omni_moe" in _PIPELINE_REGISTRY
def test_getitem_loads_correct_pipeline(self):
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
assert pipeline.model_type == "qwen3_omni_moe"
assert pipeline.model_arch == "Qwen3OmniMoeForConditionalGeneration"
def test_getitem_loads_aura_omni_pipeline(self):
pipeline = _PIPELINE_REGISTRY["aura_omni"]
assert pipeline.model_type == "aura_omni"
assert [stage.model_stage for stage in pipeline.stages] == [
"asr",
"aura",
"qwen3_tts",
"code2wav",
]
assert pipeline.stages[0].model_arch == "Qwen3ASRForConditionalGeneration"
assert pipeline.stages[1].model_arch == "AuraQwen3VLForConditionalGeneration"
assert pipeline.stages[2].model_arch == "Qwen3TTSTalkerForConditionalGeneration"
assert pipeline.stages[3].model_arch == "Qwen3TTSCode2Wav"
def test_unknown_model_type_returns_none_via_get(self):
assert _PIPELINE_REGISTRY.get("not_a_real_pipeline") is None
def test_unknown_model_type_raises_keyerror_via_getitem(self):
with pytest.raises(KeyError):
_PIPELINE_REGISTRY["not_a_real_pipeline"]
def test_iteration_yields_registered_pipelines(self):
keys = set(_PIPELINE_REGISTRY)
assert "qwen2_5_omni" in keys
assert "qwen3_omni_moe" in keys
def test_iteration_resilience(self, monkeypatch, caplog):
# Synthetic failure to verify iteration robustness against load errors.
target_logger = logging.getLogger("vllm_omni.config.stage_config")
target_logger.addHandler(caplog.handler)
prev_level = target_logger.level
target_logger.setLevel(logging.ERROR)
try:
monkeypatch.setattr(
_PIPELINE_REGISTRY,
"_lazy_map",
{
"healthy": ("vllm_omni.model_executor.models.qwen2_5_omni.pipeline", "QWEN2_5_OMNI_PIPELINE"),
"broken": ("vllm_omni.non_existent_module", "SomeConfig"),
"attr_fail": ("vllm_omni.config.stage_config", "NON_EXISTENT_VAR"),
},
)
monkeypatch.setattr(_PIPELINE_REGISTRY, "_loaded", {})
# Assert that list(_PIPELINE_REGISTRY.values()) and dict(_PIPELINE_REGISTRY.items())
# both yield only healthy entries.
healthy_instance = _PIPELINE_REGISTRY["healthy"]
assert list(_PIPELINE_REGISTRY.values()) == [healthy_instance]
assert dict(_PIPELINE_REGISTRY.items()) == {"healthy": healthy_instance}
assert "Failed to import pipeline module" in caplog.text
assert "Pipeline variable 'NON_EXISTENT_VAR' not found" in caplog.text
finally:
target_logger.removeHandler(caplog.handler)
target_logger.setLevel(prev_level)
def test_register_pipeline_config(clean_pipeline_registry):
"""Ensure that we can register a custom pipeline config to OMNI_PIPELINES."""
new_model_type = "new_model_type"
pipe_cfg = PipelineConfig(model_type=new_model_type)
assert new_model_type not in OMNI_PIPELINES
register_pipeline(pipe_cfg)
assert new_model_type in OMNI_PIPELINES
assert OMNI_PIPELINES[new_model_type] is pipe_cfg
class TestDynamicRegistration:
"""``register_pipeline()`` still works for plugins and tests."""
def test_register_pipeline_config_with_model_type(clean_pipeline_registry):
"""Ensure that we can register a custom pipeline config with an explicit model_type to OMNI_PIPELINES."""
new_model_type = "new_model_type"
unused_model_type = "foo"
pipe_cfg = PipelineConfig(model_type=unused_model_type)
assert new_model_type not in OMNI_PIPELINES
assert unused_model_type not in OMNI_PIPELINES
def test_register_adds_to_registry(self):
custom = PipelineConfig(
model_type="_test_dynamic_registration",
model_arch="DynamicTestModel",
stages=(
StagePipelineConfig(
stage_id=0,
model_stage="test",
execution_type=StageExecutionType.LLM_AR,
input_sources=(),
final_output=True,
),
),
)
register_pipeline(custom)
try:
assert "_test_dynamic_registration" in _PIPELINE_REGISTRY
assert _PIPELINE_REGISTRY["_test_dynamic_registration"] is custom
finally:
# Don't leak the test registration into other tests.
if "_test_dynamic_registration" in _PIPELINE_REGISTRY:
del _PIPELINE_REGISTRY["_test_dynamic_registration"]
# Registering with an explicitly provided model_type uses
# the passed value instead of the pipeline_cfg.model_type
register_pipeline(pipe_cfg, new_model_type)
assert new_model_type in OMNI_PIPELINES
assert unused_model_type not in OMNI_PIPELINES
assert OMNI_PIPELINES[new_model_type] is pipe_cfg
def test_dynamic_registration_overrides_lazy_entry(self):
# Build a substitute for qwen3_omni_moe that we can distinguish.
original = _PIPELINE_REGISTRY["qwen3_omni_moe"]
override = PipelineConfig(
model_type="qwen3_omni_moe",
model_arch="OverriddenArch",
stages=original.stages,
)
register_pipeline(override)
try:
assert _PIPELINE_REGISTRY["qwen3_omni_moe"].model_arch == "OverriddenArch"
finally:
# Remove the dynamic override so later tests see the original.
if "qwen3_omni_moe" in _PIPELINE_REGISTRY._loaded:
del _PIPELINE_REGISTRY["qwen3_omni_moe"]
def test_register_resolver(custom_resolver, clean_pipeline_registry):
"""Ensure that we can register a custom resolver to OMNI_PIPELINES."""
new_model_type = "new_model_type"
assert new_model_type not in OMNI_PIPELINES
register_pipeline(custom_resolver, new_model_type)
assert new_model_type in OMNI_PIPELINES
assert OMNI_PIPELINES[new_model_type] is custom_resolver
def test_register_resolver_requires_model_type(custom_resolver, clean_pipeline_registry):
"""Ensure that registering a custom resolver to OMNI_PIPELINES requires an explicit model_type."""
with pytest.raises(ValueError):
register_pipeline(custom_resolver)
+1
View File
@@ -17,6 +17,7 @@ from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules
bootstrap_vllm_layer_custom_op_modules()
pytest_plugins = (
"tests.helpers.fixtures.config",
"tests.helpers.fixtures.env",
"tests.helpers.fixtures.log",
"tests.helpers.fixtures.run_args",
@@ -115,6 +115,7 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase):
("MiMoAudioModel", "code2wav"),
("Qwen3TTSCode2Wav", "code2wav"),
("CosyVoice3Model", "cosyvoice3_code2wav"),
("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),
("DyninOmniForConditionalGeneration", "token2image"),
("DyninOmniForConditionalGeneration", "token2audio"),
}
@@ -621,7 +621,6 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
{"consent": "c1", "name": "v1"},
("audio_sample", "speaker_embedding", "provided"),
id="neither_audio_nor_embedding",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_ok",
@@ -673,7 +672,6 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
{"consent": "consent_ok", "name": "evil/name"},
("voice name", "invalid voice name"),
id="invalid_name_path_sep",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_ok",
@@ -698,21 +696,18 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
{"consent": "consent_ok", "name": "v_bad_mime_{uuid}"},
("mime", "unsupported", "audio/mp4"),
id="audio_unsupported_mime",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_short",
{"consent": "consent_ok", "name": "v_short_{uuid}"},
("too short", "reference audio"),
id="audio_too_short",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_bad_body",
{"consent": "consent_ok", "name": "v_bad_wav_{uuid}"},
("decode", "audio", "Format not recognised"),
id="audio_decode_error",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
None,
@@ -723,7 +718,6 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
},
("speaker_embedding", "valid JSON"),
id="speaker_embedding_invalid_json",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
None,
@@ -734,7 +728,6 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
},
("speaker_embedding", "non-empty"),
id="speaker_embedding_empty",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
None,
@@ -745,7 +738,6 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
},
("speaker_embedding", "finite"),
id="speaker_embedding_nan",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
None,
@@ -756,7 +748,6 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
},
("expected", "dimensions", "1024"),
id="speaker_embedding_wrong_dim",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_ok",
@@ -767,7 +758,6 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
},
("mutually exclusive", "audio_sample", "speaker_embedding"),
id="audio_and_embedding_mutually_exclusive",
marks=_SKIP_ISSUE_3649,
),
],
)
@@ -19,9 +19,19 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
class StubScheduler:
def __init__(self, timesteps: list[int] | None = None, *, flow_shift: float = 1.0) -> None:
def __init__(
self,
timesteps: list[int] | None = None,
*,
flow_shift: float = 1.0,
use_karras_sigmas: bool = True,
) -> None:
self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64)
self.config = SimpleNamespace(num_train_timesteps=1000, flow_shift=flow_shift)
self.config = SimpleNamespace(
num_train_timesteps=1000,
flow_shift=flow_shift,
use_karras_sigmas=use_karras_sigmas,
)
self.set_timesteps_calls: list[tuple[int, torch.device]] = []
self.step_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
@@ -194,6 +204,8 @@ def make_cosmos3_pipeline():
pipeline._base_scheduler_config = pipeline.scheduler.config
pipeline._engine_init_flow_shift = 1.0
pipeline._current_flow_shift = 1.0
pipeline._base_scheduler_use_karras_sigmas = True
pipeline._current_scheduler_use_karras_sigmas = True
pipeline._guidance_scale = None
pipeline._num_timesteps = None
pipeline._cosmos3_branch_caches = None
@@ -292,8 +304,8 @@ def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch):
class _StubDiffusersScheduler:
from_config_calls: list[dict[str, Any]] = []
def __init__(self, *, flow_shift: float = 1.0) -> None:
self.config = SimpleNamespace(flow_shift=flow_shift)
def __init__(self, *, flow_shift: float = 1.0, use_karras_sigmas: bool = True) -> None:
self.config = SimpleNamespace(flow_shift=flow_shift, use_karras_sigmas=use_karras_sigmas)
@classmethod
def from_pretrained(cls, *args, **kwargs):
@@ -302,7 +314,13 @@ def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch):
@classmethod
def from_config(cls, config, **kwargs):
cls.from_config_calls.append({"config": config, "kwargs": dict(kwargs)})
return cls(flow_shift=kwargs.get("flow_shift", getattr(config, "flow_shift", 1.0)))
return cls(
flow_shift=kwargs.get("flow_shift", getattr(config, "flow_shift", 1.0)),
use_karras_sigmas=kwargs.get(
"use_karras_sigmas",
getattr(config, "use_karras_sigmas", True),
),
)
class _StubVideoProcessor:
def __init__(self, *args, **kwargs) -> None:
@@ -349,10 +367,13 @@ def test_pipeline_init_skips_tokenizer_when_sound_disabled(stub_real_pipeline_in
pipeline = Cosmos3OmniDiffusersPipeline(od_config=_make_od_config(sound_gen=False))
assert stub_real_pipeline_init.from_config_calls == []
assert pipeline._sound_tokenizer is None
assert pipeline.transformer.sound_gen is False
assert not hasattr(pipeline.transformer, "audio_proj_in")
assert not hasattr(pipeline.transformer, "audio_proj_out")
assert pipeline._base_scheduler_use_karras_sigmas is True
assert pipeline._current_scheduler_use_karras_sigmas is True
def test_pipeline_init_rebuilds_scheduler_with_cosmos3_defaults(stub_real_pipeline_init) -> None:
@@ -366,8 +387,52 @@ def test_pipeline_init_rebuilds_scheduler_with_cosmos3_defaults(stub_real_pipeli
assert stub_real_pipeline_init.from_config_calls
call = stub_real_pipeline_init.from_config_calls[-1]
assert getattr(call["config"], "flow_shift") == 1.0
assert call["kwargs"] == {"flow_shift": 2.5, "use_karras_sigmas": False}
assert call["kwargs"] == {"flow_shift": 2.5}
assert pipeline._engine_init_flow_shift == 2.5
assert pipeline._base_scheduler_use_karras_sigmas is True
def test_set_flow_shift_restores_checkpoint_scheduler_mode(
make_cosmos3_pipeline,
monkeypatch: pytest.MonkeyPatch,
) -> None:
from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3
class RecordingScheduler:
from_config_calls: list[dict[str, Any]] = []
def __init__(self, *, flow_shift: float, use_karras_sigmas: bool) -> None:
self.config = SimpleNamespace(flow_shift=flow_shift, use_karras_sigmas=use_karras_sigmas)
@classmethod
def from_config(cls, config, **kwargs):
cls.from_config_calls.append({"config": config, "kwargs": dict(kwargs)})
return cls(
flow_shift=kwargs.get("flow_shift", getattr(config, "flow_shift", 1.0)),
use_karras_sigmas=kwargs.get(
"use_karras_sigmas",
getattr(config, "use_karras_sigmas", True),
),
)
monkeypatch.setattr(pipeline_cosmos3, "UniPCMultistepScheduler", RecordingScheduler)
pipeline = make_cosmos3_pipeline()
pipeline._set_flow_shift(10.0, use_karras_sigmas=False)
assert RecordingScheduler.from_config_calls[-1]["kwargs"] == {
"flow_shift": 10.0,
"use_karras_sigmas": False,
}
assert pipeline._current_flow_shift == 10.0
assert pipeline._current_scheduler_use_karras_sigmas is False
pipeline._set_flow_shift(10.0)
assert RecordingScheduler.from_config_calls[-1]["kwargs"] == {"flow_shift": 10.0}
assert pipeline._current_flow_shift == 10.0
assert pipeline._current_scheduler_use_karras_sigmas is True
pipeline._set_flow_shift(10.0)
assert len(RecordingScheduler.from_config_calls) == 2
def test_pipeline_init_passes_tokenizer_attrs_into_transformer(
@@ -1022,7 +1087,12 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint
pipeline._encode_video_tensor = fake_encode
pipeline._prepare_transfer_latents = fake_prepare
pipeline.diffuse_transfer = fake_diffuse_transfer
pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target)
def fake_set_flow_shift(target, *, use_karras_sigmas=None):
captured.setdefault("flow_shifts", []).append(target)
captured.setdefault("scheduler_use_karras_sigmas", []).append(use_karras_sigmas)
pipeline._set_flow_shift = fake_set_flow_shift
pipeline._decode_latents = lambda latents: torch.zeros(1, 3, 5, 16, 16, device="meta")
control = torch.zeros(3, 5, 16, 16, dtype=torch.uint8)
@@ -1055,6 +1125,8 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint
assert captured["format_frame_rate"] == expected_fps
assert captured["shared_kwargs"]["fps"] == expected_fps
assert captured["flow_shifts"] == [10.0]
# Transfer uses the V2V flow-sigma schedule (karras off) so flow_shift applies.
assert captured["scheduler_use_karras_sigmas"] == [False]
assert output.custom_output["fps"] == expected_fps
assert output.output["video"].device.type == "meta"
@@ -1068,7 +1140,7 @@ def test_forward_transfer_runs_multichunk_overlap_path(
pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16)
pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (_ids(2), _mask(), _ids(1), _mask())
pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target)
pipeline._set_flow_shift = lambda target, **_kwargs: captured.setdefault("flow_shifts", []).append(target)
original_prepare = pipeline._prepare_transfer_latents
@@ -1213,7 +1285,12 @@ class TestForwardRouting:
pipeline._format_and_tokenize_prompts = fake_format
pipeline._prepare_latents = fake_prepare
pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target)
def fake_set_flow_shift(target, *, use_karras_sigmas=None):
captured.setdefault("flow_shifts", []).append(target)
captured.setdefault("scheduler_use_karras_sigmas", []).append(use_karras_sigmas)
pipeline._set_flow_shift = fake_set_flow_shift
pipeline.diffuse = fake_diffuse
pipeline._decode_latents = lambda latents: latents
return captured
@@ -1227,7 +1304,6 @@ class TestForwardRouting:
{
"key": "image",
"is_t2i": True,
"use_system_prompt": True,
"flow": [3.0],
"steps": [50, 50],
"frames": 1,
@@ -1239,7 +1315,6 @@ class TestForwardRouting:
{
"key": "video",
"is_t2i": False,
"use_system_prompt": True,
"flow": [1.0],
"steps": [35],
"frames": 189,
@@ -1261,9 +1336,9 @@ class TestForwardRouting:
assert expected["key"] in output.output
assert captured["format"]["is_t2i"] is expected["is_t2i"]
assert captured["format"]["use_system_prompt"] is expected["use_system_prompt"]
assert captured["format"]["num_frames"] == expected["frames"]
assert captured["flow_shifts"] == expected["flow"]
assert captured["scheduler_use_karras_sigmas"] == [None]
assert [call[0] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"]
def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> None:
@@ -1290,7 +1365,6 @@ class TestForwardRouting:
)
)
assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask
assert captured["format"]["use_system_prompt"] is True
video_tensor = torch.zeros(1, 3, 5, 16, 16)
v2v_condition = torch.full((1, 2, 2, 1, 1), 4.0)
@@ -1316,8 +1390,8 @@ class TestForwardRouting:
)
)
assert captured["flow_shifts"][-1] == 10.0
assert captured["scheduler_use_karras_sigmas"][-1] is False
assert captured["format"]["negative_prompt"] == ""
assert captured["format"]["use_system_prompt"] is True
assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is v2v_mask
assert captured["diffuse_calls"][-1]["condition_latents"] is v2v_condition
@@ -1333,7 +1407,6 @@ class TestForwardRouting:
)
)
assert captured["diffuse_calls"][-1]["sound_latents"] is sound_latents
assert captured["format"]["use_system_prompt"] is True
assert output.output["audio_sample_rate"] == 10
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4)
@@ -1359,7 +1432,6 @@ class TestForwardRouting:
)
)
assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
assert captured["format"]["use_system_prompt"] is False
assert output.custom_output["action"].shape == (1, 2, 2)
assert "action_only_output" not in output.custom_output
@@ -0,0 +1,54 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for HunyuanImage3 AR image RGB conversion."""
from types import SimpleNamespace
import pytest
import torch
from PIL import Image
from vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3 import (
HunyuanImage3Processor,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _FixedResolutionGroup:
def get_target_size(self, width: int, height: int):
return width, height
def get_base_size_and_ratio_index(self, width: int, height: int):
return 1024, 0
def test_ar_processor_uses_pil_rgb_conversion_for_rgba_images():
captured: dict[str, object] = {}
def fake_vit_processor(image: Image.Image):
captured["mode"] = image.mode
captured["pixel"] = image.getpixel((0, 0))
return {
"pixel_values": torch.zeros((1, 1, 3), dtype=torch.float32),
"pixel_attention_mask": torch.ones((1, 1), dtype=torch.bool),
"spatial_shapes": torch.tensor([[1, 1]], dtype=torch.long),
}
def fake_vae_processor(_image: Image.Image):
return torch.zeros((1, 3, 1, 1), dtype=torch.float32)
processor = object.__new__(HunyuanImage3Processor)
processor.hf_config = SimpleNamespace(
vit={"num_channels": 3},
vae_downsample_factor=(1, 1),
patch_size=1,
)
processor.reso_group = _FixedResolutionGroup()
processor.vision_encoder_processor = fake_vit_processor
processor.vae_processor = fake_vae_processor
processor.process_image(Image.new("RGBA", (1, 1), color=(255, 0, 0, 0)))
assert captured["mode"] == "RGB"
assert captured["pixel"] == (255, 0, 0)
@@ -20,6 +20,7 @@ RATIO_START = 200
RATIO_END = 210
RATIO_OTHER_START = 220
RATIO_OTHER_END = 223
TIMESTEP = 224
class FakeSamplerModel:
@@ -60,6 +61,67 @@ class FakeSamplerModel:
_apply_ratio_restriction = _Real._apply_ratio_restriction
class FakeTokenEmbedding:
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return torch.stack(
[
input_ids.to(dtype=torch.float32),
input_ids.to(dtype=torch.float32) + 0.5,
],
dim=-1,
)
class FakeEmbeddingModel:
def __init__(self):
self.model = FakeTokenEmbedding()
self._timestep_token_id = TIMESTEP
from vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3 import (
HunyuanImage3ForConditionalGeneration as _Real,
)
embed_input_ids = _Real.embed_input_ids
def _timestep_encode(self, timestep: torch.Tensor) -> torch.Tensor:
return torch.full(
(1, 2),
-7.0,
device=timestep.device,
dtype=timestep.dtype,
)
class TestTimestepEmbedding:
def test_timestep_embedding_avoids_scalar_item_sync(self, monkeypatch):
model = FakeEmbeddingModel()
input_ids = torch.tensor([1, TIMESTEP, 2, TIMESTEP], dtype=torch.long)
def fail_item(_self):
raise AssertionError("embed_input_ids should not synchronize timestep count")
monkeypatch.setattr(torch.Tensor, "item", fail_item)
actual = model.embed_input_ids(input_ids)
expected = model.model.embed_input_ids(input_ids)
expected[input_ids == TIMESTEP] = -7.0
torch.testing.assert_close(actual, expected)
def test_timestep_embedding_preserves_no_timestep_inputs(self, monkeypatch):
model = FakeEmbeddingModel()
input_ids = torch.tensor([1, 2, 3], dtype=torch.long)
def fail_item(_self):
raise AssertionError("embed_input_ids should not synchronize timestep count")
monkeypatch.setattr(torch.Tensor, "item", fail_item)
actual = model.embed_input_ids(input_ids)
torch.testing.assert_close(actual, model.model.embed_input_ids(input_ids))
class TestGetForcedToken:
"""Tests for the stateless _get_forced_token method."""
@@ -9,6 +9,7 @@ from unittest.mock import Mock
import pytest
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler
@@ -33,6 +34,7 @@ def _make_engine() -> DiffusionEngine:
engine._rpc_lock = threading.RLock()
engine._cv = threading.Condition(engine._rpc_lock)
engine._out_queue = {}
engine._out_queue_streaming = {}
engine._closed = False
engine._shutdown_complete = False
engine.abort_queue = queue.Queue()
@@ -57,6 +59,23 @@ def test_close_completes_pending_async_waiters() -> None:
event_loop.close()
def test_close_completes_pending_streaming_waiters() -> None:
engine = _make_engine()
event_loop = asyncio.new_event_loop()
try:
engine.main_loop = event_loop
queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
engine._out_queue_streaming["pending-stream"] = queue
engine.close()
output = queue.get_nowait()
assert output.error == "DiffusionEngine is closed."
assert output.finished is True
finally:
event_loop.close()
def test_handle_finished_requests_ignores_already_drained_waiter() -> None:
class RacingOutQueue(dict):
def __contains__(self, key):
@@ -9,6 +9,7 @@ import torch
import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
pytestmark = [pytest.mark.diffusion]
@@ -31,6 +32,46 @@ class _DummyPipeline:
return self._output
class _ChunkStepPipeline:
device = torch.device("cpu")
supports_step_execution = True
def __init__(self, outputs):
self._outputs = outputs
self.prepare_calls = 0
self.decode_calls = 0
def prepare_encode(self, state):
self.prepare_calls += 1
state.prompt_embeds = torch.zeros(1, 1, 1)
state.latents = torch.zeros(1, 1)
state.timesteps = torch.tensor([1.0, 0.0])
state.step_index = 0
state.step_in_chunk = 0
state.chunk_num_steps = 2
state.total_chunks = len(self._outputs)
return state
def denoise_step(self, input_batch, states):
del states
return torch.ones_like(input_batch.latents)
def step_scheduler(self, state, noise_pred):
state.latents = noise_pred
state.step_index += 1
state.step_in_chunk += 1
def post_decode(self, state):
output = self._outputs[self.decode_calls]
self.decode_calls += 1
state.chunk_index += 1
state.step_index = 0
state.step_in_chunk = 0
if not state.request_denoise_completed:
state.latents = torch.zeros(1, 1)
return output
def _make_request(skip_cache_refresh: bool = True):
sampling_params = SimpleNamespace(
generator=None,
@@ -52,11 +93,13 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
runner.pipeline = _DummyPipeline(output=SimpleNamespace(output="ok"))
runner.cache_backend = cache_backend
runner.offload_backend = None
runner.state_cache = {}
runner.prompt_embed_cache = None
runner.od_config = SimpleNamespace(
cache_backend=cache_backend_name,
enable_cache_dit_summary=enable_cache_dit_summary,
parallel_config=SimpleNamespace(use_hsdp=False),
streaming_output=False,
)
runner.kv_transfer_manager = SimpleNamespace(
receive_kv_cache=lambda req, target_device=None: None,
@@ -66,6 +109,49 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
return runner
@pytest.mark.core_model
@pytest.mark.cpu
def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):
"""Step streaming returns empty step results until a chunk decode boundary."""
chunks = [
DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2),
DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2),
]
runner = _make_runner(cache_backend=None, cache_backend_name=None)
runner.pipeline = _ChunkStepPipeline(chunks)
runner.od_config.streaming_output = True
runner.od_config.step_execution = True
req = _make_request(skip_cache_refresh=True)
req.request_id = "req"
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
scheduler_output = SimpleNamespace(
finished_req_ids=set(),
scheduled_new_reqs=[SimpleNamespace(request_id="req", req=req)],
scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
)
first = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
assert first.get_request_output("req").result is None
scheduler_output = SimpleNamespace(
finished_req_ids=set(),
scheduled_new_reqs=[],
scheduled_cached_reqs=SimpleNamespace(request_ids=["req"]),
)
second = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
assert second.get_request_output("req").result == chunks[0]
assert second.get_request_output("req").finished is False
DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
fourth = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
assert fourth.get_request_output("req").result == chunks[1]
assert fourth.get_request_output("req").finished is True
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypatch):
@@ -76,6 +162,9 @@ def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypa
cache_summary_calls = []
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
monkeypatch.setattr(
model_runner_module,
"cache_summary",
@@ -101,6 +190,9 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp
cache_summary_calls = []
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
monkeypatch.setattr(
model_runner_module,
"cache_summary",
@@ -157,6 +249,7 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
cache_config={},
model_class_name="NextStep11Pipeline",
enforce_eager=True,
streaming_output=False,
)
monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
+100 -3
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
import queue
import threading
from types import SimpleNamespace
@@ -238,6 +239,35 @@ class TestRequestScheduler:
assert finished == {req_id}
assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
def test_streaming_output_keeps_request_running_until_final_chunk(self) -> None:
scheduler = StepScheduler()
scheduler.initialize(SimpleNamespace())
req_id = scheduler.add_request(_make_request("stream"))
sched_output = scheduler.schedule()
chunk = RunnerOutput(
request_id=req_id,
step_index=1,
finished=False,
result=DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2),
)
finished = scheduler.update_from_output(sched_output, chunk)
assert finished == set()
assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.RUNNING
assert scheduler.has_requests() is True
final_chunk = RunnerOutput(
request_id=req_id,
step_index=2,
finished=True,
result=DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2),
)
finished = scheduler.update_from_output(sched_output, final_chunk)
assert finished == {req_id}
assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
def test_fifo_single_request_scheduling(self) -> None:
req_id_a = self.scheduler.add_request(_make_request("a"))
req_id_b = self.scheduler.add_request(_make_request("b"))
@@ -369,6 +399,7 @@ class TestRequestScheduler:
class TestDiffusionEngine:
def test_add_req_and_wait_for_response_single_path(self, mocker: MockerFixture) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.od_config = SimpleNamespace(streaming_output=False)
engine.scheduler = RequestScheduler()
engine.scheduler.initialize(SimpleNamespace())
engine._rpc_lock = threading.RLock()
@@ -391,6 +422,7 @@ class TestDiffusionEngine:
scheduler = _StubScheduler(request, runner_output)
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.od_config = SimpleNamespace(streaming_output=False)
engine.scheduler = scheduler
engine._rpc_lock = threading.RLock()
engine._cv = threading.Condition(engine._rpc_lock)
@@ -410,7 +442,7 @@ class TestDiffusionEngine:
) -> None:
request = _make_request("init")
scheduler = _StubScheduler(request, DiffusionOutput(output=None))
od_config = SimpleNamespace(model_class_name="mock_model")
od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False)
fake_executor_cls = mocker.Mock(return_value=mocker.Mock())
monkeypatch.setattr(
@@ -472,7 +504,7 @@ class TestDiffusionEngine:
def test_finalize_finished_request_returns_aborted_output(self) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.scheduler = RequestScheduler()
engine.scheduler = StepScheduler()
engine.scheduler.initialize(SimpleNamespace())
req_id = engine.scheduler.add_request(_make_request("req-finalize"))
@@ -483,12 +515,77 @@ class TestDiffusionEngine:
assert output.aborted is True
assert output.abort_message == "Request req-finalize aborted."
@pytest.mark.asyncio
async def test_streaming_runner_output_notifies_each_chunk(self) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.scheduler = StepScheduler()
engine.scheduler.initialize(SimpleNamespace())
engine._rpc_lock = threading.RLock()
engine._cv = threading.Condition(engine._rpc_lock)
engine._out_queue_streaming = {}
engine.main_loop = asyncio.get_running_loop()
req_id = engine.scheduler.add_request(_make_request("stream-engine"))
queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
engine._out_queue_streaming[req_id] = queue
sched_output = engine.scheduler.schedule()
chunk = RunnerOutput(
request_id=req_id,
step_index=1,
finished=False,
result=DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2),
)
finished_req_ids = engine.scheduler.update_from_output(sched_output, chunk)
engine._handle_step_streaming_runner_output(finished_req_ids, sched_output.scheduled_request_ids, chunk)
notified_chunk = await asyncio.wait_for(queue.get(), timeout=1)
assert notified_chunk.output == "chunk-0"
assert notified_chunk.finished is False
assert engine.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.RUNNING
final_chunk = RunnerOutput(
request_id=req_id,
step_index=2,
finished=True,
result=DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2),
)
finished_req_ids = engine.scheduler.update_from_output(sched_output, final_chunk)
engine._handle_step_streaming_runner_output(finished_req_ids, sched_output.scheduled_request_ids, final_chunk)
notified_final = await asyncio.wait_for(queue.get(), timeout=1)
assert notified_final.output == "chunk-1"
assert notified_final.finished is True
assert engine.scheduler.get_request_state(req_id) is None
@pytest.mark.asyncio
async def test_finished_streaming_request_without_runner_output_notifies_waiter(self) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.scheduler = RequestScheduler()
engine.scheduler.initialize(SimpleNamespace())
engine._rpc_lock = threading.RLock()
engine._cv = threading.Condition(engine._rpc_lock)
engine._out_queue_streaming = {}
engine.main_loop = asyncio.get_running_loop()
req_id = engine.scheduler.add_request(_make_request("stream-abort"))
queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
engine._out_queue_streaming[req_id] = queue
engine.scheduler.finish_requests(req_id, DiffusionRequestStatus.FINISHED_ABORTED)
engine._handle_empty_streaming_requests({req_id})
output = await asyncio.wait_for(queue.get(), timeout=1)
assert output.aborted is True
assert output.finished is True
assert engine.scheduler.get_request_state(req_id) is None
def test_initializes_step_scheduler_when_step_execution_enabled(
self,
monkeypatch: pytest.MonkeyPatch,
mocker: MockerFixture,
) -> None:
od_config = SimpleNamespace(model_class_name="mock_model")
od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False)
od_config.step_execution = True
fake_executor = mocker.Mock()
fake_executor_cls = mocker.Mock(return_value=fake_executor)
@@ -249,6 +249,7 @@ def _make_runner():
runner.od_config = SimpleNamespace(
cache_backend=None,
parallel_config=SimpleNamespace(use_hsdp=False),
streaming_output=False,
)
runner.device = torch.device("cpu")
runner.pipeline = _StepPipeline()
@@ -267,6 +268,7 @@ def _make_distributed_runner(mode: str, device: torch.device):
runner.od_config = SimpleNamespace(
cache_backend=None,
parallel_config=SimpleNamespace(use_hsdp=False),
streaming_output=False,
)
runner.device = device
runner.pipeline = _DistributedStepPipeline(mode=mode, device=device)
@@ -327,7 +329,7 @@ def _make_cached_scheduler_output(request_id="req-1", step_id=1, finished_req_id
def _make_engine(scheduler, execute_fn=None) -> DiffusionEngine:
engine = object.__new__(DiffusionEngine)
engine.od_config = SimpleNamespace(model_class_name="QwenImagePipeline")
engine.od_config = SimpleNamespace(model_class_name="QwenImagePipeline", streaming_output=False)
engine.pre_process_func = None
engine.post_process_func = None
engine.scheduler = scheduler
@@ -563,6 +565,7 @@ class TestRunner:
step_execution=True,
model_class_name="RequestOnlyPipeline",
parallel_config=SimpleNamespace(use_hsdp=False),
streaming_output=False,
)
runner.device = torch.device("cpu")
runner.pipeline = None
@@ -710,6 +713,7 @@ class TestExecutor:
def test_execute_step_passes_through_runner_output(self, mocker: MockerFixture):
executor = object.__new__(MultiprocDiffusionExecutor)
executor.od_config = SimpleNamespace(streaming_output=False)
executor._ensure_open = lambda: None
expected = RunnerOutput(request_id="req-step", step_index=1, finished=False, result=None)
executor.collective_rpc = mocker.Mock(return_value=expected)
@@ -0,0 +1,525 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Integration tests for diffusion pipeline streaming output mode.
Covers:
- Mock pipeline -> StageDiffusionClient (ZMQ subprocess path)
- Mock pipeline -> InlineStageDiffusionClient -> Orchestrator -> AsyncOmni
- Mock pipeline -> InlineStageDiffusionClient -> Orchestrator -> AsyncOmni -> `/v1/realtime/video`
"""
import asyncio
import contextlib
import queue
import time
from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock, patch
import numpy as np
import pytest
import torch
import zmq
from fastapi import FastAPI, WebSocket
from pytest_mock import MockerFixture
from starlette.testclient import TestClient
from vllm.utils.network_utils import get_open_zmq_ipc_path
import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
from tests.engine.test_orchestrator import OrchestratorFixture, _build_harness, _wait_for
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.distributed.omni_connectors.utils.serialization import (
OmniMsgpackDecoder,
OmniMsgpackEncoder,
)
from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
from vllm_omni.engine.messages import ShutdownRequestMessage, StageSubmissionMessage
from vllm_omni.engine.stage_init_utils import StageMetadata
from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni
from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
class _StepStreamingPipeline:
supports_step_execution = True
def __init__(self, outputs: list[DiffusionOutput]) -> None:
self.outputs = outputs
self.requests: list[Any] = []
def step_outputs(self, request):
self.requests.append(request)
return list(self.outputs)
class _FailingStreamingPipeline:
supports_step_execution = True
def __init__(self) -> None:
self.requests: list[Any] = []
def step_outputs(self, request):
self.requests.append(request)
raise RuntimeError("stream failed midway")
class _PipelineBackedEngine:
"""DiffusionEngine stand-in that exposes step-streaming chunks from a mock pipeline."""
def __init__(self, pipeline: _StepStreamingPipeline | _FailingStreamingPipeline) -> None:
self.pipeline = pipeline
self.executor = SimpleNamespace(
register_failure_callback=MagicMock(),
check_health=MagicMock(),
)
async def step_streaming(self, request):
try:
outputs = self.pipeline.step_outputs(request)
except Exception as exc:
outputs = [
DiffusionOutput(custom_output={"chunk": 0}, finished=False),
DiffusionOutput(error=str(exc), finished=True),
]
for output in outputs:
if output.error is not None:
yield [OmniRequestOutput.from_error(request_id=request.request_id, error_message=output.error)]
continue
custom_output = output.custom_output or {}
yield [
OmniRequestOutput.from_diffusion(
request_id=request.request_id,
images=list(custom_output.get("images") or []),
custom_output=custom_output,
finished=output.finished,
)
]
def abort(self, request_id: str) -> None:
del request_id
class TestPipelineStreamingOutputToStageDiffusionClient:
"""Streaming pipeline output over ZMQ into ``StageDiffusionClient``."""
@pytest.mark.asyncio
async def test_streaming_output_reaches_stage_client_from_mock_pipeline(self) -> None:
"""Mock pipeline chunks over ZMQ reach StageDiffusionClient with correct finished flags."""
pipeline = _StepStreamingPipeline(
[
DiffusionOutput(custom_output={"chunk": 0}, finished=False),
DiffusionOutput(custom_output={"chunk": 1}, finished=True),
]
)
outputs = await self._run_streaming_client_proc(pipeline, request_id="req-stream")
assert [output.request_id for output in outputs] == ["req-stream", "req-stream"]
assert [output.finished for output in outputs] == [False, True]
assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
assert pipeline.requests[0].request_id == "req-stream"
@pytest.mark.asyncio
async def test_streaming_midway_error_reaches_stage_client_after_prior_chunk(self) -> None:
"""A pipeline error after the first chunk still delivers that chunk then a terminal error output."""
pipeline = _FailingStreamingPipeline()
outputs = await self._run_streaming_client_proc(pipeline, request_id="req-error")
assert len(outputs) == 2
assert outputs[0].request_id == "req-error"
assert outputs[0].finished is False
assert outputs[0].custom_output["chunk"] == 0
assert outputs[1].request_id == "req-error"
assert outputs[1].finished is True
assert outputs[1].error is not None
assert "stream failed midway" in outputs[1].error
@classmethod
async def _run_streaming_client_proc(
cls,
pipeline: _StepStreamingPipeline | _FailingStreamingPipeline,
*,
request_id: str,
) -> list[OmniRequestOutput]:
request_address = get_open_zmq_ipc_path()
response_address = get_open_zmq_ipc_path()
proc = object.__new__(StageDiffusionProc)
proc._od_config = SimpleNamespace(streaming_output=True)
proc._engine = _PipelineBackedEngine(pipeline)
proc._closed = False
client = cls._make_client(request_address, response_address)
proc_task = asyncio.create_task(proc.run_loop(request_address, response_address))
try:
await asyncio.sleep(0.05)
await client.add_request_async(
request_id,
"prompt",
OmniDiffusionSamplingParams(),
)
deadline = asyncio.get_running_loop().time() + 2.0
outputs: list[OmniRequestOutput] = []
while (not outputs or not outputs[-1].finished) and asyncio.get_running_loop().time() < deadline:
output = client.get_diffusion_output_nowait()
if output is not None:
outputs.append(output)
else:
await asyncio.sleep(0.01)
return outputs
finally:
try:
client._request_socket.send(client._encoder.encode({"type": "shutdown"}))
except Exception:
pass
try:
await asyncio.wait_for(proc_task, timeout=2)
finally:
cls._close_client(client)
@staticmethod
def _make_client(request_address: str, response_address: str) -> StageDiffusionClient:
client = object.__new__(StageDiffusionClient)
client.stage_id = 0
client.replica_id = 0
client.final_output = True
client.final_output_type = "image"
client.default_sampling_params = {}
client.requires_multimodal_data = False
client.custom_process_input_func = None
client.engine_input_source = []
client._proc = None
client._proc_manager = None
client._owns_process = False
client._zmq_ctx = zmq.Context()
client._request_socket = client._zmq_ctx.socket(zmq.PUSH)
client._request_socket.bind(request_address)
client._response_socket = client._zmq_ctx.socket(zmq.PULL)
client._response_socket.bind(response_address)
client._encoder = OmniMsgpackEncoder()
client._decoder = OmniMsgpackDecoder()
client._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
client._rpc_results = {}
client._pending_rpcs = set()
client._tasks = {}
client._shutting_down = False
client._engine_dead = False
return client
@staticmethod
def _close_client(client: StageDiffusionClient) -> None:
client._request_socket.close(linger=0)
client._response_socket.close(linger=0)
client._zmq_ctx.term()
class TestPipelineStreamingOutputToEntrypoint:
"""Streaming pipeline output through inline stage, orchestrator, and entrypoints."""
@pytest.mark.asyncio
async def test_streaming_output_reaches_async_omni_from_mock_pipeline(self) -> None:
"""Mock pipeline streaming chunks reach AsyncOmni.generate() via inline stage and orchestrator."""
pipeline = _StepStreamingPipeline(
[
DiffusionOutput(custom_output={"chunk": 0}, finished=False),
DiffusionOutput(custom_output={"chunk": 1}, finished=True),
]
)
inline_client = self._make_inline_pipeline_client(pipeline)
fixture = _build_harness([inline_client])
omni = self._make_async_omni(self._OrchestratorBridgeEngine(fixture))
try:
outputs: list[OmniRequestOutput] = []
async for output in omni.generate(
prompt={"prompt": "a cat"},
request_id="req-omni",
sampling_params_list=[OmniDiffusionSamplingParams()],
output_modalities=["image"],
):
outputs.append(output)
await _wait_for(lambda: len(pipeline.requests) == 1)
assert pipeline.requests[0].request_id.startswith("req-omni-")
assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
assert [getattr(output.request_output, "finished", output.finished) for output in outputs] == [False, True]
finally:
await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client)
@pytest.mark.asyncio
async def test_streaming_output_reaches_api_websocket_from_mock_pipeline(self, mocker: MockerFixture) -> None:
"""Mock pipeline streaming chunks reach the video WebSocket handler through AsyncOmni.generate()."""
frames = [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(2)]
pipeline = _StepStreamingPipeline(
[
DiffusionOutput(custom_output={"chunk": 0, "images": [frames]}, finished=False),
DiffusionOutput(custom_output={"chunk": 1, "images": [frames]}, finished=True),
]
)
inline_client = self._make_inline_pipeline_client(pipeline)
fixture = _build_harness([inline_client])
bridge_engine = self._OrchestratorBridgeEngine(fixture)
omni = self._make_async_omni(bridge_engine)
encoded = [b"pipeline-fmp4-0", b"pipeline-fmp4-1"]
class FakeStreamingVideoEncoder:
encode_calls: list[int] = []
def encode(self, video):
del video
idx = len(self.encode_calls)
self.encode_calls.append(idx)
return encoded[idx]
def close(self):
return b""
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
lambda *, output_format, fps, video_codec_options=None: FakeStreamingVideoEncoder(),
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
return_value="diffusion",
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
return_value=[OmniDiffusionSamplingParams()],
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list",
return_value=[OmniDiffusionSamplingParams()],
)
handler = OmniStreamingVideoOutputHandler(
engine_client=omni,
model_name="test-model",
stage_configs=bridge_engine.stage_configs,
stall_timeout=5.0,
start_timeout=5.0,
)
app = FastAPI()
@app.websocket("/v1/realtime/video")
async def ws_endpoint(websocket: WebSocket):
await handler.handle_session(websocket)
try:
with TestClient(app) as client:
with client.websocket_connect("/v1/realtime/video") as ws:
ws.send_json({"type": "session.start", "prompt": "integration test"})
assert ws.receive_json()["type"] == "video.start"
assert ws.receive_bytes() == b"pipeline-fmp4-0"
assert ws.receive_bytes() == b"pipeline-fmp4-1"
done = ws.receive_json()
assert done["type"] == "session.done"
assert done["chunks"] == 2
await _wait_for(lambda: len(pipeline.requests) == 1)
assert pipeline.requests[0].request_id
finally:
await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client)
class _OrchestratorBridgeEngine:
"""Minimal AsyncOmni engine facade backed by a live Orchestrator harness."""
def __init__(self, fixture: OrchestratorFixture) -> None:
self._fixture = fixture
self.stage_metadata = [
StageRuntimeInfo(
stage_type="diffusion",
final_output=True,
final_output_type="image",
)
]
self.stage_configs: list[Any] = [SimpleNamespace(stage_type="diffusion")]
self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
self.num_stages = 1
self.supported_tasks = ("generate",)
self._alive = True
async def add_request_async(
self,
*,
request_id: str,
prompt: Any,
sampling_params_list: list[Any],
final_stage_id: int,
**kwargs: Any,
) -> None:
self._fixture.request_sync_q.put_nowait(
StageSubmissionMessage(
type="add_request",
request_id=request_id,
prompt=prompt,
original_prompt=prompt,
output_prompt_text=None,
sampling_params_list=sampling_params_list,
final_stage_id=final_stage_id,
final_output_stage_ids=kwargs.get("final_output_stage_ids"),
preprocess_ms=0.0,
request_timestamp=kwargs.get("arrival_time", time.time()),
enqueue_ts=time.perf_counter(),
)
)
async def try_get_output_async(self) -> Any | None:
try:
return self._fixture.output_sync_q.get_nowait()
except queue.Empty:
return None
def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo:
return self.stage_metadata[stage_id]
def is_alive(self) -> bool:
return self._fixture.thread.is_alive()
async def abort_async(self, request_ids: list[str]) -> None:
del request_ids
@classmethod
def _make_inline_pipeline_client(cls, pipeline: _StepStreamingPipeline) -> InlineStageDiffusionClient:
"""Build an inline diffusion stage client whose engine runs the given mock pipeline."""
metadata = StageMetadata(
stage_id=0,
stage_type="diffusion",
engine_output_type="image",
is_comprehension=False,
requires_multimodal_data=False,
engine_input_source=[],
final_output=True,
final_output_type="image",
default_sampling_params=OmniDiffusionSamplingParams(),
custom_process_input_func=None,
model_stage=None,
runtime_cfg=None,
)
pipeline_engine = _PipelineBackedEngine(pipeline)
with patch.object(InlineStageDiffusionClient, "_enrich_config"):
with patch(
"vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine.make_engine",
return_value=pipeline_engine,
):
od_config = MagicMock(spec=OmniDiffusionConfig)
od_config.streaming_output = True
return InlineStageDiffusionClient(
model="test_model",
od_config=od_config,
metadata=metadata,
batch_size=1,
)
@staticmethod
def _make_async_omni(engine: _OrchestratorBridgeEngine) -> AsyncOmni:
omni = object.__new__(AsyncOmni)
omni.engine = engine
omni.log_stats = False
omni._pause_cond = asyncio.Condition()
omni._paused = False
omni.request_states = {}
omni.final_output_task = None
omni.event_resolver = AsyncEventResolver()
omni._enable_ar_profiler = False
omni._is_sleeping = False
omni.prom_metrics = MagicMock()
omni.mod_metrics = MagicMock()
omni.resolve_sampling_params_list = lambda params, allow_delta_coercion: params
omni._compute_final_stage_id = lambda output_modalities: 0
omni._compute_final_output_stage_ids = lambda output_modalities: [0]
omni.default_sampling_params_list = engine.default_sampling_params_list
omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None)
return omni
@staticmethod
async def _shutdown_pipeline_omni_harness(
omni: AsyncOmni,
fixture: OrchestratorFixture,
inline_client: InlineStageDiffusionClient,
) -> None:
if omni.final_output_task is not None:
omni.final_output_task.cancel()
await asyncio.gather(omni.final_output_task, return_exceptions=True)
inline_client.shutdown()
fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
await asyncio.to_thread(fixture.thread.join, 5)
def _make_vllm_config():
@contextlib.contextmanager
def set_priority(*args, **kwargs):
del args, kwargs
yield
return SimpleNamespace(
kernel_config=SimpleNamespace(ir_op_priority=SimpleNamespace(set_priority=set_priority)),
compilation_config=SimpleNamespace(ir_enable_torch_wrap=True),
)
class TestSupportedPipelines:
"""Streaming-output protocol checks for supported pipelines."""
def test_helios_supports_step_execution_for_streaming_output(self) -> None:
from vllm_omni.diffusion.models.helios.pipeline_helios import HeliosPipeline
from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
# Avoid loading model weights; protocol membership depends on the class contract.
pipeline = object.__new__(HeliosPipeline)
assert pipeline.supports_step_execution is True
assert supports_step_execution(pipeline) is True
assert isinstance(pipeline, SupportsStepExecution) is True
def test_load_model_rejects_streaming_output_without_step_execution(self, monkeypatch) -> None:
class _NoStepPipeline:
def forward(self): ...
class _FakeDiffusersPipelineLoader:
def __init__(self, *args, **kwargs): ...
def load_model(self, **kwargs):
return _NoStepPipeline()
class _FakeDeviceMemoryProfiler:
consumed_memory = 0
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb):
return False
runner = object.__new__(DiffusionModelRunner)
runner.vllm_config = _make_vllm_config()
runner.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
enforce_eager=True,
cache_backend=None,
cache_config=None,
streaming_output=True,
step_execution=False,
model_class_name="NoStepPipeline",
parallel_config=SimpleNamespace(use_hsdp=False),
)
runner.device = torch.device("cpu")
runner.pipeline = None
runner.cache_backend = None
runner.offload_backend = None
runner.kv_transfer_manager = SimpleNamespace()
monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeDiffusersPipelineLoader)
monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeDeviceMemoryProfiler)
monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None)
monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None)
with pytest.raises(ValueError, match="NoStepPipeline"):
DiffusionModelRunner.load_model(runner)
@@ -46,6 +46,7 @@ def client(mock_engine):
)
with patch.object(InlineStageDiffusionClient, "_enrich_config"):
od_config = MagicMock(spec=OmniDiffusionConfig)
od_config.streaming_output = False
c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1)
yield c
c.shutdown()
@@ -72,6 +73,35 @@ async def test_inline_dispatch_request_success(client, mock_engine):
mock_engine.step.assert_called_once()
@pytest.mark.asyncio
async def test_inline_dispatch_request_streaming_success(client, mock_engine):
"""Inline StageDiffusionClient correctly receives streaming chunk output from Diffusion Engine"""
chunks = [
OmniRequestOutput.from_diffusion(request_id="req-stream", images=[], finished=False),
OmniRequestOutput.from_diffusion(request_id="req-stream", images=[], finished=True),
]
async def _step_streaming(_request):
for chunk in chunks:
yield [chunk]
client.od_config.streaming_output = True
mock_engine.step_streaming = _step_streaming
await client.add_request_async("req-stream", "A test prompt", OmniDiffusionSamplingParams())
outputs = []
for _ in range(20):
output = client.get_diffusion_output_nowait()
if output is not None:
outputs.append(output)
if output.finished:
break
await asyncio.sleep(0.01)
assert outputs == chunks
@pytest.mark.asyncio
async def test_inline_dispatch_request_error(client, mock_engine):
# Setup mock engine step to raise an exception
@@ -21,6 +21,7 @@ from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE
from vllm_omni.diffusion.sched import RequestScheduler
from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc
from vllm_omni.diffusion.worker.utils import RunnerOutput
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
@@ -49,7 +50,7 @@ def _make_executor(num_gpus: int = 1):
Returns ``(executor, request_queue, result_queue)``.
"""
od_cfg = SimpleNamespace(num_gpus=num_gpus)
od_cfg = SimpleNamespace(num_gpus=num_gpus, streaming_output=False)
monkeypatch = pytest.MonkeyPatch()
monkeypatch.setattr(MultiprocDiffusionExecutor, "_init_executor", lambda self: None)
executor = MultiprocDiffusionExecutor(od_cfg)
@@ -75,6 +76,7 @@ def _make_engine(num_gpus: int = 1):
"""Create a lightweight ``DiffusionEngine`` wired to mocked executor."""
executor, req_q, res_q = _make_executor(num_gpus)
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.od_config = SimpleNamespace(streaming_output=False)
sched = RequestScheduler()
sched.initialize(SimpleNamespace())
engine.scheduler = sched
@@ -581,6 +583,58 @@ class TestMultiprocExecutorRaisesEngineDeadError:
)
class TestMultiprocExecutorStepStreamingOutput:
"""Streaming output uses step execution and one worker reply per step."""
def test_execute_step_allows_streaming_output_mode(self):
executor, req_q, res_q = _make_executor()
executor.od_config = SimpleNamespace(streaming_output=True) # pyright: ignore[reportAttributeAccessIssue]
runner_outputs = [
RunnerOutput(
request_id="sched-stream",
step_index=1,
finished=False,
result=DiffusionOutput(output={"chunk": 0}, finished=False, chunk_index=0, total_chunks=2),
),
RunnerOutput(
request_id="sched-stream",
step_index=2,
finished=True,
result=DiffusionOutput(output={"chunk": 1}, finished=True, chunk_index=1, total_chunks=2),
),
]
scheduler_output = SimpleNamespace(
scheduled_request_ids=["sched-stream"],
)
def _worker():
for runner_output in runner_outputs:
req_q.get(timeout=10)
res_q.put(runner_output)
thread = threading.Thread(target=_worker, daemon=True)
thread.start()
first: RunnerOutput = MultiprocDiffusionExecutor.execute_step( # pyright: ignore[reportAssignmentType]
executor,
scheduler_output, # pyright: ignore[reportArgumentType]
)
second: RunnerOutput = MultiprocDiffusionExecutor.execute_step( # pyright: ignore[reportAssignmentType]
executor,
scheduler_output, # pyright: ignore[reportArgumentType]
)
assert first is runner_outputs[0]
assert first.result is not None
assert first.result.output == {"chunk": 0}
assert first.finished is False
assert second is runner_outputs[1]
assert second.result is not None
assert second.result.output == {"chunk": 1}
assert second.finished is True
thread.join(timeout=2)
class TestDiffusionEngineDeadErrorPassthrough:
"""``DiffusionEngine.add_req_and_wait_for_response`` re-raises
``EngineDeadError`` from executor and wraps other errors."""
@@ -11,6 +11,7 @@ import pytest
from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -54,6 +55,7 @@ def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
proc = object.__new__(StageDiffusionProc)
proc._engine = SimpleNamespace(step=step)
proc._od_config = SimpleNamespace(streaming_output=False)
proc._executor = ThreadPoolExecutor(max_workers=1)
try:
@@ -102,6 +104,40 @@ class MockDiffusionEngine:
return [MockOmniRequestOutput(request_id=request.request_id)]
@pytest.mark.asyncio
async def test_proc_streaming_request_yields_each_engine_chunk():
"""Ensure that the streaming output chunks from DiffusionEngine reaches StageDiffusionProc"""
captured = {}
chunks = [
OmniRequestOutput.from_diffusion(request_id="", images=[], finished=False),
OmniRequestOutput.from_diffusion(request_id="", images=[], finished=True),
]
class _StreamingEngine:
async def step_streaming(self, request):
captured["request"] = request
for chunk in chunks:
yield [chunk]
stage_proc = object.__new__(StageDiffusionProc)
stage_proc._engine = _StreamingEngine()
outputs = [
output
async for output in stage_proc._process_streaming_request(
request_id="req-stream",
prompt="prompt",
sampling_params_dict=asdict(OmniDiffusionSamplingParams()),
kv_sender_info={0: {"host": "127.0.0.1"}},
)
]
assert outputs == chunks
assert [output.request_id for output in outputs] == ["req-stream", "req-stream"]
assert [output.finished for output in outputs] == [False, True]
assert captured["request"].kv_sender_info == {0: {"host": "127.0.0.1"}}
@pytest.mark.asyncio
async def test_proc_process_request_with_batching_async_output():
stage_proc = object.__new__(StageDiffusionProc)
+2 -1
View File
@@ -305,7 +305,7 @@ def assert_video_similarity_metrics(
offline_path: Path,
ssim_threshold: float,
psnr_threshold: float,
) -> None:
) -> tuple[float, float]:
ssim_output = run_ffmpeg_similarity("ssim", online_path, offline_path)
psnr_output = run_ffmpeg_similarity("psnr", online_path, offline_path)
ssim_score = parse_ssim_score(ssim_output)
@@ -339,6 +339,7 @@ def assert_video_similarity_metrics(
f"PSNR below threshold: got {psnr_score:.6f}, expected >= {psnr_threshold:.6f}. "
f"online={online_path} offline={offline_path}"
)
return ssim_score, psnr_score
def assert_similarity(
@@ -33,7 +33,7 @@ Prerequisites
* **Seed-TTS** optional extras for WER/SIM/UTMOS::
pip install 'vllm-omni[seed-tts-eval]'
pip install 'vllm-omni[dev]'
Examples
--------
@@ -0,0 +1,143 @@
"""
Online serving smoke test for streaming video output similarity.
"""
import json
from pathlib import Path
import pytest
from tests.e2e.accuracy.helpers import assert_video_similarity_metrics, probe_video
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
MODEL = "BestWishYsh/Helios-Distilled"
PROMPT = "A serene lakeside sunrise with mist over the water."
WIDTH = 640
HEIGHT = 384
FPS = 16
NUM_FRAMES = 65
NUM_INFERENCE_STEPS = 5
SEED = 42
MIN_PSNR_DB = 28.0
MIN_SSIM = 0.95
SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
HELIOS_DISTILLED_EXTRA_PARAMS = {
"is_enable_stage2": True,
"pyramid_num_stages": 3,
"pyramid_num_inference_steps_list": [1, 1, 1],
"is_amplify_first_chunk": True,
}
def _server_args(*, streaming_output: bool) -> list[str]:
args = [
"--stage-init-timeout",
"600",
"--init-timeout",
"900",
"--log-stats",
]
if streaming_output:
args.append("--diffusion-streaming-output")
return args
def _request_config(model: str, *, serialize_extra_params: bool) -> dict:
extra_params = (
json.dumps(HELIOS_DISTILLED_EXTRA_PARAMS) if serialize_extra_params else HELIOS_DISTILLED_EXTRA_PARAMS
)
return {
"model": model,
"form_data": {
"prompt": PROMPT,
"width": WIDTH,
"height": HEIGHT,
"fps": FPS,
"num_frames": NUM_FRAMES,
"num_inference_steps": NUM_INFERENCE_STEPS,
"guidance_scale": 1.0,
"seed": SEED,
"extra_params": extra_params,
},
}
def _client_for(server: OmniServer, *, run_level: str) -> OpenAIClientHandler:
return OpenAIClientHandler(
host=server.host,
port=server.port,
api_key="EMPTY",
run_level=run_level,
log_stats=server.log_stats,
)
def _write_video_artifact(response: DiffusionResponse, output_path: Path) -> Path:
assert response.videos, "No video artifact returned"
output_path.write_bytes(response.videos[0])
return output_path
@pytest.mark.parametrize(
"model",
[
pytest.param(
MODEL,
id="helios_distilled",
marks=SINGLE_CARD_MARKS,
)
],
)
def test_helios_streaming_output_matches_non_streaming(
model: str,
model_prefix: str,
run_level: str,
tmp_path: Path,
record_property,
):
model = model_prefix + model
with OmniServer(model, _server_args(streaming_output=True), use_omni=True) as server:
client = _client_for(server, run_level=run_level)
streaming_response = client.send_streaming_video_diffusion_request(
_request_config(server.model, serialize_extra_params=False)
)[0]
streaming_path = _write_video_artifact(streaming_response, tmp_path / "helios_streaming.mp4")
record_property("helios_streaming_video", str(streaming_path))
record_property("helios_streaming_latency_s", streaming_response.e2e_latency)
with OmniServer(model, _server_args(streaming_output=False), use_omni=True) as server:
client = _client_for(server, run_level=run_level)
non_streaming_response = client.send_video_diffusion_request(
_request_config(server.model, serialize_extra_params=True)
)[0]
non_streaming_path = _write_video_artifact(non_streaming_response, tmp_path / "helios_non_streaming.mp4")
record_property("helios_non_streaming_video", str(non_streaming_path))
record_property("helios_non_streaming_latency_s", non_streaming_response.e2e_latency)
streaming_metadata = probe_video(streaming_path)
non_streaming_metadata = probe_video(non_streaming_path)
assert streaming_metadata == non_streaming_metadata, (
f"Video metadata mismatch for helios_streaming_output:\n"
f"streaming={streaming_metadata}\n"
f"non_streaming={non_streaming_metadata}\n"
f"streaming_path={streaming_path}\n"
f"non_streaming_path={non_streaming_path}"
)
ssim, psnr = assert_video_similarity_metrics(
label="helios_streaming_output",
online_path=streaming_path,
offline_path=non_streaming_path,
ssim_threshold=MIN_SSIM,
psnr_threshold=MIN_PSNR_DB,
)
record_property("helios_streaming_vs_non_streaming_psnr_db", psnr)
record_property("helios_streaming_vs_non_streaming_ssim_percent", ssim * 100.0)
@@ -0,0 +1,96 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""E2E offline inference tests for IndexTTS2 two-stage pipeline.
Stage 0 (GPT AR) → Stage 1 (S2Mel + BigVGAN). Output is 22050 Hz mono WAV.
Covers: basic TTS.
"""
from __future__ import annotations
from pathlib import Path
import pytest
import torch
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni import Omni
from vllm_omni.model_executor.models.indextts2.prompt_utils import (
build_indextts2_prefill_prompt_ids,
)
MODEL_NAME = "IndexTeam/IndexTTS-2"
STAGE_CONFIG = get_deploy_config_path("indextts2.yaml")
_OMNI_RUNNER_PARAM = (
MODEL_NAME,
STAGE_CONFIG,
)
ASSETS_DIR = Path(__file__).resolve().parents[2] / "assets" / "indextts2"
REF_AUDIO = str(ASSETS_DIR / "ref_audio.wav")
pytestmark = [
pytest.mark.advanced_model,
pytest.mark.tts,
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
]
SAMPLE_RATE = 22050
def _audio_from_mm(mm: dict) -> torch.Tensor | None:
audio = mm.get("audio")
if audio is None:
audio = mm.get("model_outputs")
if isinstance(audio, list):
chunks = [chunk.reshape(-1) for chunk in audio if isinstance(chunk, torch.Tensor) and chunk.numel() > 0]
return torch.cat(chunks, dim=0) if chunks else None
return audio if isinstance(audio, torch.Tensor) else None
def _sample_rate_from_mm(mm: dict) -> int:
sr = mm.get("sr")
if isinstance(sr, list):
sr = sr[-1] if sr else None
if hasattr(sr, "item"):
return int(sr.item())
return int(sr) if sr is not None else SAMPLE_RATE
def _build_request(text: str, ref_audio: str, **extra) -> dict:
info = {
"text": [text],
"voice": [ref_audio],
}
info.update(extra)
return {
"prompt_token_ids": build_indextts2_prefill_prompt_ids(MODEL_NAME, text),
"additional_information": info,
}
def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
for omni_out in omni.generate(request):
mm = omni_out.multimodal_output
assert mm is not None, "Expected multimodal_output"
audio = _audio_from_mm(mm)
assert audio is not None, "Expected audio output"
return audio.cpu(), _sample_rate_from_mm(mm)
raise AssertionError("No outputs received")
def _assert_valid_audio(audio: torch.Tensor, sr: int) -> None:
assert sr == SAMPLE_RATE
assert audio.numel() > 0
assert torch.any(audio != 0).item()
@hardware_test(res={"cuda": "L4"})
def test_basic_english(omni_runner: OmniRunner) -> None:
"""English TTS: text + ref_audio, no emotion conditioning."""
req = _build_request("Hello, this is a voice synthesis test.", REF_AUDIO)
audio, sr = _collect_audio(omni_runner.omni, req)
_assert_valid_audio(audio, sr)
+17 -1
View File
@@ -14,7 +14,7 @@ from tests.helpers.stage_config import get_deploy_config_path, modify_stage_conf
from vllm_omni.platforms import current_omni_platform
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
thinker_only_models = ["Qwen/Qwen3-Omni-30B-A3B-Captioner"]
# Single CI deploy YAML; rocm/xpu deltas are picked automatically via the
# platforms: section. Only CUDA needs an extra enforce_eager tweak.
@@ -40,6 +40,10 @@ else:
# Create parameter combinations for model and stage config
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
# we can use the same config for a model that only has thinker (i.e., does not
# enable audio output) because the resolver should figure out that it doesn't
# need the full pipeline based on the HF config.
thinker_test_params = [(model, stage_config) for model in thinker_only_models for stage_config in stage_configs]
def get_question(prompt_type="video"):
@@ -61,3 +65,15 @@ def test_video_to_audio(omni_runner, omni_runner_handler) -> None:
# Test single completion
omni_runner_handler.send_omni_request(request_config)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
@pytest.mark.parametrize("omni_runner", thinker_test_params, indirect=True)
def test_thinker_only_model_request(omni_runner, omni_runner_handler) -> None:
"""Test that we can load and run a request through a model that only has the thinker stage."""
request_config = {"prompts": "what color is the sky?", "modalities": ["text"]}
# Test single completion
omni_runner_handler.send_omni_request(request_config)
@@ -0,0 +1,241 @@
"""E2E offline inference tests for SoulX-Singer (single-stage, preprocess inline)."""
import functools
import importlib
import json
import os
import shutil
from pathlib import Path
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
import numpy as np
import pytest
from tests.helpers.mark import hardware_test
from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
PROMPT_AUDIO = get_asset_path("soulxsinger/zh_prompt.mp3")
TARGET_AUDIO = get_asset_path("soulxsinger/music.mp3")
PROMPT_METADATA = get_asset_path("soulxsinger/zh_prompt.json")
TARGET_METADATA = get_asset_path("soulxsinger/music.json")
SAMPLE_RATE = 24_000
# phone_set.json (phoneme vocab) is not shipped on HuggingFace; SVS loads it
# from the model dir. Stage it from this pinned upstream commit so the SVS DiT
# can run in CI; the SVS tests skip gracefully if it cannot be fetched.
_PHONE_SET_URL = (
"https://raw.githubusercontent.com/Soul-AILab/SoulX-Singer/"
"81aeb3ae772c70093c3de74dc23c92d983801ae4/soulxsinger/utils/phoneme/phone_set.json"
)
if not PROMPT_AUDIO.is_file() or not TARGET_AUDIO.is_file():
pytest.skip(
f"Missing SoulX-Singer audio assets: {PROMPT_AUDIO.name}, {TARGET_AUDIO.name}",
allow_module_level=True,
)
pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.tts]
_CASES = (
pytest.param(
"SoulXSingerPipeline",
"soulxsinger_svs.yaml",
{
"language": "Mandarin",
"vocal_sep": False,
"control": "score",
"auto_shift": False,
"pitch_shift": 0,
},
("g2pM", "g2p_en"),
id="svs",
),
pytest.param(
"SoulXSingerSVCPipeline",
"soulxsinger_svc.yaml",
{"vocal_sep": False, "auto_shift": False, "pitch_shift": 0},
(),
id="svc",
),
)
@functools.lru_cache(maxsize=1)
def _resolve_weights() -> tuple[Path, Path, Path]:
from huggingface_hub import snapshot_download
base = Path(snapshot_download("Soul-AILab/SoulX-Singer", allow_patterns=["*"]))
# phone_set.json is not on HF; best-effort stage it from pinned upstream so
# SVS can load. SVS tests skip (not fail) if both this and a manual copy are absent.
phone_set = base / "phoneme" / "phone_set.json"
if not phone_set.is_file() and not (base / "phone_set.json").is_file():
try:
from urllib.request import urlretrieve
phone_set.parent.mkdir(exist_ok=True)
urlretrieve(_PHONE_SET_URL, phone_set)
except Exception:
pass
# make a temporary svc directory
svc_dir = base.parent / "SoulX-Singer-SVC"
svc_dir.mkdir(parents=True, exist_ok=True)
shutil.copyfile(base / "config.yaml", svc_dir / "config.yaml")
shutil.move(base / "model-svc.pt", svc_dir / "model-svc.pt")
if raw := os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR"):
pre = Path(raw).expanduser().resolve()
if (pre / "rmvpe" / "rmvpe.pt").is_file():
return base, svc_dir, pre
from huggingface_hub import snapshot_download
pre = Path(snapshot_download("Soul-AILab/SoulX-Singer-Preprocess", allow_patterns=["*"]))
return base, svc_dir, pre
@pytest.fixture(scope="session")
def soulx_weights() -> tuple[Path, Path, Path]:
try:
return _resolve_weights()
except Exception as exc:
pytest.skip(f"Set SOULXSINGER_MODEL_DIR / SOULX_PREPROCESS_WEIGHTS_DIR. ({exc})")
def _flatten_audio(audio_val) -> np.ndarray:
import torch
if isinstance(audio_val, list):
chunks = [c.detach().cpu().float().numpy().reshape(-1) for c in audio_val if c is not None]
return np.concatenate(chunks) if chunks else np.array([], dtype=np.float32)
if isinstance(audio_val, torch.Tensor):
return audio_val.detach().cpu().float().numpy().reshape(-1)
return np.asarray(audio_val, dtype=np.float32).reshape(-1)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("architecture,deploy_yaml,extra_args,py_deps", _CASES)
def test_soulxsinger_multistage_from_audio(
soulx_weights: tuple[Path, Path, Path],
architecture: str,
deploy_yaml: str,
extra_args: dict,
py_deps: tuple[str, ...],
) -> None:
for mod in py_deps:
try:
importlib.import_module(mod)
except ImportError as exc:
pytest.fail(f"SoulX SVS requires {mod}: {exc}")
base_dir, svc_dir, preprocess_dir = soulx_weights
# SVS mode requires phone_set.json in the model directory
if architecture == "SoulXSingerPipeline":
json.dump(
{"model_type": "soulxsinger", "architectures": ["SoulXSingerPipeline"], "max_num_seqs": 1},
open(base_dir / "config.json", "w"),
indent=4,
)
model = str(base_dir)
if not (base_dir / "phoneme" / "phone_set.json").is_file() and not (base_dir / "phone_set.json").is_file():
pytest.skip(
"SoulX-Singer SVS test requires phoneme/phone_set.json. "
"Copy it from github.com/Soul-AILab/SoulX-Singer into the model dir. "
"See `examples/offline_inference/text_to_speech/README.md` for details."
)
if architecture == "SoulXSingerSVCPipeline":
json.dump(
{"model_type": "soulxsinger", "architectures": ["SoulXSingerSVCPipeline"], "max_num_seqs": 1},
open(svc_dir / "config.json", "w"),
indent=4,
)
model = str(svc_dir)
with OmniRunner(
model,
stage_configs_path=get_deploy_config_path(deploy_yaml),
async_chunk=False,
) as runner:
sampling = OmniDiffusionSamplingParams(
num_inference_steps=4,
guidance_scale=3.0,
seed=42,
extra_args={
"prompt_audio": str(PROMPT_AUDIO),
"target_audio": str(TARGET_AUDIO),
"preprocess_weights_dir": str(preprocess_dir),
**extra_args,
},
)
prompt = {"prompt_token_ids": [0]}
outputs = runner.generate([prompt], sampling)
assert outputs and outputs[0].error is None, outputs[0].error if outputs else "no output"
mm = outputs[0].multimodal_output
assert isinstance(mm, dict) and "audio" in mm
audio = _flatten_audio(mm["audio"])
assert 12_000 <= audio.size
assert np.isfinite(audio).all() and float(np.max(np.abs(audio))) > 1e-4
duration_s = audio.size / SAMPLE_RATE
assert 50.0 <= duration_s <= 52.0, f"duration={duration_s:.1f}s"
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_soulxsinger_svs_precomputed(soulx_weights: tuple[Path, Path, Path]) -> None:
"""SVS from precomputed note/lyric metadata.
Exercises the SVS DiT + vocoder serving path while skipping the heavy
upstream preprocess (funasr ASR + ROSVOT note transcription), so it runs in
the merge gate without those non-pip dependencies.
"""
base_dir, _svc_dir, preprocess_dir = soulx_weights
if not PROMPT_METADATA.is_file() or not TARGET_METADATA.is_file():
pytest.skip(f"Missing SoulX-Singer metadata fixtures: {PROMPT_METADATA.name}, {TARGET_METADATA.name}")
if not (base_dir / "phoneme" / "phone_set.json").is_file() and not (base_dir / "phone_set.json").is_file():
pytest.skip(
"SoulX-Singer SVS requires phoneme/phone_set.json (upstream fetch unavailable). "
"See `examples/offline_inference/text_to_speech/README.md`."
)
json.dump(
{"model_type": "soulxsinger", "architectures": ["SoulXSingerPipeline"], "max_num_seqs": 1},
open(base_dir / "config.json", "w"),
indent=4,
)
with OmniRunner(
str(base_dir),
stage_configs_path=get_deploy_config_path("soulxsinger_svs.yaml"),
async_chunk=False,
) as runner:
sampling = OmniDiffusionSamplingParams(
num_inference_steps=4,
guidance_scale=3.0,
seed=42,
extra_args={
"language": "Mandarin",
"control": "score",
"prompt_metadata_path": str(PROMPT_METADATA),
"target_metadata_path": str(TARGET_METADATA),
"audio_path": str(PROMPT_AUDIO),
"preprocess_weights_dir": str(preprocess_dir),
},
)
prompt = {"prompt_token_ids": [0]}
outputs = runner.generate([prompt], sampling)
assert outputs and outputs[0].error is None, outputs[0].error if outputs else "no output"
mm = outputs[0].multimodal_output
assert isinstance(mm, dict) and "audio" in mm
audio = _flatten_audio(mm["audio"])
assert 12_000 <= audio.size
assert np.isfinite(audio).all() and float(np.max(np.abs(audio))) > 1e-4
duration_s = audio.size / SAMPLE_RATE
assert 50.0 <= duration_s <= 52.0, f"duration={duration_s:.1f}s"
@@ -0,0 +1,82 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""E2E online serving tests for IndexTTS2 via /v1/audio/speech endpoint.
Two-stage pipeline: GPT AR → S2Mel + BigVGAN. Output is 22050 Hz mono WAV.
Covers: basic TTS and the speech endpoint's stream=True compatibility path.
"""
from __future__ import annotations
import os
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
import pytest
from tests.helpers.mark import hardware_test
from tests.helpers.media import load_test_audio_data_url
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path
pytestmark = [pytest.mark.advanced_model, pytest.mark.tts]
MODEL = "IndexTeam/IndexTTS-2"
REF_AUDIO_URL = load_test_audio_data_url("indextts2/ref_audio.wav")
tts_server_params = [
pytest.param(
OmniServerParams(
model=MODEL,
stage_config_path=get_deploy_config_path("indextts2.yaml"),
server_args=["--disable-log-stats"],
),
id="indextts2",
)
]
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
def test_basic_english(omni_server, openai_client) -> None:
"""
Basic English TTS: text + ref_audio, no emotion.
Deploy Setting: default yaml
Input Modal: text + reference audio
Output Modal: audio (22050 Hz, WAV)
Input Setting: stream=False
"""
request_config = {
"model": omni_server.model,
"input": "Hello, this is a voice cloning demo.",
"stream": False,
"response_format": "wav",
"ref_audio": REF_AUDIO_URL,
"min_audio_bytes": 1024,
}
openai_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
def test_streaming(omni_server, openai_client) -> None:
"""
stream=True compatibility: IndexTTS2 itself is non-async-chunking, but the
OpenAI speech endpoint should still accept the streaming request path and
return a valid PCM response.
Deploy Setting: default yaml
Input Modal: text + reference audio
Output Modal: audio (22050 Hz, PCM)
Input Setting: stream=True
"""
request_config = {
"model": omni_server.model,
"input": "Hello, this is a voice cloning demo.",
"stream": True,
"response_format": "pcm",
"ref_audio": REF_AUDIO_URL,
"min_audio_bytes": 1024,
}
openai_client.send_audio_speech_request(request_config)
@@ -0,0 +1,95 @@
# 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 LTX2TwoStagesPipeline (text-to-video)
- HSDP with LTX2ImageToVideoTwoStagesPipeline (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 = []
# T2V: LTX2TwoStagesPipeline
cases.append(
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
*HSDP_ARGS,
"--model-class-name",
"LTX2TwoStagesPipeline",
],
),
id="t2v_hsdp",
marks=PARALLEL_MARKS,
)
)
# I2V: LTX2ImageToVideoTwoStagesPipeline
cases.append(
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
*HSDP_ARGS,
"--model-class-name",
"LTX2ImageToVideoTwoStagesPipeline",
],
),
id="i2v_hsdp",
marks=PARALLEL_MARKS,
)
)
return cases
@pytest.mark.parametrize("omni_server", _cases(), indirect=True)
def test_ltx2_two_stage_hsdp(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
):
is_i2v = any("ImageToVideo" in arg for arg in omni_server.serve_args)
# The two-stage pipeline generates at the requested resolution then 2x
# upsamples via LTX2LatentUpsamplerModel, so output dimensions differ
# from the request. Omit height/width from form_data so the assertion
# helper skips the dimension check (it only asserts when the key is
# present). The pipeline falls back to its own defaults.
form_data = {
"prompt": PROMPT,
"model": omni_server.model,
"num_frames": 9,
"fps": 8,
"num_inference_steps": 2,
"guidance_scale": 1.0,
"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)
+2 -2
View File
@@ -9,13 +9,13 @@ from vllm.inputs import PromptType
from tests.helpers.mark import hardware_test
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.config.stage_config import _PIPELINE_REGISTRY
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
from vllm_omni.entrypoints.async_omni import AsyncOmni
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
SEED = 42
pipeline = _PIPELINE_REGISTRY["qwen2_5_omni_thinker_only"]
pipeline = OMNI_PIPELINES["qwen2_5_omni_thinker_only"]
# Single-stage thinker-only deploy, materialized from tests.helpers.stage_config._CI_OVERLAYS.
stage_config = get_deploy_config_path("ci/qwen2_5_omni_thinker_only.yaml")
+61
View File
@@ -523,6 +523,67 @@ async def test_run_single_stage_diffusion(orchestrator_factory) -> None:
await _shutdown_orchestrator(orchestrator_fixture)
@pytest.mark.asyncio
async def test_run_single_stage_diffusion_streaming_forwards_intermediate_chunks(orchestrator_factory) -> None:
"""Intermediate diffusion chunks (finished=False) reach the frontend before the final chunk."""
stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
orchestrator_fixture = orchestrator_factory([stage0])
params = OmniDiffusionSamplingParams()
try:
await _enqueue_add_request(
orchestrator_fixture,
request_id="req-stream",
prompt={"prompt": "draw a cat"},
original_prompt={"prompt": "draw a cat"},
sampling_params_list=[params],
final_stage_id=0,
)
await _wait_for(lambda: len(stage0.add_request_calls) == 1)
stage0.push_diffusion_output(
OmniRequestOutput.from_diffusion(
request_id="req-stream",
images=[],
final_output_type="image",
custom_output={"chunk": 0},
finished=False,
)
)
stage0.push_diffusion_output(
OmniRequestOutput.from_diffusion(
request_id="req-stream",
images=[],
final_output_type="image",
custom_output={"chunk": 1},
finished=True,
)
)
output_msgs: list[OutputMessage] = []
deadline = time.monotonic() + 2.0
while not output_msgs or not output_msgs[-1].finished:
if time.monotonic() >= deadline:
raise AssertionError(
f"Timed out waiting for finished orchestrator output, got {len(output_msgs)} message(s)"
)
try:
msg = orchestrator_fixture.output_sync_q.get_nowait()
except queue.Empty:
await asyncio.sleep(0.01)
continue
if isinstance(msg, OutputMessage):
output_msgs.append(msg)
assert [msg.request_id for msg in output_msgs] == ["req-stream", "req-stream"]
assert [msg.finished for msg in output_msgs] == [False, True]
assert [msg.engine_outputs.finished for msg in output_msgs] == [False, True]
assert [msg.engine_outputs.custom_output["chunk"] for msg in output_msgs] == [0, 1]
await _wait_for(lambda: "req-stream" not in orchestrator_fixture.orchestrator.request_states)
finally:
await _shutdown_orchestrator(orchestrator_fixture)
@pytest.mark.asyncio
async def test_run_llm_to_diffusion(orchestrator_factory) -> None:
stage0 = FakeStageClient(stage_type="llm", final_output=False)
+20
View File
@@ -4,6 +4,7 @@
from unittest.mock import MagicMock
import numpy as np
import pytest
import torch
from vllm.outputs import PoolingRequestOutput
@@ -250,6 +251,25 @@ def test_no_detokenizer_make_request_output():
assert AUDIO in result.outputs[0].multimodal_output
def test_no_detokenizer_make_request_output_with_routed_experts():
"""make_request_output accepts the routed_experts arg that the multimodal
output channel (_process_mm_only_outputs) passes positionally, and attaches
it to the completion output.
Regression: the call site passes 6 positional args
(..., kv_transfer_params, routed_experts) but the override previously took
only 5, raising TypeError on every generation-stage output.
"""
s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
routed_experts = np.zeros((2, 3), dtype=np.int32)
# Mirror the exact call shape of _process_mm_only_outputs.
result = s.make_request_output([], None, FinishReason.STOP, None, None, routed_experts)
assert result is not None
assert not isinstance(result, PoolingRequestOutput)
assert result.outputs[0].routed_experts is routed_experts
def test_no_detokenizer_stream_interval_skipped():
"""stream_interval > 1 does not crash when detokenizer is None."""
kwargs = {**_NO_DETOK_STATE_KWARGS, "stream_interval": 5}
@@ -1301,6 +1301,58 @@ def test_generate_images_rejects_model_mismatch(test_client):
assert "model mismatch" in response.json()["detail"].lower()
def test_image_file_response_format_multiple(test_client):
"""Test response_format=file with n>1 returns ZIP archive"""
response = test_client.post(
"/v1/images/generations",
json={
"prompt": "a dog",
"n": 3,
"response_format": "file",
},
)
assert response.status_code == 200
assert response.headers["content-type"] == "application/zip"
assert "attachment" in response.headers.get("content-disposition", "")
assert ".zip" in response.headers.get("content-disposition", "")
# Verify it's a valid ZIP with 3 PNG files
import zipfile
zip_buffer = io.BytesIO(response.content)
with zipfile.ZipFile(zip_buffer, "r") as zf:
files = zf.namelist()
assert len(files) == 3
assert all(f.endswith(".png") for f in files)
# Verify each file is a valid PNG
for filename in files:
img_bytes = zf.read(filename)
img = Image.open(io.BytesIO(img_bytes))
assert img.format == "PNG"
def test_image_file_response_format_single(test_client):
"""Test response_format=file with n=1 returns a single image file."""
response = test_client.post(
"/v1/images/generations",
json={
"prompt": "a dog",
"n": 1,
"response_format": "file",
},
)
assert response.status_code == 200
assert response.headers["content-type"] == "image/png"
assert "attachment" in response.headers.get("content-disposition", "")
assert ".png" in response.headers.get("content-disposition", "")
img = Image.open(io.BytesIO(response.content))
assert img.format == "PNG"
def make_test_image_bytes(size=(64, 64)) -> bytes:
img = Image.new(
"RGB",
@@ -29,6 +29,7 @@ from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
from vllm_omni.entrypoints.openai.protocol.audio import (
BatchSpeechRequest,
CreateAudio,
OpenAICreateAudioGenerateRequest,
OpenAICreateSpeechRequest,
SpeechBatchItem,
StreamingSpeechSessionConfig,
@@ -2292,26 +2293,260 @@ def test_api_server_create_speech_wraps_error_response_status(mocker: MockerFixt
)
)
raw_request = _make_api_server_request(handler, path="/v1/audio/speech")
request = OpenAICreateSpeechRequest(input="Hello")
response = asyncio.run(api_server_module.create_speech(request, raw_request))
_assert_openai_error_response(response, status_code=400, message="bad request")
def _make_api_server_request(handler, *, method: str = "POST", path: str = "/v1/audio/voices") -> Request:
app = FastAPI()
app.state.openai_serving_speech = handler
scope = {
"type": "http",
"app": app,
"method": "POST",
"path": "/v1/audio/speech",
"method": method,
"path": path,
"headers": [],
"query_string": b"",
"client": ("127.0.0.1", 12345),
"server": ("testserver", 80),
"scheme": "http",
}
raw_request = Request(scope)
return Request(scope)
def _patch_api_server_base(mocker: MockerFixture):
def _fake_create_error_response(message, err_type="BadRequestError", status_code=400, param=None):
return ErrorResponse(
error=ErrorInfo(
message=message,
type=err_type,
param=param,
code=getattr(status_code, "value", status_code),
)
)
fake_base = mocker.MagicMock()
fake_base.create_error_response.side_effect = _fake_create_error_response
mocker.patch.object(api_server_module, "base", return_value=fake_base)
return fake_base
def _assert_openai_error_response(
response: JSONResponse,
*,
status_code: int,
message: str,
err_type: str = "BadRequestError",
) -> None:
assert isinstance(response, JSONResponse)
assert response.status_code == status_code
body = json.loads(response.body)
assert body["error"]["code"] == status_code
assert body["error"]["type"] == err_type
assert message in body["error"]["message"]
def test_api_server_list_voices_without_speech_handler_returns_404(mocker: MockerFixture):
_patch_api_server_base(mocker)
raw_request = _make_api_server_request(None, method="GET")
response = asyncio.run(api_server_module.list_voices(raw_request))
_assert_openai_error_response(
response, status_code=404, message="does not support Speech API", err_type="NotFoundError"
)
def test_api_server_upload_voice_value_error_returns_400(mocker: MockerFixture):
_patch_api_server_base(mocker)
handler = mocker.MagicMock()
handler.upload_voice = mocker.AsyncMock(side_effect=ValueError("Unsupported MIME type: audio/x-m4a"))
raw_request = _make_api_server_request(handler)
response = asyncio.run(
api_server_module.upload_voice(
raw_request,
audio_sample=mocker.MagicMock(),
speaker_embedding=None,
consent="cons_test",
name="probe",
)
)
_assert_openai_error_response(response, status_code=400, message="Unsupported MIME type")
def test_api_server_upload_voice_without_speech_handler_returns_404(mocker: MockerFixture):
_patch_api_server_base(mocker)
raw_request = _make_api_server_request(None)
response = asyncio.run(
api_server_module.upload_voice(
raw_request,
consent="cons_test",
name="probe",
)
)
_assert_openai_error_response(
response, status_code=404, message="does not support Speech API", err_type="NotFoundError"
)
def test_api_server_upload_voice_without_input_returns_400(mocker: MockerFixture):
_patch_api_server_base(mocker)
raw_request = _make_api_server_request(mocker.MagicMock())
response = asyncio.run(
api_server_module.upload_voice(
raw_request,
audio_sample=None,
speaker_embedding=None,
consent="cons_test",
name="probe",
)
)
_assert_openai_error_response(response, status_code=400, message="must be provided")
def test_api_server_upload_voice_with_audio_and_embedding_returns_400(mocker: MockerFixture):
_patch_api_server_base(mocker)
raw_request = _make_api_server_request(mocker.MagicMock())
response = asyncio.run(
api_server_module.upload_voice(
raw_request,
audio_sample=mocker.MagicMock(),
speaker_embedding="[0.1]",
consent="cons_test",
name="probe",
)
)
_assert_openai_error_response(response, status_code=400, message="mutually exclusive")
def test_api_server_upload_voice_exception_returns_500(mocker: MockerFixture):
_patch_api_server_base(mocker)
handler = mocker.MagicMock()
handler.upload_voice = mocker.AsyncMock(side_effect=RuntimeError("disk failed"))
raw_request = _make_api_server_request(handler)
response = asyncio.run(
api_server_module.upload_voice(
raw_request,
audio_sample=mocker.MagicMock(),
speaker_embedding=None,
consent="cons_test",
name="probe",
)
)
_assert_openai_error_response(
response,
status_code=500,
message="Failed to upload voice",
err_type="InternalServerError",
)
def test_api_server_delete_voice_without_speech_handler_returns_404(mocker: MockerFixture):
_patch_api_server_base(mocker)
raw_request = _make_api_server_request(None, method="DELETE", path="/v1/audio/voices/probe")
response = asyncio.run(api_server_module.delete_voice("probe", raw_request))
_assert_openai_error_response(
response, status_code=404, message="does not support Speech API", err_type="NotFoundError"
)
def test_api_server_delete_voice_value_error_returns_400(mocker: MockerFixture):
_patch_api_server_base(mocker)
handler = mocker.MagicMock()
handler.delete_voice = mocker.AsyncMock(side_effect=ValueError("Invalid voice name"))
raw_request = _make_api_server_request(handler, method="DELETE", path="/v1/audio/voices/probe")
response = asyncio.run(api_server_module.delete_voice("probe", raw_request))
_assert_openai_error_response(response, status_code=400, message="Invalid voice name")
def test_api_server_delete_voice_not_found_returns_404(mocker: MockerFixture):
_patch_api_server_base(mocker)
handler = mocker.MagicMock()
handler.delete_voice = mocker.AsyncMock(return_value=False)
raw_request = _make_api_server_request(handler, method="DELETE", path="/v1/audio/voices/missing")
response = asyncio.run(api_server_module.delete_voice("missing", raw_request))
_assert_openai_error_response(
response,
status_code=404,
message="Voice 'missing' not found",
err_type="NotFoundError",
)
def test_api_server_delete_voice_exception_returns_500(mocker: MockerFixture):
_patch_api_server_base(mocker)
handler = mocker.MagicMock()
handler.delete_voice = mocker.AsyncMock(side_effect=RuntimeError("disk failed"))
raw_request = _make_api_server_request(handler, method="DELETE", path="/v1/audio/voices/probe")
response = asyncio.run(api_server_module.delete_voice("probe", raw_request))
_assert_openai_error_response(
response,
status_code=500,
message="Failed to delete voice",
err_type="InternalServerError",
)
def test_api_server_create_speech_without_handler_returns_404(mocker: MockerFixture):
fake_base = _patch_api_server_base(mocker)
raw_request = _make_api_server_request(None, path="/v1/audio/speech")
raw_request.app.state.openai_serving_tokenization = fake_base
request = OpenAICreateSpeechRequest(input="Hello")
response = asyncio.run(api_server_module.create_speech(request, raw_request))
assert isinstance(response, JSONResponse)
assert response.status_code == 400
_assert_openai_error_response(
response, status_code=404, message="does not support Speech API", err_type="NotFoundError"
)
def test_api_server_create_speech_batch_without_handler_returns_404(mocker: MockerFixture):
fake_base = _patch_api_server_base(mocker)
raw_request = _make_api_server_request(None, path="/v1/audio/speech/batch")
raw_request.app.state.openai_serving_tokenization = fake_base
request = BatchSpeechRequest(items=[SpeechBatchItem(input="hi")])
response = asyncio.run(api_server_module.create_speech_batch(request, raw_request))
_assert_openai_error_response(
response, status_code=404, message="does not support Speech API", err_type="NotFoundError"
)
def test_api_server_create_audio_generate_without_handler_returns_404(mocker: MockerFixture):
fake_base = _patch_api_server_base(mocker)
raw_request = _make_api_server_request(None, path="/v1/audio/generate")
raw_request.app.state.openai_serving_audio_generate = None
raw_request.app.state.openai_serving_tokenization = fake_base
request = OpenAICreateAudioGenerateRequest(input="a bird singing")
response = asyncio.run(api_server_module.create_audio_generate(request, raw_request))
_assert_openai_error_response(
response, status_code=404, message="does not support Audio Generate API", err_type="NotFoundError"
)
def test_api_server_create_speech_engine_error_response_includes_request_and_stage_id(mocker: MockerFixture):
@@ -2325,26 +2560,13 @@ def test_api_server_create_speech_engine_error_response_includes_request_and_sta
terminate_mock = mocker.patch.object(api_server_module, "terminate_if_errored")
app = FastAPI()
app.state.args = SimpleNamespace(log_error_stack=False)
app.state.openai_serving_speech = handler
app.state.engine_client = SimpleNamespace(
raw_request = _make_api_server_request(handler, path="/v1/audio/speech")
raw_request.app.state.args = SimpleNamespace(log_error_stack=False)
raw_request.app.state.engine_client = SimpleNamespace(
engine=SimpleNamespace(is_alive=lambda: False),
errored=True,
)
app.state.server = SimpleNamespace()
scope = {
"type": "http",
"app": app,
"method": "POST",
"path": "/v1/audio/speech",
"headers": [],
"query_string": b"",
"client": ("127.0.0.1", 12345),
"server": ("testserver", 80),
"scheme": "http",
}
raw_request = Request(scope)
raw_request.app.state.server = SimpleNamespace()
raw_request.state.request_metadata = SimpleNamespace(request_id="speech-req-1")
request = OpenAICreateSpeechRequest(input="Hello")
@@ -0,0 +1,446 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for `/v1/realtime/video` WebSocket video output streaming."""
from __future__ import annotations
import asyncio
import json
from collections.abc import AsyncGenerator, Callable
from types import SimpleNamespace
from unittest.mock import MagicMock
import numpy as np
import pytest
from fastapi import FastAPI, WebSocket
from pytest_mock import MockerFixture
from starlette.testclient import TestClient
from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _fake_video_frames(num_frames: int = 2) -> list[np.ndarray]:
return [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(num_frames)]
def _build_test_app(
*,
mocker: MockerFixture,
streaming_chunks: list[tuple[bytes, bool]],
mock_generate: Callable[..., AsyncGenerator[OmniRequestOutput, None]],
final_chunk: bytes = b"",
stall_timeout: float = 5.0,
) -> tuple[FastAPI, OmniStreamingVideoOutputHandler, MagicMock]:
engine_client = mocker.MagicMock()
engine_client.abort = mocker.AsyncMock()
engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
engine_client.generate = mock_generate
class FakeStreamingVideoEncoder:
encode_calls: list[int] = []
def encode(self, video):
del video
idx = len(self.encode_calls)
self.encode_calls.append(idx)
return streaming_chunks[idx][0]
def close(self):
return final_chunk
def _make_encoder(*, output_format, fps, video_codec_options=None):
del output_format, fps, video_codec_options
return FakeStreamingVideoEncoder()
encoder_factory = mocker.MagicMock(side_effect=_make_encoder)
engine_client.streaming_encoder_factory = encoder_factory
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
encoder_factory,
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
return_value="diffusion",
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
return_value=[OmniDiffusionSamplingParams()],
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list",
return_value=[OmniDiffusionSamplingParams()],
)
handler = OmniStreamingVideoOutputHandler(
engine_client=engine_client,
model_name="test-model",
stage_configs=engine_client.stage_configs,
stall_timeout=stall_timeout,
start_timeout=5.0,
)
app = FastAPI()
@app.websocket("/v1/realtime/video")
async def ws_endpoint(websocket: WebSocket):
await handler.handle_session(websocket)
return app, handler, engine_client
class TestStreamingVideoOutputWebSocket:
"""WebSocket protocol tests for streaming generated video output."""
def test_streaming_session_emits_video_start_binary_chunks_and_done(self, mocker: MockerFixture):
"""A full session delivers video.start, binary chunks, then session.done."""
async def mock_generate(*_args, **_kwargs):
for frames, finished in [
(_fake_video_frames(2), False),
(_fake_video_frames(3), True),
]:
yield OmniRequestOutput.from_diffusion(
request_id="req-ws",
images=[frames],
final_output_type="image",
finished=finished,
)
app, _handler, engine_client = _build_test_app(
mocker=mocker,
streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)],
mock_generate=mock_generate,
)
with TestClient(app) as client:
with client.websocket_connect("/v1/realtime/video") as ws:
ws.send_json({"type": "session.start", "prompt": "A cat walking in the rain"})
start = ws.receive_json()
assert start["type"] == "video.start"
assert start["format"] == "m4s"
assert "format" not in start["config"]
assert "request_id" in start
assert ws.receive_bytes() == b"mp4-chunk-0"
assert ws.receive_bytes() == b"mp4-chunk-1"
done = ws.receive_json()
assert done["type"] == "session.done"
assert done["chunks"] == 2
assert done["stopped"] is False
engine_client.abort.assert_not_called()
engine_client.streaming_encoder_factory.assert_called_once()
assert engine_client.streaming_encoder_factory.call_args.kwargs["output_format"] == "m4s"
def test_streaming_session_rejects_unknown_format(self, mocker: MockerFixture):
async def mock_generate(*_args, **_kwargs):
if False:
yield None
app, _handler, engine_client = _build_test_app(
mocker=mocker,
streaming_chunks=[],
mock_generate=mock_generate,
)
with TestClient(app) as client:
with client.websocket_connect("/v1/realtime/video") as ws:
ws.send_json({"type": "session.start", "prompt": "hello", "format": "webm"})
err = ws.receive_json()
assert err["type"] == "error"
assert "format" in err["message"]
engine_client.abort.assert_not_called()
def test_streaming_session_emits_final_encoder_delta_before_done(self, mocker: MockerFixture):
"""A non-empty encoder close delta is delivered as the final binary chunk."""
async def mock_generate(*_args, **_kwargs):
yield OmniRequestOutput.from_diffusion(
request_id="req-ws",
images=[_fake_video_frames(2)],
final_output_type="image",
finished=True,
)
app, _handler, engine_client = _build_test_app(
mocker=mocker,
streaming_chunks=[(b"mp4-chunk-0", True)],
mock_generate=mock_generate,
final_chunk=b"mp4-trailer",
)
with TestClient(app) as client:
with client.websocket_connect("/v1/realtime/video") as ws:
ws.send_json({"type": "session.start", "prompt": "A cat walking in the rain"})
assert ws.receive_json()["type"] == "video.start"
assert ws.receive_bytes() == b"mp4-chunk-0"
assert ws.receive_bytes() == b"mp4-trailer"
done = ws.receive_json()
assert done["type"] == "session.done"
assert done["chunks"] == 2
assert done["stopped"] is False
engine_client.abort.assert_not_called()
def test_generation_error_output_returns_websocket_error(self, mocker: MockerFixture):
"""An erroneous OmniRequestOutput from generate is sent to the client as a WebSocket error."""
error_message = "diffusion engine exploded"
async def mock_generate(*_args, **_kwargs):
yield OmniRequestOutput.from_diffusion(
request_id="req-ws",
images=[_fake_video_frames(2)],
final_output_type="image",
finished=False,
)
yield OmniRequestOutput.from_error("req-ws", error_message)
app, _handler, engine_client = _build_test_app(
mocker=mocker,
streaming_chunks=[(b"mp4-chunk-0", False)],
mock_generate=mock_generate,
)
with TestClient(app) as client:
with client.websocket_connect("/v1/realtime/video") as ws:
ws.send_json({"type": "session.start", "prompt": "hello"})
assert ws.receive_json()["type"] == "video.start"
assert ws.receive_bytes() == b"mp4-chunk-0"
err = ws.receive_json()
assert err["type"] == "error"
assert error_message in err["message"]
engine_client.abort.assert_not_called()
class _MockVideoOutputWebSocket:
"""Feeds scripted client messages and records server sends."""
def __init__(self, messages: list[str]) -> None:
self._messages: asyncio.Queue[str] = asyncio.Queue()
for message in messages:
self._messages.put_nowait(message)
self.sent: list[dict | bytes] = []
def enqueue(self, message: str) -> None:
self._messages.put_nowait(message)
async def accept(self) -> None:
return None
async def receive_text(self) -> str:
return await self._messages.get()
async def send_json(self, data: dict) -> None:
self.sent.append(data)
async def send_bytes(self, data: bytes) -> None:
self.sent.append(data)
@property
def sent_json(self) -> list[dict]:
return [m for m in self.sent if isinstance(m, dict)]
def _build_handler_for_async_tests(
mocker: MockerFixture,
*,
mock_generate: Callable[..., AsyncGenerator[OmniRequestOutput, None]],
streaming_chunks: list[tuple[bytes, bool]] | None = None,
stall_timeout: float = 5.0,
) -> tuple[OmniStreamingVideoOutputHandler, MagicMock, _MockVideoOutputWebSocket]:
engine_client = mocker.MagicMock()
engine_client.abort = mocker.AsyncMock()
engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
engine_client.generate = mock_generate
chunks = streaming_chunks or [(b"mp4-chunk-0", True)]
class FakeStreamingVideoEncoder:
def __init__(self) -> None:
self.encode_calls = 0
def encode(self, video):
del video
idx = self.encode_calls
self.encode_calls += 1
return chunks[idx][0]
def close(self):
return b""
encoder_factory = mocker.MagicMock(side_effect=lambda **_: FakeStreamingVideoEncoder())
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
encoder_factory,
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
return_value="diffusion",
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
return_value=[OmniDiffusionSamplingParams()],
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list",
return_value=[OmniDiffusionSamplingParams()],
)
handler = OmniStreamingVideoOutputHandler(
engine_client=engine_client,
model_name="test-model",
stage_configs=engine_client.stage_configs,
stall_timeout=stall_timeout,
start_timeout=5.0,
)
start_msg = json.dumps({"type": "session.start", "prompt": "async test"})
ws = _MockVideoOutputWebSocket([start_msg])
return handler, engine_client, ws
class TestStreamingVideoOutputStallTimeout:
"""Stall-timeout and session.ping behavior."""
def test_long_generation_without_client_messages_succeeds(self, mocker: MockerFixture):
"""Engine progress between yields keeps the session alive without client pings."""
async def mock_generate(*_args, **_kwargs):
yield OmniRequestOutput.from_diffusion(
request_id="req-ws",
images=[_fake_video_frames(2)],
final_output_type="image",
finished=False,
)
await asyncio.sleep(0.3)
yield OmniRequestOutput.from_diffusion(
request_id="req-ws",
images=[_fake_video_frames(2)],
final_output_type="image",
finished=True,
)
app, _handler, engine_client = _build_test_app(
mocker=mocker,
streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)],
mock_generate=mock_generate,
stall_timeout=0.5,
)
with TestClient(app) as client:
with client.websocket_connect("/v1/realtime/video") as ws:
ws.send_json({"type": "session.start", "prompt": "slow generation"})
assert ws.receive_json()["type"] == "video.start"
assert ws.receive_bytes() == b"mp4-chunk-0"
assert ws.receive_bytes() == b"mp4-chunk-1"
done = ws.receive_json()
assert done["type"] == "session.done"
assert done["stopped"] is False
engine_client.abort.assert_not_called()
def test_stall_timeout_aborts_when_engine_silent(self, mocker: MockerFixture):
"""No engine output for longer than stall_timeout aborts the request."""
async def mock_generate(*_args, **_kwargs):
await asyncio.sleep(1.0)
yield OmniRequestOutput.from_diffusion(
request_id="req-ws",
images=[_fake_video_frames(2)],
final_output_type="image",
finished=True,
)
app, _handler, engine_client = _build_test_app(
mocker=mocker,
streaming_chunks=[(b"mp4-chunk-0", True)],
mock_generate=mock_generate,
stall_timeout=0.2,
)
with TestClient(app) as client:
with client.websocket_connect("/v1/realtime/video") as ws:
ws.send_json({"type": "session.start", "prompt": "stalled"})
assert ws.receive_json()["type"] == "video.start"
err = ws.receive_json()
assert err["type"] == "error"
assert "Stall timeout" in err["message"]
engine_client.abort.assert_called()
@pytest.mark.asyncio
async def test_session_ping_returns_pong(self, mocker: MockerFixture):
"""session.ping during generation receives session.pong."""
async def mock_generate(*_args, **_kwargs):
await asyncio.sleep(0.5)
yield OmniRequestOutput.from_diffusion(
request_id="req-ws",
images=[_fake_video_frames(2)],
final_output_type="image",
finished=True,
)
handler, _engine_client, ws = _build_handler_for_async_tests(
mocker,
mock_generate=mock_generate,
stall_timeout=5.0,
)
async def _send_ping_after_delay() -> None:
await asyncio.sleep(0.05)
ws.enqueue(json.dumps({"type": "session.ping"}))
ping_task = asyncio.create_task(_send_ping_after_delay())
await handler.handle_session(ws) # type: ignore[arg-type]
await ping_task
assert any(m.get("type") == "session.pong" for m in ws.sent_json)
@pytest.mark.asyncio
async def test_ping_prevents_stall_during_silent_engine(self, mocker: MockerFixture):
"""Client pings refresh the stall clock while the engine produces no output."""
release = asyncio.Event()
async def mock_generate(*_args, **_kwargs):
await release.wait()
yield OmniRequestOutput.from_diffusion(
request_id="req-ws",
images=[_fake_video_frames(2)],
final_output_type="image",
finished=True,
)
handler, engine_client, ws = _build_handler_for_async_tests(
mocker,
mock_generate=mock_generate,
stall_timeout=0.3,
)
async def _pinger() -> None:
while not release.is_set():
ws.enqueue(json.dumps({"type": "session.ping"}))
await asyncio.sleep(0.1)
pinger_task = asyncio.create_task(_pinger())
session_task = asyncio.create_task(handler.handle_session(ws)) # type: ignore[arg-type]
await asyncio.sleep(0.8)
release.set()
await session_task
await pinger_task
errors = [m for m in ws.sent_json if m.get("type") == "error"]
assert not any("Stall timeout" in m.get("message", "") for m in errors)
assert any(m.get("type") == "session.pong" for m in ws.sent_json)
assert any(m.get("type") == "session.done" for m in ws.sent_json)
engine_client.abort.assert_not_called()
@@ -48,6 +48,107 @@ def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch):
assert mux_calls[0]["audio"] is None
def test_fragmented_mp4_video_encoder_reuses_single_muxer(monkeypatch):
muxers = []
class FakeFragmentedMP4Muxer:
def __init__(self, *, width, height, fps, video_codec_options=None):
self.calls = []
muxers.append(
{
"width": width,
"height": height,
"fps": fps,
"video_codec_options": video_codec_options,
"instance": self,
}
)
def mux_video_frames(self, frames):
self.calls.append(frames)
return f"fragment-{len(self.calls)}".encode()
def close(self):
return b"trailer"
monkeypatch.setattr(
"vllm_omni.diffusion.utils.media_utils.FragmentedMP4Muxer",
FakeFragmentedMP4Muxer,
)
encoder = video_api_utils.FragmentedMP4VideoEncoder(
fps=12,
video_codec_options={"preset": "ultrafast"},
)
assert encoder.encode([np.zeros((4, 6, 3), dtype=np.float32)]) == b"fragment-1"
assert encoder.encode([np.ones((4, 6, 3), dtype=np.float32)]) == b"fragment-2"
assert encoder.close() == b"trailer"
assert len(muxers) == 1
assert muxers[0]["width"] == 6
assert muxers[0]["height"] == 4
assert muxers[0]["fps"] == 12.0
assert muxers[0]["video_codec_options"] == {"preset": "ultrafast"}
assert len(muxers[0]["instance"].calls) == 2
def test_create_streaming_video_encoder_selects_requested_format():
assert isinstance(
video_api_utils.create_streaming_video_encoder(output_format="m4s", fps=12),
video_api_utils.FragmentedMP4VideoEncoder,
)
def test_finalize_streaming_mp4_bytes_produces_progressive_mp4():
"""Fragment MP4 chunks are remuxed into decodable progressive MP4 bytes."""
import numpy as np
from vllm_omni.diffusion.utils.media_utils import FragmentedMP4Muxer, finalize_streaming_video_bytes
def _read_mp4_video_info(mp4_bytes: bytes) -> tuple[int, float, float | None, int, int]:
import io
import av
with av.open(io.BytesIO(mp4_bytes), mode="r", format="mp4") as container:
stream = container.streams.video[0]
frame_count = sum(1 for _ in container.decode(stream))
assert stream.average_rate is not None
fps = float(stream.average_rate)
duration_sec = None
if stream.duration is not None:
assert stream.time_base is not None
duration_sec = float(stream.duration * stream.time_base)
return frame_count, fps, duration_sec, stream.width, stream.height
width = 32
height = 32
fps = 16.0
input_frame_count = 2
muxer = FragmentedMP4Muxer(width=width, height=height, fps=fps)
frames = np.zeros((input_frame_count, height, width, 3), dtype=np.uint8)
streamed = muxer.mux_video_frames(frames) + muxer.close()
finalized = finalize_streaming_video_bytes(streamed, input_format="m4s", fps=fps)
assert finalized
assert finalized != streamed
streamed_info = _read_mp4_video_info(streamed)
final_info = _read_mp4_video_info(finalized)
expected_duration = input_frame_count / fps
assert streamed_info[0] == input_frame_count
assert final_info[0] == input_frame_count
assert streamed_info[1] == pytest.approx(fps)
assert final_info[1] == pytest.approx(fps)
assert streamed_info[3:] == (width, height)
assert final_info[3:] == (width, height)
assert final_info[2] == pytest.approx(expected_duration, rel=0.05)
assert streamed_info[2] == pytest.approx(expected_duration, rel=0.05)
def test_rife_model_inference_runs_on_dummy_tensors():
model = rife_interpolator.Model().eval()
img0 = torch.rand(1, 3, 32, 32)
@@ -1334,7 +1334,7 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture)
extra_params = {
"is_enable_stage2": True,
"pyramid_num_stages": 3,
"pyramid_num_inference_steps_list": [20, 20, 20],
"pyramid_num_inference_steps_list": [1, 1, 1],
"use_cfg_zero_star": True,
}
response = test_client.post(
@@ -1352,7 +1352,7 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture)
captured = engine.captured_sampling_params_list[0]
assert captured.extra_args["is_enable_stage2"] is True
assert captured.extra_args["pyramid_num_stages"] == 3
assert captured.extra_args["pyramid_num_inference_steps_list"] == [20, 20, 20]
assert captured.extra_args["pyramid_num_inference_steps_list"] == [1, 1, 1]
assert captured.extra_args["use_cfg_zero_star"] is True
+43
View File
@@ -8,6 +8,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams
from tests.helpers.mark import hardware_test
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model]
@@ -198,6 +199,48 @@ def test_generate_accepts_request_after_repeated_cancellations():
asyncio.run(run_test())
@pytest.mark.cpu
def test_generate_yields_streaming_diffusion_chunks_before_final():
"""AsyncOmni.generate yields every intermediate diffusion chunk before the final one."""
async def streaming_process_results(request_id, metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts):
del metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts
yield OmniRequestOutput.from_diffusion(
request_id=request_id,
images=[],
final_output_type="image",
custom_output={"chunk": 0},
finished=False,
)
yield OmniRequestOutput.from_diffusion(
request_id=request_id,
images=[],
final_output_type="image",
custom_output={"chunk": 1},
finished=True,
)
async def run_test():
omni = get_async_omni_instance()
omni._process_orchestrator_results = streaming_process_results
outputs = []
async for output in AsyncOmni.generate(
omni,
prompt={"prompt": "a cat"},
request_id="req-stream",
sampling_params_list=[SimpleNamespace()],
output_modalities=["image"],
):
outputs.append(output)
assert len(outputs) == 2
assert [output.finished for output in outputs] == [False, True]
assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
asyncio.run(run_test())
@pytest.mark.cpu
@pytest.mark.parametrize(
"output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY, RequestOutputKind.CUMULATIVE]
@@ -24,7 +24,7 @@ def test_dreamzero_vla_resolves_to_dreamzero_config(monkeypatch):
def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
monkeypatch.setattr(
"vllm_omni.config.stage_config.StageConfigFactory._auto_detect_model_type",
"vllm_omni.config.config_factory.StageConfigFactory._auto_detect_model_type",
classmethod(lambda _cls, _model, trust_remote_code=True: ("vla", None)),
)
monkeypatch.setattr(
@@ -1,11 +1,9 @@
"""
Offline inference tests: vace video generation.
See examples/offline_inference/vace/vace_video_generation.md
"""
"""VACE coverage through the shared text-to-video and image-to-video examples."""
from pathlib import Path
import pytest
from PIL import Image
from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet
from tests.helpers.assertions import assert_video_valid
@@ -14,22 +12,83 @@ from tests.helpers.mark import hardware_marks
pytestmark = [
pytest.mark.usefixtures("clean_gpu_memory_between_tests"),
pytest.mark.full_model,
pytest.mark.diffusion,
pytest.mark.example,
*hardware_marks(res={"cuda": "H100"}),
]
VACE_SCRIPT = EXAMPLES / "offline_inference" / "vace" / "vace_video_generation.py"
README_PATH = VACE_SCRIPT.with_name("vace_video_generation.md")
TEXT_TO_VIDEO_README = EXAMPLES / "offline_inference" / "text_to_video" / "text_to_video.md"
IMAGE_TO_VIDEO_README = EXAMPLES / "offline_inference" / "image_to_video" / "README.md"
EXAMPLE_OUTPUT_SUBFOLDER = "example_offline_vace"
README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH)
TEXT_TO_VIDEO_SNIPPET = next(
snippet
for snippet in ReadmeSnippet.extract_readme_snippets(TEXT_TO_VIDEO_README)
if "vace_t2v_output.mp4" in snippet.code
)
CONDITIONAL_OUTPUTS = (
"vace_i2v_output.mp4",
"vace_v2lf_output.mp4",
"vace_flf2v_output.mp4",
"vace_inpaint_output.mp4",
"vace_r2v_output.mp4",
)
IMAGE_TO_VIDEO_SNIPPETS = [
snippet
for snippet in ReadmeSnippet.extract_readme_snippets(IMAGE_TO_VIDEO_README)
if snippet.h2_title == "Wan2.1 VACE Conditional Tasks"
and any(output in snippet.code for output in CONDITIONAL_OUTPUTS)
]
@pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id)
def test_vace_video_generation(snippet: ReadmeSnippet, example_runner: ExampleRunner):
should_skip, reason = snippet.skip
if should_skip:
pytest.skip(reason)
def _smoke_snippet(snippet: ReadmeSnippet, asset_dir: Path) -> ReadmeSnippet:
# Six fresh model launches share a 90-minute X2V job. Keep the documented
# quality settings intact while reducing only the CI executions.
replacements = {
"astronaut.jpg": str(asset_dir / "astronaut.jpg"),
"vace_first_frame.png": str(asset_dir / "vace_first_frame.png"),
"vace_last_frame.png": str(asset_dir / "vace_last_frame.png"),
"vace_center_mask.png": str(asset_dir / "vace_center_mask.png"),
"--num-frames 81": "--num-frames 5",
"--num-inference-steps 30": "--num-inference-steps 2",
}
code = snippet.code
for source, target in replacements.items():
code = code.replace(source, target)
return snippet._replace(code=code)
@pytest.fixture
def vace_asset_dir(tmp_path: Path) -> Path:
asset_dir = tmp_path / "vace_assets"
asset_dir.mkdir(parents=True, exist_ok=True)
Image.new("RGB", (832, 480), (180, 180, 180)).save(asset_dir / "astronaut.jpg")
Image.new("RGB", (512, 512), (80, 140, 220)).save(asset_dir / "vace_first_frame.png")
Image.new("RGB", (512, 512), (220, 140, 80)).save(asset_dir / "vace_last_frame.png")
mask = Image.new("L", (832, 480), 0)
mask.paste(255, (336, 0, 496, 480))
mask.save(asset_dir / "vace_center_mask.png")
return asset_dir
def test_vace_text_to_video_shared_example(example_runner: ExampleRunner) -> None:
snippet = TEXT_TO_VIDEO_SNIPPET._replace(
code=TEXT_TO_VIDEO_SNIPPET.code.replace("--num-frames 81", "--num-frames 5").replace(
"--num-inference-steps 30", "--num-inference-steps 2"
)
)
result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER))
for asset in result.assets:
assert_video_valid(asset)
@pytest.mark.parametrize("snippet", IMAGE_TO_VIDEO_SNIPPETS, ids=lambda snippet: snippet.test_id)
def test_vace_conditional_shared_examples(
snippet: ReadmeSnippet,
example_runner: ExampleRunner,
vace_asset_dir: Path,
) -> None:
snippet = _smoke_snippet(snippet, vace_asset_dir)
result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER))
for asset in result.assets:
assert_video_valid(asset)
+18 -4
View File
@@ -17,8 +17,8 @@ class _FakeSummarizer:
self.chunks += 1
return f"summary-of-chunk-{chunk_index}"
async def compress_to_long_term(self, existing, mids):
return (existing + " | " if existing else "") + f"compressed({len(mids)})"
async def compress_to_long_term(self, mids):
return f"compressed({len(mids)})"
class _FakeDelegation:
@@ -42,7 +42,7 @@ def test_query_freshness_and_archive_on_change():
b.record_response("1 bottle")
assert b.update_query("now describe") is True
assert b.memory.qa_history[-1].query == "count bottles"
assert b.memory.qa_history[-1].responses == [("0.0s", "1 bottle")]
assert b.memory.qa_history[-1].responses == [("0.0 seconds", "1 bottle")]
def test_query_in_frame_on_arrival_then_prefix_on_later_chunks():
@@ -71,7 +71,7 @@ async def test_flush_archives_qa_and_summarizes():
assert summ.chunks == 1
assert b.chunk_index == 2
assert b.should_flush() is False
assert b.memory.qa_history[-1].responses == [("2.0s", "a person enters")]
assert b.memory.qa_history[-1].responses == [("2.0 seconds", "a person enters")]
assert len(b.memory.mid_term_summaries) == 1
@@ -85,6 +85,20 @@ async def test_long_term_compression_after_n_chunks():
assert b.memory.mid_term_summaries == []
@pytest.mark.asyncio
async def test_long_term_window_caps_blocks():
# Long-term memory must not grow unbounded over hours: oldest blocks are dropped
# once the sliding window is full (reference window = 15; here 2 for the test).
summ = _FakeSummarizer()
b = InteractionBrain(
summarizer=summ, chunk_frames=1, long_term_every_n_chunks=1, long_term_window=2, frame_seconds=1.0
)
for _ in range(4):
await b.flush([("0.0s", "u")])
assert len(b.memory.long_term_blocks) == 2 # capped, not 4
assert b.memory.long_term_memory.count("compressed(") == 2
@pytest.mark.asyncio
async def test_delegation_submit_and_fold():
b = InteractionBrain(delegation=_FakeDelegation())
+1 -1
View File
@@ -38,7 +38,7 @@ def test_commit_records_only_spoken():
assert p.commit("</silence>").action.value == "silence"
assert p.brain.response_records == []
assert p.commit("</response> hello").action.value == "response"
assert p.brain.response_records == [("0.0s", "hello")]
assert p.brain.response_records == [("0.0 seconds", "hello")]
def test_commit_dedup_default_is_exact_match_only():
+5
View File
@@ -115,6 +115,11 @@ def assert_video_diffusion_response(
expected_height = _maybe_int(form_data.get("height"))
expected_fps = _maybe_int(form_data.get("fps"))
# Skip num_frames assertion for Helios models because they round up frames
model = request_config.get("model", "")
if "Helios" in model:
expected_frames = None
for vid_bytes in response.videos:
assert_video_valid(
vid_bytes,
+14
View File
@@ -0,0 +1,14 @@
"""Fixtures config creation & out of tree registry management."""
import pytest
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
@pytest.fixture
def clean_pipeline_registry():
"""Ensure the OMNI_PIPELINES are in a clean state for a test that mutates it."""
snapshot = dict(OMNI_PIPELINES)
yield
OMNI_PIPELINES.clear()
OMNI_PIPELINES.update(snapshot)
+97 -1
View File
@@ -19,7 +19,7 @@ from collections.abc import Generator
from dataclasses import asdict, dataclass
from io import BytesIO
from pathlib import Path
from typing import Any, NamedTuple
from typing import Any, NamedTuple, cast
from urllib.parse import quote
import psutil
@@ -2168,6 +2168,102 @@ class OpenAIClientHandler:
timeout=timeout,
)
def send_streaming_video_diffusion_request(
self,
request_config: dict[str, Any],
request_num: int = 1,
*,
timeout_seconds: float = 600.0,
) -> list[DiffusionResponse]:
"""
Send a native ``/v1/realtime/video`` WebSocket request and return one
finalized MP4 artifact assembled from the streamed binary fragments.
"""
if request_num != 1:
raise NotImplementedError("Concurrent streaming video diffusion requests are not currently implemented")
response = asyncio.run(
self._send_streaming_video_diffusion_request_once(
request_config,
timeout_seconds=timeout_seconds,
)
)
assert_diffusion_response(response, request_config, run_level=self.run_level)
if response.e2e_latency is not None:
self._print_client_stat(f"[diffusion.stream] request#1 success in {response.e2e_latency:.3f}s")
else:
self._print_client_stat("[diffusion.stream] request#1 completed")
return [response]
async def _send_streaming_video_diffusion_request_once(
self,
request_config: dict[str, Any],
*,
timeout_seconds: float,
) -> DiffusionResponse:
form_data = request_config.get("form_data")
if not isinstance(form_data, dict):
raise ValueError("Video request_config must contain 'form_data'")
payload: dict[str, Any] = {
"type": "session.start",
**{key: value for key, value in form_data.items() if value is not None},
}
model = request_config.get("model")
if model is not None:
payload["model"] = model
payload.setdefault("format", "m4s")
fps = float(payload.get("fps") or 16)
stream_format = payload["format"]
url = self._build_ws_url("/v1/realtime/video")
result = DiffusionResponse()
chunks: list[bytes] = []
start_time = time.perf_counter()
deadline = start_time + timeout_seconds
import websockets
async with websockets.connect(url, max_size=None) as websocket:
await websocket.send(json.dumps(payload))
while True:
remaining = deadline - time.perf_counter()
if remaining <= 0:
raise TimeoutError(f"Streaming video request did not complete within {timeout_seconds}s")
message = await asyncio.wait_for(websocket.recv(), timeout=remaining)
if isinstance(message, bytes):
chunks.append(message)
continue
msg = json.loads(message)
msg_type = msg.get("type")
if msg_type == "video.start":
stream_format = msg.get("format") or stream_format
continue
if msg_type == "session.done":
break
if msg_type == "error":
raise RuntimeError(str(msg.get("message", msg)))
from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes
from vllm_omni.entrypoints.openai.video_api_utils import StreamingVideoFormat
streamed_bytes = b"".join(chunks)
if not streamed_bytes:
raise RuntimeError("Streaming video request completed without binary video chunks")
result.videos = [
finalize_streaming_video_bytes(
streamed_bytes,
input_format=cast(StreamingVideoFormat, stream_format),
fps=fps,
)
]
result.e2e_latency = time.perf_counter() - start_time
result.success = True
return result
def _wait_until_video_completed(
self, video_id: str, poll_interval_seconds: int = 2, timeout_seconds: int = 300
) -> None:
@@ -1,76 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for GLM-Image model __init__.py lazy import pattern.
The ``__init__.py`` uses ``__getattr__`` for lazy loading to avoid importing
``transformers.models.glm_image`` at module init, which may not be available
in all environments.
"""
import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestLazyImport:
"""Test the __getattr__ lazy import pattern in glm_image __init__.py."""
def test_getattr_exists_and_is_callable(self):
"""Verify __getattr__ exists and is callable."""
import vllm_omni.model_executor.models.glm_image as glm_image_pkg
assert hasattr(glm_image_pkg, "__getattr__")
assert callable(glm_image_pkg.__getattr__)
def test_getattr_returns_class_for_known_attribute(self):
"""Verify __getattr__ returns GlmImageForConditionalGeneration for known attribute."""
import vllm_omni.model_executor.models.glm_image as glm_image_pkg
# Call __getattr__ directly to test the lazy import logic
result = glm_image_pkg.__getattr__("GlmImageForConditionalGeneration")
# Verify we get a class (the actual GlmImageForConditionalGeneration)
assert result is not None
assert isinstance(result, type)
assert result.__name__ == "GlmImageForConditionalGeneration"
def test_getattr_raises_for_unknown_attribute(self):
"""Verify __getattr__ raises AttributeError for unknown attributes."""
import vllm_omni.model_executor.models.glm_image as glm_image_pkg
# Test unknown attribute via __getattr__ directly
with pytest.raises(AttributeError, match="has no attribute"):
glm_image_pkg.__getattr__("UnknownClass")
def test___all___exports_correct_symbols(self):
"""Verify __all__ contains the expected exported symbols."""
import vllm_omni.model_executor.models.glm_image as glm_image_pkg
assert hasattr(glm_image_pkg, "__all__")
assert "GlmImageForConditionalGeneration" in glm_image_pkg.__all__
class TestLazyImportDoesNotImportTransformersAtInit:
"""Verify that importing the package does not eagerly load transformers."""
def test_glm_image_module_has_getattr(self):
"""Test that the module has __getattr__ for lazy loading."""
import vllm_omni.model_executor.models.glm_image as glm_image_pkg
# The module should have __getattr__
assert hasattr(glm_image_pkg, "__getattr__")
assert callable(glm_image_pkg.__getattr__)
# And __all__
assert hasattr(glm_image_pkg, "__all__")
def test_module_has_proper_structure(self):
"""Test that the module has proper Python module structure."""
import vllm_omni.model_executor.models.glm_image as glm_image_pkg
# Should be a module
assert hasattr(glm_image_pkg, "__name__")
assert glm_image_pkg.__name__ == "vllm_omni.model_executor.models.glm_image"
# Should have __file__ (even if it's a package)
assert hasattr(glm_image_pkg, "__file__")
@@ -19,9 +19,8 @@ from types import SimpleNamespace
import pytest
from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
from vllm_omni.config.stage_config import (
_PIPELINE_REGISTRY,
PipelineConfig,
StageExecutionType,
)
@@ -34,13 +33,13 @@ _PIPELINE_KEY = "minicpmo_4_5"
class TestRegistryDeclaration:
def test_declared_in_omni_pipelines(self) -> None:
assert _PIPELINE_KEY in _OMNI_PIPELINES
assert _PIPELINE_KEY in OMNI_PIPELINES
def test_visible_in_central_registry(self) -> None:
assert _PIPELINE_KEY in _PIPELINE_REGISTRY
assert _PIPELINE_KEY in OMNI_PIPELINES
def test_lazy_load_returns_pipeline_config(self) -> None:
pipeline = _PIPELINE_REGISTRY[_PIPELINE_KEY]
pipeline = OMNI_PIPELINES[_PIPELINE_KEY]
assert isinstance(pipeline, PipelineConfig)
assert pipeline.model_type == _PIPELINE_KEY
assert pipeline.model_arch == "MiniCPMO45OmniForConditionalGeneration"
@@ -49,7 +48,7 @@ class TestRegistryDeclaration:
class TestPipelineTopology:
@pytest.fixture(scope="class")
def pipeline(self) -> PipelineConfig:
return _PIPELINE_REGISTRY[_PIPELINE_KEY]
return OMNI_PIPELINES[_PIPELINE_KEY]
def test_two_stages(self, pipeline: PipelineConfig) -> None:
assert len(pipeline.stages) == 2
@@ -99,7 +98,7 @@ class TestArchAliases:
@pytest.fixture(scope="class")
def pipeline(self) -> PipelineConfig:
return _PIPELINE_REGISTRY[_PIPELINE_KEY]
return OMNI_PIPELINES[_PIPELINE_KEY]
def test_shared_minicpmo_alias_present(self, pipeline: PipelineConfig) -> None:
# MiniCPM-o 4.5 ships ``architectures=["MiniCPMO"]`` in its HF config.
@@ -123,7 +122,7 @@ class TestHfConfigPredicate:
@pytest.fixture(scope="class")
def predicate(self):
pipeline = _PIPELINE_REGISTRY[_PIPELINE_KEY]
pipeline = OMNI_PIPELINES[_PIPELINE_KEY]
assert pipeline.hf_config_predicate is not None, (
"MiniCPM-o 4.5 pipeline must declare hf_config_predicate to "
"avoid misrouting MiniCPM-o 2.6 checkpoints into the 4.5 path."
@@ -375,6 +375,22 @@ def test_malformed_codec_length_warning_is_rate_limited():
assert out2.multimodal_outputs["model_outputs"][0].numel() == 0
def test_forward_emits_zero_samples_for_empty_codec_payload():
"""Consumer side of #4463.
The empty-but-finished payload that ``talker2code2wav_full_payload`` now
returns on a degenerate take carries a zero-length ``codes.audio``. Code2Wav
must turn that into a 0-sample output (finishing the request) so the Stage-1
wait gate releases, rather than stalling the pipeline to the connector
timeout.
"""
model = _make_model()
out = model.forward(input_ids=torch.zeros(0, dtype=torch.long))
assert out.multimodal_outputs["model_outputs"][0].numel() == 0
def test_forward_batches_equal_length_requests_in_one_decoder_call():
model = _make_model()
@@ -95,7 +95,7 @@ def test_aura2tts_builds_qwen3_tts_prompt_information():
[tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt])
assert len(tts_input["prompt_token_ids"]) >= 32
assert len(tts_input["prompt_token_ids"]) > 0
assert tts_input["additional_information"]["text"] == ["Hello."]
assert PRECOMPUTED_TEXT_IDS_KEY not in tts_input["additional_information"]
assert tts_input["additional_information"]["task_type"] == ["Base"]
@@ -166,6 +166,7 @@ def test_aura2tts_supports_custom_voice_mode():
assert tts_input["additional_information"]["task_type"] == ["CustomVoice"]
assert tts_input["additional_information"]["speaker"] == ["Vivian"]
assert "ref_audio" not in tts_input["additional_information"]
assert len(tts_input["prompt_token_ids"]) == 14
def test_aura2tts_passes_token_ids_to_qwen3_tts_when_enabled():
@@ -0,0 +1,336 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import torch
from vllm_omni.model_executor.stage_input_processors.indextts2 import (
_strip_stop_token,
talker2s2mel,
talker2s2mel_full_payload,
talker2s2mel_token_only,
)
STOP_MEL_TOKEN = 8193
LATENT_DIM = 16
def _make_talker_output(
*,
mel_codes: torch.Tensor,
latent: torch.Tensor,
s_ref: torch.Tensor | None = None,
ref_mel: torch.Tensor | None = None,
style: torch.Tensor | None = None,
finished: bool = True,
):
meta = {}
if s_ref is not None:
meta["S_ref"] = s_ref
if ref_mel is not None:
meta["ref_mel"] = ref_mel
if style is not None:
meta["style"] = style
mm = {
"codes": {"mel": mel_codes},
"hidden_states": {"latent": latent},
"meta": meta,
}
return SimpleNamespace(
finished=finished,
outputs=[SimpleNamespace(multimodal_output=mm)],
)
# ── _strip_stop_token ──
def test_strip_stop_token_basic():
codes = torch.tensor([10, 20, 30, STOP_MEL_TOKEN, 99])
latent = torch.randn(5, LATENT_DIM)
c, lat, lens = _strip_stop_token(codes, latent)
assert c.shape == (1, 3)
assert lat.shape == (1, 3, LATENT_DIM)
assert lens.tolist() == [3]
assert (c[0] == torch.tensor([10, 20, 30])).all()
def test_strip_stop_token_no_stop():
codes = torch.tensor([10, 20, 30])
latent = torch.randn(3, LATENT_DIM)
c, lat, lens = _strip_stop_token(codes, latent)
assert c.shape == (1, 3)
assert lens.tolist() == [3]
def test_strip_stop_token_at_start():
codes = torch.tensor([STOP_MEL_TOKEN, 10, 20])
latent = torch.randn(3, LATENT_DIM)
c, lat, lens = _strip_stop_token(codes, latent)
assert c.shape == (1, 0)
assert lat.shape == (1, 0, LATENT_DIM)
assert lens.tolist() == [0]
def test_strip_stop_token_batch_with_padding():
codes = torch.tensor(
[
[10, 20, STOP_MEL_TOKEN, 0],
[30, 40, 50, STOP_MEL_TOKEN],
]
)
latent = torch.randn(2, 4, LATENT_DIM)
c, lat, lens = _strip_stop_token(codes, latent)
assert c.shape == (2, 3)
assert lat.shape == (2, 3, LATENT_DIM)
assert lens.tolist() == [2, 3]
assert c[0, 2].item() == 0 # padded with a valid semantic code
def test_strip_stop_token_all_stop():
codes = torch.tensor([STOP_MEL_TOKEN])
latent = torch.randn(1, LATENT_DIM)
c, lat, lens = _strip_stop_token(codes, latent)
assert lens.tolist() == [0]
assert c.shape[1] == 0
assert lat.shape[1] == 0
# ── talker2s2mel ──
def test_talker2s2mel_basic():
mel_codes = torch.tensor([10, 20, 30, STOP_MEL_TOKEN])
latent = torch.randn(4, LATENT_DIM)
s_ref = torch.randn(1, 128)
ref_mel = torch.randn(1, 80, 200)
style = torch.randn(1, 256)
outputs = talker2s2mel(
source_outputs=[
_make_talker_output(
mel_codes=mel_codes,
latent=latent,
s_ref=s_ref,
ref_mel=ref_mel,
style=style,
)
]
)
assert len(outputs) == 1
info = outputs[0]["additional_information"]
assert info["mel_codes"].shape == (1, 3)
assert info["latent"].shape == (1, 3, LATENT_DIM)
assert info["code_lens"].tolist() == [3]
assert "S_ref" in info
assert "ref_mel" in info
assert "style" in info
assert outputs[0]["prompt_token_ids"] == [0]
def test_talker2s2mel_skips_unfinished():
mel_codes = torch.tensor([10, 20])
latent = torch.randn(2, LATENT_DIM)
outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, finished=False)])
assert len(outputs) == 0
def test_talker2s2mel_skips_empty_mel():
outputs = talker2s2mel(
source_outputs=[
_make_talker_output(
mel_codes=torch.tensor([]),
latent=torch.randn(0, LATENT_DIM),
)
]
)
assert len(outputs) == 0
def test_talker2s2mel_missing_optional_meta():
mel_codes = torch.tensor([10, 20])
latent = torch.randn(2, LATENT_DIM)
outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent)])
assert len(outputs) == 1
info = outputs[0]["additional_information"]
assert "S_ref" not in info
assert "ref_mel" not in info
assert "style" not in info
def test_talker2s2mel_batch():
out0 = _make_talker_output(
mel_codes=torch.tensor([1, 2, STOP_MEL_TOKEN]),
latent=torch.randn(3, LATENT_DIM),
s_ref=torch.randn(1, 64),
)
out1 = _make_talker_output(
mel_codes=torch.tensor([5, 6, 7, 8]),
latent=torch.randn(4, LATENT_DIM),
ref_mel=torch.randn(1, 80, 100),
)
outputs = talker2s2mel(source_outputs=[out0, out1])
assert len(outputs) == 2
assert outputs[0]["additional_information"]["code_lens"].tolist() == [2]
assert outputs[1]["additional_information"]["code_lens"].tolist() == [4]
def test_talker2s2mel_tensors_on_cpu():
device = torch.device("cpu")
mel_codes = torch.tensor([10, 20], device=device)
latent = torch.randn(2, LATENT_DIM, device=device)
s_ref = torch.randn(1, 64, device=device)
outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, s_ref=s_ref)])
info = outputs[0]["additional_information"]
assert info["latent"].device.type == "cpu"
assert info["mel_codes"].device.type == "cpu"
assert info["S_ref"].device.type == "cpu"
def test_talker2s2mel_no_multimodal_dict():
output = SimpleNamespace(
finished=True,
outputs=[SimpleNamespace(multimodal_output="not_a_dict")],
)
results = talker2s2mel(source_outputs=[output])
assert len(results) == 0
# ── full-payload connector path ──
def _conditioning_tensors():
return {
"S_ref": torch.randn(1, 3, LATENT_DIM),
"ref_mel": torch.randn(1, 80, 5),
"style": torch.randn(1, 4),
}
def test_talker2s2mel_token_only_placeholder_finished_only():
finished = _make_talker_output(mel_codes=torch.tensor([1]), latent=torch.randn(1, LATENT_DIM), finished=True)
unfinished = _make_talker_output(mel_codes=torch.tensor([2]), latent=torch.randn(1, LATENT_DIM), finished=False)
outputs = talker2s2mel_token_only([unfinished, finished])
assert len(outputs) == 1
assert outputs[0]["prompt_token_ids"] == [0]
assert outputs[0]["additional_information"] is None
assert outputs[0]["multi_modal_data"] is None
def test_talker2s2mel_full_payload_flat_keys_builds_s2mel_contract():
cond = _conditioning_tensors()
latent = torch.randn(3, LATENT_DIM)
payload = {
"codes.mel": torch.tensor([[10], [20], [30]]),
"hidden_states.latent": latent,
"meta.S_ref": cond["S_ref"],
"meta.ref_mel": cond["ref_mel"],
"meta.style": cond["style"],
}
result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-flat"))
assert result is not None
assert result["mel_codes"].shape == (1, 3)
assert result["mel_codes"].tolist() == [[10, 20, 30]]
assert result["latent"].shape == (1, 3, LATENT_DIM)
assert result["code_lens"].tolist() == [3]
assert result["S_ref"].device.type == "cpu"
assert result["ref_mel"].device.type == "cpu"
assert result["style"].device.type == "cpu"
assert result["S_ref"].data_ptr() != cond["S_ref"].data_ptr()
def test_talker2s2mel_full_payload_nested_fallback_input():
cond = _conditioning_tensors()
payload = {
"codes": {"mel": torch.tensor([[4], [5]])},
"hidden_states": {"latent": torch.randn(2, LATENT_DIM)},
"meta": cond,
}
result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-nested"))
assert result is not None
assert result["mel_codes"].tolist() == [[4, 5]]
assert result["latent"].shape == (1, 2, LATENT_DIM)
def test_talker2s2mel_full_payload_normalizes_one_by_t_mel_row():
cond = _conditioning_tensors()
payload = {
"codes.mel": torch.tensor([[7, 8, 9]]),
"hidden_states.latent": torch.randn(3, LATENT_DIM),
"meta": cond,
}
result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-row"))
assert result is not None
assert result["mel_codes"].tolist() == [[7, 8, 9]]
def test_talker2s2mel_full_payload_trims_stop_token_and_matching_latent():
cond = _conditioning_tensors()
latent = torch.arange(4 * LATENT_DIM, dtype=torch.float32).reshape(4, LATENT_DIM)
payload = {
"codes.mel": torch.tensor([[1], [2], [STOP_MEL_TOKEN], [99]]),
"hidden_states.latent": latent,
"meta": cond,
}
result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-stop"))
assert result is not None
assert result["mel_codes"].tolist() == [[1, 2]]
assert result["latent"].shape == (1, 2, LATENT_DIM)
assert result["latent"][0, 0].tolist() == latent[0].tolist()
assert result["latent"][0, 1].tolist() == latent[1].tolist()
assert result["code_lens"].tolist() == [2]
def test_talker2s2mel_full_payload_crops_mel_latent_to_common_length():
cond = _conditioning_tensors()
payload = {
"codes.mel": torch.tensor([[1], [2], [3], [4]]),
"hidden_states.latent": torch.randn(2, LATENT_DIM),
"meta": cond,
}
result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-crop"))
assert result is not None
assert result["mel_codes"].tolist() == [[1, 2]]
assert result["latent"].shape == (1, 2, LATENT_DIM)
assert result["code_lens"].tolist() == [2]
def test_talker2s2mel_full_payload_missing_required_fields_returns_none():
assert talker2s2mel_full_payload(None, {"hidden_states.latent": torch.randn(1, LATENT_DIM)}, None) is None
assert talker2s2mel_full_payload(None, {"codes.mel": torch.tensor([[1]])}, None) is None
@@ -486,3 +486,32 @@ def test_full_payload_omits_left_context_size_without_ref():
assert payload["meta"]["finished"].item() is True
assert "left_context_size" not in payload["meta"]
assert len(payload["codes"]["audio"]) == _Q * 2
@pytest.mark.parametrize(
"pooling_output",
[
pytest.param(SimpleNamespace(codes="not-a-dict"), id="non_dict_output"),
pytest.param({}, id="missing_codes_audio"),
pytest.param({"codes.audio": torch.zeros((3, _Q), dtype=torch.long)}, id="all_codes_filtered"),
],
)
def test_full_payload_emits_empty_finished_payload_on_degenerate_take(pooling_output):
"""Regression for #4463.
A degenerate talker take used to return ``None`` from
``talker2code2wav_full_payload``. The connector treats ``None`` as "drop the
request", but Stage-1 was already scheduled to receive it, so its wait gate
polls to ``connector_get_max_wait`` (~300s) and the orchestrator aborts — one
stuck request stalls the whole two-stage pipeline. Each degenerate case
(non-dict pooling_output, missing ``codes.audio``, all codec frames dropped by
the filter) must instead return an empty-but-finished payload so the gate
releases and the request finishes immediately with zero-length audio.
"""
request = SimpleNamespace(request_id="r", output_token_ids=[0, 1, 2])
payload = talker2code2wav_full_payload(transfer_manager=None, pooling_output=pooling_output, request=request)
assert payload is not None
assert payload["meta"]["finished"].item() is True
assert payload["codes"]["audio"].numel() == 0
+133
View File
@@ -10,6 +10,7 @@ from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_extras import (
build_image_to_image_prompt,
build_image_to_video_prompt,
build_text_to_image_prompt,
get_extra_body_params,
get_extra_output_params,
@@ -57,6 +58,57 @@ def test_sensenova_extra_registry_declares_request_and_response_params() -> None
assert should_init_extra_args_for_non_diffusion_stages("SenseNovaU1Pipeline") is False
@pytest.mark.core_model
@pytest.mark.cpu
def test_cosmos3_extra_registry_declares_request_and_response_params() -> None:
assert get_extra_body_params("Cosmos3OmniDiffusersPipeline") == frozenset(
{
"flow_shift",
"max_sequence_length",
"use_resolution_template",
"use_duration_template",
"use_system_prompt",
"system_prompt",
"negative_prompt",
"guardrails",
"condition_frame_indexes_vision",
"condition_video_keep",
"generate_sound",
"sound_gen",
"sound_duration",
"audio_duration",
"action_mode",
"action",
"domain_name",
"domain_id",
"raw_action_dim",
"action_chunk_size",
"action_space",
"action_fps",
"image_height",
"image_width",
"history_length",
"conditioning_fps",
"resolution",
"image_size",
"use_state",
"observation",
"robot_obs",
"deterministic_seed",
"session_id",
}
)
assert get_extra_output_params("Cosmos3OmniDiffusersPipeline") == frozenset(
{
"action",
"raw_action_dim",
"domain_id",
"action_mode",
}
)
assert should_init_extra_args_for_non_diffusion_stages("Cosmos3OmniDiffusersPipeline") is False
@pytest.mark.core_model
@pytest.mark.cpu
def test_magi_human_extra_registry_declares_request_and_response_params() -> None:
@@ -74,6 +126,27 @@ def test_magi_human_extra_registry_declares_request_and_response_params() -> Non
assert should_init_extra_args_for_non_diffusion_stages("MagiHumanPipeline") is False
@pytest.mark.core_model
@pytest.mark.cpu
def test_cosmos3_text_to_image_prompt_builder_selects_image_modality() -> None:
assert build_text_to_image_prompt(
"Cosmos3OmniDiffusersPipeline",
prompt="a red sports car at golden hour",
negative_prompt="blurry, distorted",
height=1024,
width=1024,
) == {
"prompt": "a red sports car at golden hour",
"modalities": ["image"],
"negative_prompt": "blurry, distorted",
}
assert build_text_to_image_prompt(
"Cosmos3OmniDiffusersPipeline",
prompt="a red sports car",
negative_prompt=None,
) == {"prompt": "a red sports car", "modalities": ["image"]}
@pytest.mark.core_model
@pytest.mark.cpu
def test_helios_extra_registry_declares_request_and_response_params() -> None:
@@ -104,6 +177,14 @@ def test_helios_extra_registry_declares_request_and_response_params() -> None:
assert should_init_extra_args_for_non_diffusion_stages(cls) is False
@pytest.mark.core_model
@pytest.mark.cpu
def test_vace_extra_registry_has_no_pipeline_params() -> None:
assert get_extra_body_params("WanVACEPipeline") == frozenset()
assert get_extra_output_params("WanVACEPipeline") == frozenset()
assert should_init_extra_args_for_non_diffusion_stages("WanVACEPipeline") is False
@pytest.mark.core_model
@pytest.mark.cpu
def test_unknown_pipeline_has_empty_extra_registry() -> None:
@@ -181,6 +262,58 @@ def test_unknown_pipeline_uses_default_image_to_image_prompt() -> None:
}
def _build_vace_prompt(media_inputs: dict[str, object], *, num_frames: int = 5) -> dict:
return build_image_to_video_prompt(
"WanVACEPipeline",
prompt="a bird flying",
negative_prompt=None,
media_inputs=media_inputs,
height=16,
width=320,
num_frames=num_frames,
)
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize(
"media_inputs",
[
{"image": Image.new("RGB", (320, 16), "red")},
{"last_image": Image.new("RGB", (320, 16), "blue")},
{
"image": Image.new("RGB", (320, 16), "red"),
"last_image": Image.new("RGB", (320, 16), "blue"),
},
{"image": Image.new("RGB", (320, 16), "red"), "mask": Image.new("L", (320, 16), 0)},
{"reference_images": [Image.new("RGB", (64, 64), "red")]},
],
ids=["i2v", "v2lf", "flf2v", "inpaint", "r2v"],
)
def test_vace_image_to_video_prompt_builder(media_inputs: dict[str, object]) -> None:
result = _build_vace_prompt(media_inputs)
mmd = result["multi_modal_data"]
if "reference_images" in mmd:
assert mmd["reference_images"] is media_inputs["reference_images"]
else:
assert len(mmd["video"]) == len(mmd["mask"]) == 5
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize(
("media_inputs", "message"),
[
({}, "requires a conditioning media input"),
({"mask": Image.new("L", (320, 16))}, "mask input requires an image"),
({"control_image": Image.new("RGB", (320, 16))}, "Unsupported VACE media input"),
],
)
def test_vace_rejects_invalid_media_combinations(media_inputs: dict[str, object], message: str) -> None:
with pytest.raises(ValueError, match=message):
_build_vace_prompt(media_inputs)
@pytest.mark.core_model
@pytest.mark.cpu
def test_declared_extra_args_apply_to_existing_sampling_params() -> None:
+252 -73
View File
@@ -12,14 +12,15 @@ from pathlib import Path
from unittest.mock import patch
import pytest
from transformers import PretrainedConfig, Qwen3OmniMoeConfig
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.config.config_factory import StageConfigFactory
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline
from vllm_omni.config.stage_config import (
_PIPELINE_REGISTRY,
DeployConfig,
PipelineConfig,
StageConfig,
StageConfigFactory,
StageDeployConfig,
StageExecutionType,
StagePipelineConfig,
@@ -30,14 +31,16 @@ from vllm_omni.config.stage_config import (
build_stage_runtime_overrides,
load_deploy_config,
merge_pipeline_deploy,
register_pipeline,
pipeline_cfg_resolver,
strip_parent_engine_args,
)
from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True)
Q3_OMNI_THINKER_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=False)
class TestStageType:
@@ -511,26 +514,39 @@ class TestStageResolutionHelpers:
class TestPipelineDiscovery:
"""Tests for the central pipeline registry (``pipeline_registry._OMNI_PIPELINES``)."""
"""Tests for the central pipeline registry (``OMNI_PIPELINES``)."""
def test_registry_has_known_models(self):
"""Built-in pipelines are lazy-loaded from the central declaration
on first access; no eager import or discovery walk needed."""
# ``in`` triggers the lazy-map lookup without forcing a load.
assert "qwen2_5_omni" in _PIPELINE_REGISTRY
assert "qwen3_omni_moe" in _PIPELINE_REGISTRY
assert "qwen3_tts" in _PIPELINE_REGISTRY
"""Check that specific models are in OMNI_PIPELINES."""
assert "qwen2_5_omni" in OMNI_PIPELINES
assert "qwen3_omni_moe" in OMNI_PIPELINES
assert "qwen3_tts" in OMNI_PIPELINES
def test_registry_loads_pipeline_on_getitem(self):
"""Looking up a registered model_type returns the matching PipelineConfig."""
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
def test_registry_resolver_qwen3_omni_all_stages(self):
"""Test that providing the HF config for qwen3 omni with audio enabled uses all stages."""
pipeline = StageConfigFactory.resolve_pipeline_config(
"qwen3_omni_moe",
Q3_OMNI_ALL_STAGES_HF_CONFIG,
)
assert isinstance(pipeline, PipelineConfig)
assert pipeline.model_type == "qwen3_omni_moe"
assert len(pipeline.stages) == 3 # thinker + talker + code2wav
def test_registry_resolver_qwen3_omni_thinker_only(self):
"""Test that providing the HF config for qwen3 omni without audio is thinker only."""
pipeline = StageConfigFactory.resolve_pipeline_config(
"qwen3_omni_moe",
Q3_OMNI_THINKER_HF_CONFIG,
)
assert isinstance(pipeline, PipelineConfig)
assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
assert len(pipeline.stages) == 1 # thinker only
def test_registry_returns_none_for_unknown(self):
"""Unknown model_types aren't found; ``get()`` returns None."""
assert "definitely_not_a_real_model" not in _PIPELINE_REGISTRY
assert _PIPELINE_REGISTRY.get("definitely_not_a_real_model") is None
"""Unknown model_types aren't found and resolve to `None`."""
assert "definitely_not_a_real_model" not in OMNI_PIPELINES
assert OMNI_PIPELINES.get("definitely_not_a_real_model") is None
assert StageConfigFactory.resolve_pipeline_config("definitely_not_a_real_model") is None
def test_pipeline_config_supports_hf_architectures(self):
"""PipelineConfig accepts hf_architectures for HF-arch fallback
@@ -579,6 +595,67 @@ class TestPipelineConfigNew:
assert any("no stages" in e.lower() for e in p.validate())
class TestPipelineRegistration:
def test_pipeline_registration(self, clean_pipeline_registry):
"""Ensure that we can register and create a custom pipeline config."""
new_model_type = "new_model_type"
pipe_cfg = PipelineConfig(model_type=new_model_type)
# Register the new PipelineConfig
assert new_model_type not in OMNI_PIPELINES
register_pipeline(pipe_cfg)
assert new_model_type in OMNI_PIPELINES
assert OMNI_PIPELINES[new_model_type] is pipe_cfg
class FakeConfig(PretrainedConfig):
model_type = new_model_type
# Create the model
with (
patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()),
patch.object(StageConfigFactory, "_create_from_registry") as mock_create,
):
StageConfigFactory.create_from_model("fake/model")
mock_create.assert_called_once()
# Ensure that the passed pipeline_config is the right type
pipeline_cfg = mock_create.call_args[0][1]
assert isinstance(pipeline_cfg, PipelineConfig)
assert pipe_cfg.model_type == new_model_type
def test_resolver_registration(self, clean_pipeline_registry):
"""Ensure that we can register and create a custom resolver."""
new_model_type = "new_model_type"
resolved_type = "resolved_type"
class FakeConfig(PretrainedConfig):
model_type = new_model_type
@pipeline_cfg_resolver(config_type=FakeConfig)
def custom_resolver(
hf_config: FakeConfig,
) -> PipelineConfig:
return PipelineConfig(model_type=resolved_type)
# Register the new PipelineConfig
assert new_model_type not in OMNI_PIPELINES
register_pipeline(custom_resolver, model_type=new_model_type)
assert new_model_type in OMNI_PIPELINES
assert OMNI_PIPELINES[new_model_type] is custom_resolver
# Create the model
with (
patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()),
patch.object(StageConfigFactory, "_create_from_registry") as mock_create,
):
StageConfigFactory.create_from_model("fake/model")
mock_create.assert_called_once()
# Ensure that the passed pipeline_config is the right type
pipeline_cfg = mock_create.call_args[0][1]
assert isinstance(pipeline_cfg, PipelineConfig)
assert pipeline_cfg.model_type == resolved_type
class TestResolveScheduler:
def test_all_execution_types_handled(self):
for et in StageExecutionType:
@@ -603,18 +680,6 @@ class TestResolveScheduler:
assert _resolve_scheduler(StageExecutionType.DIFFUSION) is None
class TestPipelineRegistry:
def test_register_and_lookup(self):
p = PipelineConfig(
model_type="__test_only__",
model_arch="A",
stages=(StagePipelineConfig(stage_id=0, model_stage="a"),),
)
register_pipeline(p)
assert _PIPELINE_REGISTRY["__test_only__"] is p
del _PIPELINE_REGISTRY["__test_only__"]
class TestDeployConfigLoading:
def test_custom_voice_dir_is_pipeline_wide_engine_arg(self, tmp_path):
deploy_path = tmp_path / "qwen3_tts_custom_voice.yaml"
@@ -633,8 +698,9 @@ stages:
)
deploy = load_deploy_config(deploy_path)
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
pipeline = StageConfigFactory.resolve_pipeline_config("qwen3_tts")
stages = merge_pipeline_deploy(pipeline, deploy)
assert isinstance(pipeline, PipelineConfig)
assert deploy.custom_voice_dir == str(custom_voice_dir)
assert {s.yaml_engine_args.get("custom_voice_dir") for s in stages} == {str(custom_voice_dir)}
@@ -664,11 +730,16 @@ stages:
assert "compilation_config" not in deploy.stages[0].engine_extras
def test_merge_pipeline_deploy(self):
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
if not deploy_path.exists():
pytest.skip("Deploy config not found")
pipeline = StageConfigFactory.resolve_pipeline_config(
"qwen3_omni_moe",
Q3_OMNI_ALL_STAGES_HF_CONFIG,
)
assert isinstance(pipeline, PipelineConfig)
deploy = load_deploy_config(deploy_path)
stages = merge_pipeline_deploy(pipeline, deploy)
@@ -680,7 +751,12 @@ stages:
assert s0.yaml_extras["default_sampling_params"]["detokenize"] is True
def test_merge_pipeline_deploy_preserves_num_replicas(self, tmp_path):
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
pipeline = StageConfigFactory.resolve_pipeline_config(
"qwen3_omni_moe",
Q3_OMNI_ALL_STAGES_HF_CONFIG,
)
assert isinstance(pipeline, PipelineConfig)
base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
if not base.exists():
pytest.skip("Deploy config not found")
@@ -894,14 +970,24 @@ stages:
class TestQwen3OmniPipeline:
def test_registered(self):
p = _PIPELINE_REGISTRY.get("qwen3_omni_moe")
assert p is not None
p = StageConfigFactory.resolve_pipeline_config(
"qwen3_omni_moe",
Q3_OMNI_ALL_STAGES_HF_CONFIG,
)
assert isinstance(p, PipelineConfig)
assert p.model_arch == "Qwen3OmniMoeForConditionalGeneration"
assert len(p.stages) == 3
assert p.validate() == []
def test_thinker(self):
s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(0)
p = StageConfigFactory.resolve_pipeline_config(
"qwen3_omni_moe",
Q3_OMNI_ALL_STAGES_HF_CONFIG,
)
assert isinstance(p, PipelineConfig)
s = p.get_stage(0)
assert isinstance(s, StagePipelineConfig)
assert s.model_stage == "thinker"
assert s.execution_type == StageExecutionType.LLM_AR
assert s.owns_tokenizer is True
@@ -909,14 +995,28 @@ class TestQwen3OmniPipeline:
assert s.sampling_constraints["detokenize"] is True
def test_talker(self):
s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(1)
p = StageConfigFactory.resolve_pipeline_config(
"qwen3_omni_moe",
Q3_OMNI_ALL_STAGES_HF_CONFIG,
)
assert isinstance(p, PipelineConfig)
s = p.get_stage(1)
assert isinstance(s, StagePipelineConfig)
assert s.input_sources == (0,)
assert s.sampling_constraints["stop_token_ids"] == [2150]
assert s.custom_process_input_func is not None
assert s.custom_process_next_stage_input_func is not None
def test_code2wav(self):
s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(2)
p = StageConfigFactory.resolve_pipeline_config(
"qwen3_omni_moe",
Q3_OMNI_ALL_STAGES_HF_CONFIG,
)
assert isinstance(p, PipelineConfig)
s = p.get_stage(2)
assert isinstance(s, StagePipelineConfig)
assert s.execution_type == StageExecutionType.LLM_GENERATION
assert s.final_output_type == "audio"
assert s.custom_process_input_func is not None
@@ -924,14 +1024,18 @@ class TestQwen3OmniPipeline:
class TestQwen2_5OmniPipeline:
def test_registered(self):
p = _PIPELINE_REGISTRY.get("qwen2_5_omni")
assert p is not None
p = StageConfigFactory.resolve_pipeline_config("qwen2_5_omni")
assert isinstance(p, PipelineConfig)
assert p.model_arch == "Qwen2_5OmniForConditionalGeneration"
assert len(p.stages) == 3
assert p.validate() == []
def test_thinker(self):
s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(0)
p = StageConfigFactory.resolve_pipeline_config("qwen2_5_omni")
assert isinstance(p, PipelineConfig)
s = p.get_stage(0)
assert isinstance(s, StagePipelineConfig)
assert s.model_stage == "thinker"
assert s.execution_type == StageExecutionType.LLM_AR
assert s.owns_tokenizer is True
@@ -939,13 +1043,21 @@ class TestQwen2_5OmniPipeline:
assert s.requires_multimodal_data is True
def test_talker(self):
s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(1)
p = StageConfigFactory.resolve_pipeline_config("qwen2_5_omni")
assert isinstance(p, PipelineConfig)
s = p.get_stage(1)
assert isinstance(s, StagePipelineConfig)
assert s.input_sources == (0,)
assert s.sampling_constraints["stop_token_ids"] == [8294]
assert s.custom_process_input_func is not None
def test_code2wav(self):
s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(2)
p = StageConfigFactory.resolve_pipeline_config("qwen2_5_omni")
assert isinstance(p, PipelineConfig)
s = p.get_stage(2)
assert isinstance(s, StagePipelineConfig)
assert s.execution_type == StageExecutionType.LLM_GENERATION
assert s.final_output_type == "audio"
assert s.engine_output_type == "audio"
@@ -953,14 +1065,19 @@ class TestQwen2_5OmniPipeline:
class TestQwen3TTSPipeline:
def test_registered(self):
p = _PIPELINE_REGISTRY.get("qwen3_tts")
p = StageConfigFactory.resolve_pipeline_config("qwen3_tts")
assert isinstance(p, PipelineConfig)
assert p is not None
assert p.model_arch == "Qwen3TTSTalkerForConditionalGeneration"
assert len(p.stages) == 2
assert p.validate() == []
def test_talker_stage(self):
s = _PIPELINE_REGISTRY["qwen3_tts"].get_stage(0)
p = StageConfigFactory.resolve_pipeline_config("qwen3_tts")
assert isinstance(p, PipelineConfig)
s = p.get_stage(0)
assert isinstance(s, StagePipelineConfig)
assert s.model_stage == "qwen3_tts"
assert s.execution_type == StageExecutionType.LLM_AR
assert s.owns_tokenizer is True
@@ -970,7 +1087,11 @@ class TestQwen3TTSPipeline:
assert s.model_arch is None
def test_code2wav_stage_has_per_stage_model_arch(self):
s = _PIPELINE_REGISTRY["qwen3_tts"].get_stage(1)
p = StageConfigFactory.resolve_pipeline_config("qwen3_tts")
assert isinstance(p, PipelineConfig)
s = p.get_stage(1)
assert isinstance(s, StagePipelineConfig)
assert s.execution_type == StageExecutionType.LLM_GENERATION
assert s.final_output_type == "audio"
assert s.engine_output_type == "audio"
@@ -986,7 +1107,7 @@ class TestQwen3TTSPipeline:
pytest.skip("qwen3_tts deploy yaml not found")
deploy = load_deploy_config(deploy_path)
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
pipeline = OMNI_PIPELINES["qwen3_tts"]
stages = merge_pipeline_deploy(pipeline, deploy)
# Stage 0 inherits pipeline-level model_arch
@@ -1025,14 +1146,18 @@ class TestQwen3TTSPipeline:
class TestMingFlashOmniPipeline:
def test_registered(self):
p = _PIPELINE_REGISTRY.get("ming_flash_omni")
assert p is not None
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni")
assert isinstance(p, PipelineConfig)
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
assert len(p.stages) == 2
assert p.validate() == []
def test_thinker_stage(self):
s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(0)
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni")
assert isinstance(p, PipelineConfig)
s = p.get_stage(0)
assert isinstance(s, StagePipelineConfig)
assert s.model_stage == "thinker"
assert s.execution_type == StageExecutionType.LLM_AR
assert s.owns_tokenizer is True
@@ -1042,7 +1167,11 @@ class TestMingFlashOmniPipeline:
assert s.sampling_constraints["detokenize"] is True
def test_talker_stage(self):
s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(1)
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni")
assert isinstance(p, PipelineConfig)
s = p.get_stage(1)
assert isinstance(s, StagePipelineConfig)
assert s.model_stage == "ming_tts"
assert s.execution_type == StageExecutionType.LLM_GENERATION
assert s.input_sources == (0,)
@@ -1060,20 +1189,28 @@ class TestMingFlashOmniPipeline:
Lazy string references only fail at first inference otherwise — this
catches typos in the pipeline declaration at import / registration time.
"""
s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(1)
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni")
assert isinstance(p, PipelineConfig)
s = p.get_stage(1)
assert isinstance(s, StagePipelineConfig)
module_path, _, attr = s.custom_process_input_func.rpartition(".")
module = importlib.import_module(module_path)
assert callable(getattr(module, attr))
def test_tts_pipeline_registered(self):
p = _PIPELINE_REGISTRY.get("ming_flash_omni_tts")
assert p is not None
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_tts")
assert isinstance(p, PipelineConfig)
assert p.model_arch == "MingFlashOmniTalkerForConditionalGeneration"
assert len(p.stages) == 1
assert p.validate() == []
def test_tts_stage(self):
s = _PIPELINE_REGISTRY["ming_flash_omni_tts"].get_stage(0)
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_tts")
assert isinstance(p, PipelineConfig)
s = p.get_stage(0)
assert isinstance(s, StagePipelineConfig)
assert s.model_stage == "ming_tts"
assert s.execution_type == StageExecutionType.LLM_GENERATION
assert s.input_sources == ()
@@ -1096,7 +1233,8 @@ class TestMingFlashOmniPipeline:
# We won't test stage 0/1 colocation contract here,
# as there could exist more variant of custom device setup
pipeline = _PIPELINE_REGISTRY["ming_flash_omni"]
pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni")
assert isinstance(pipeline, PipelineConfig)
stages = merge_pipeline_deploy(pipeline, deploy)
assert len(stages) == 2
assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration"
@@ -1112,20 +1250,24 @@ class TestMingFlashOmniPipeline:
assert len(deploy.stages) == 1
assert deploy.pipeline == "ming_flash_omni_tts"
pipeline = _PIPELINE_REGISTRY["ming_flash_omni_tts"]
pipeline = OMNI_PIPELINES["ming_flash_omni_tts"]
stages = merge_pipeline_deploy(pipeline, deploy)
assert len(stages) == 1
assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniTalkerForConditionalGeneration"
def test_thinker_only_pipeline_registered(self):
p = _PIPELINE_REGISTRY.get("ming_flash_omni_thinker_only")
assert p is not None
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_thinker_only")
assert isinstance(p, PipelineConfig)
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
assert len(p.stages) == 1
assert p.validate() == []
def test_thinker_only_stage(self):
s = _PIPELINE_REGISTRY["ming_flash_omni_thinker_only"].get_stage(0)
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_thinker_only")
assert isinstance(p, PipelineConfig)
s = p.get_stage(0)
assert isinstance(s, StagePipelineConfig)
assert s.model_stage == "thinker"
assert s.execution_type == StageExecutionType.LLM_AR
assert s.input_sources == ()
@@ -1146,20 +1288,21 @@ class TestMingFlashOmniPipeline:
assert len(deploy.stages) == 1
assert deploy.pipeline == "ming_flash_omni_thinker_only"
pipeline = _PIPELINE_REGISTRY["ming_flash_omni_thinker_only"]
pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_thinker_only")
assert isinstance(pipeline, PipelineConfig)
stages = merge_pipeline_deploy(pipeline, deploy)
assert len(stages) == 1
assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration"
def test_image_pipeline_registered(self):
p = _PIPELINE_REGISTRY.get("ming_flash_omni_image")
p = OMNI_PIPELINES.get("ming_flash_omni_image")
assert p is not None
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
assert len(p.stages) == 2
assert p.validate() == []
def test_image_thinker_stage(self):
s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0)
s = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_image").get_stage(0)
assert s.model_stage == "thinker"
assert s.execution_type == StageExecutionType.LLM_AR
assert s.input_sources == ()
@@ -1173,7 +1316,7 @@ class TestMingFlashOmniPipeline:
assert s.prompt_expand_func is not None
def test_image_dit_stage(self):
s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1)
s = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_image").get_stage(1)
assert s.model_stage == "dit"
assert s.execution_type == StageExecutionType.DIFFUSION
assert s.input_sources == (0,)
@@ -1185,8 +1328,11 @@ class TestMingFlashOmniPipeline:
def test_image_processor_wiring_resolves(self):
"""The prompt_expand_func and custom_process_input_func strings must point to real callables."""
thinker = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0)
dit = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1)
pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_image")
assert isinstance(pipeline, PipelineConfig)
thinker = pipeline.get_stage(0)
dit = pipeline.get_stage(1)
for ref in (thinker.prompt_expand_func, dit.custom_process_input_func):
module_path, _, attr = ref.rpartition(".")
module = importlib.import_module(module_path)
@@ -1205,7 +1351,8 @@ class TestMingFlashOmniPipeline:
assert deploy.connectors is not None
assert "shared_memory_connector" in deploy.connectors
pipeline = _PIPELINE_REGISTRY["ming_flash_omni_image"]
pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_image")
assert isinstance(pipeline, PipelineConfig)
stages = merge_pipeline_deploy(pipeline, deploy)
assert len(stages) == 2
# Stage 0 thinker: AR worker that emits latents.
@@ -1356,12 +1503,16 @@ class TestCLIOverrideFlow:
"""Test --stage-overrides JSON merge into StageConfig."""
def test_stage_overrides_merge(self):
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
if not deploy_path.exists():
pytest.skip("Deploy config not found")
deploy = load_deploy_config(deploy_path)
pipeline = StageConfigFactory.resolve_pipeline_config(
"qwen3_omni_moe",
Q3_OMNI_ALL_STAGES_HF_CONFIG,
)
assert isinstance(pipeline, PipelineConfig)
stages = merge_pipeline_deploy(pipeline, deploy)
# Simulate --stage-overrides '{"0": {"gpu_memory_utilization": 0.5}}'
@@ -1370,12 +1521,16 @@ class TestCLIOverrideFlow:
assert stages[0].runtime_overrides["gpu_memory_utilization"] == 0.5
def test_global_override_applies_to_all(self):
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
if not deploy_path.exists():
pytest.skip("Deploy config not found")
deploy = load_deploy_config(deploy_path)
pipeline = StageConfigFactory.resolve_pipeline_config(
"qwen3_omni_moe",
Q3_OMNI_ALL_STAGES_HF_CONFIG,
)
assert isinstance(pipeline, PipelineConfig)
stages = merge_pipeline_deploy(pipeline, deploy)
overrides = {"enforce_eager": True}
@@ -1392,8 +1547,11 @@ class TestAuraOmniDeploy:
assert deploy.pipeline == "aura_omni"
def test_aura_omni_deploy_resolves_four_native_stages(self):
pipeline_cfg = StageConfigFactory.resolve_pipeline_config("aura_omni")
stages = StageConfigFactory._create_from_registry(
"qwen3_tts",
pipeline_cfg,
cli_overrides={},
deploy_config_path=str(Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml"),
)
@@ -1417,8 +1575,14 @@ class TestSentinelDefaultPrecedence:
to YAML / dataclass defaults (#3035)."""
def _stages(self, cli_overrides):
model_type = "qwen3_omni_moe"
pipeline_cfg = StageConfigFactory.resolve_pipeline_config(
model_type,
Q3_OMNI_ALL_STAGES_HF_CONFIG,
)
return StageConfigFactory._create_from_registry(
"qwen3_omni_moe",
pipeline_cfg,
cli_overrides=cli_overrides,
)
@@ -1483,7 +1647,8 @@ class TestSentinelDefaultPrecedence:
"""A single ``qwen3_tts`` pipeline picks per-chunk vs end-to-end
processors based on ``deploy.async_chunk``, without needing a
separate variant pipeline registration."""
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
pipeline = StageConfigFactory.resolve_pipeline_config("qwen3_tts")
assert isinstance(pipeline, PipelineConfig)
# async_chunk=True → stage 0's per-chunk processor wires up, stage 1
# has no sync input processor.
@@ -1552,9 +1717,8 @@ class TestSentinelDefaultPrecedence:
``thinker2talker_token_only`` placeholder (sync_process_input_func).
Merge under either async_chunk mode must not re-introduce a
stage-0 full-payload hook."""
from vllm_omni.config.stage_config import DeployConfig, merge_pipeline_deploy
pipeline = _PIPELINE_REGISTRY["ming_flash_omni"]
pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni")
assert isinstance(pipeline, PipelineConfig)
stage0, stage1 = pipeline.stages
assert stage0.custom_process_next_stage_input_func is None, (
@@ -1587,15 +1751,30 @@ class TestSamplingConstraintsPrecedence:
"""Test that pipeline sampling_constraints override deploy defaults."""
def test_constraints_win(self):
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
if not deploy_path.exists():
pytest.skip("Deploy config not found")
deploy = load_deploy_config(deploy_path)
pipeline = StageConfigFactory.resolve_pipeline_config(
"qwen3_omni_moe",
Q3_OMNI_ALL_STAGES_HF_CONFIG,
)
assert isinstance(pipeline, PipelineConfig)
stages = merge_pipeline_deploy(pipeline, deploy)
# Pipeline says detokenize=True for thinker, deploy can't override
assert stages[0].yaml_extras["default_sampling_params"]["detokenize"] is True
# Pipeline says stop_token_ids=[2150] for talker
assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [2150]
class TestPipelineConfigResolvers:
@pytest.mark.parametrize("resolver", [obj for obj in OMNI_PIPELINES.values() if callable(obj)])
def test_all_resolvers_reject_bad_types(self, resolver):
"""Ensure that all resolvers registered reject incorrect config types."""
class NotTheRightHfConfig(PretrainedConfig):
pass
assert resolver(NotTheRightHfConfig()) is None
+1 -1
View File
@@ -10,7 +10,7 @@ from collections.abc import Mapping
import pytest
import torch
from vllm_omni.config.stage_config import StageConfigFactory
from vllm_omni.config.config_factory import StageConfigFactory
from vllm_omni.diffusion.data import (
DiffusionParallelConfig,
OmniDiffusionConfig,
@@ -19,7 +19,7 @@ import torch
class TestHiggsAudioV3Config:
def test_default_config_loads(self):
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
HiggsAudioV3Config,
)
@@ -34,7 +34,7 @@ class TestHiggsAudioV3Config:
assert config.tie_modality_embeddings is True
def test_custom_audio_encoder_config(self):
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
HiggsAudioV3Config,
)
@@ -51,7 +51,7 @@ class TestHiggsAudioV3Config:
assert config.tie_modality_embeddings is False
def test_invalid_num_codebooks_rejected(self):
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
HiggsAudioV3Config,
)
@@ -59,7 +59,7 @@ class TestHiggsAudioV3Config:
HiggsAudioV3Config(audio_encoder_config={"num_codebooks": 0, "vocab_size": 1026})
def test_negative_num_codebooks_rejected(self):
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
HiggsAudioV3Config,
)
@@ -67,7 +67,7 @@ class TestHiggsAudioV3Config:
HiggsAudioV3Config(audio_encoder_config={"num_codebooks": -1, "vocab_size": 1026})
def test_special_token_ids_initially_none(self):
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
HiggsAudioV3Config,
)
@@ -77,7 +77,7 @@ class TestHiggsAudioV3Config:
assert config.audio_continuation_id is None
def test_hidden_size_from_text_config(self):
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
HiggsAudioV3Config,
)
@@ -797,12 +797,12 @@ class TestAudioFeedback:
class TestCodecStrictness:
def test_bundled_missing_quantizer_key_raises(self):
"""Bundled codec load must fail when a quantizer codebook key is missing."""
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
HiggsAudioV3Code2Wav,
)
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
HiggsAudioV3Config,
)
config = HiggsAudioV3Config()
c2w = HiggsAudioV3Code2Wav(config=config)
@@ -821,12 +821,12 @@ class TestCodecStrictness:
def test_bundled_missing_fc2_raises(self):
"""Bundled codec load must fail when fc2 keys are missing."""
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
HiggsAudioV3Code2Wav,
)
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
HiggsAudioV3Config,
)
config = HiggsAudioV3Config()
c2w = HiggsAudioV3Code2Wav(config=config)
@@ -844,12 +844,12 @@ class TestCodecStrictness:
def test_bundled_unknown_decoder_key_rejected(self, monkeypatch):
"""Unknown decoder-side key should be rejected as unconsumed."""
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_code2wav as c2w_mod
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
HiggsAudioV3Code2Wav,
)
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
HiggsAudioV3Config,
)
class FakeDAC(torch.nn.Module):
def __init__(self):
@@ -871,12 +871,12 @@ class TestCodecStrictness:
def test_bundled_encoder_side_keys_accepted(self, monkeypatch):
"""Known encoder-side keys should not trigger unconsumed-key rejection."""
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_code2wav as c2w_mod
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
HiggsAudioV3Code2Wav,
)
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
HiggsAudioV3Config,
)
class FakeDAC(torch.nn.Module):
def __init__(self):
@@ -1199,9 +1199,9 @@ class TestRegistry:
assert "HiggsAudioV3Code2WavForConditionalGeneration" in _OMNI_MODELS
def test_pipeline_registered(self):
from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
assert "higgs_multimodal_qwen3" in _OMNI_PIPELINES
assert "higgs_multimodal_qwen3" in OMNI_PIPELINES
def test_deploy_yaml_exists(self):
import os
+9 -4
View File
@@ -8,11 +8,11 @@ import pytest
import yaml
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.config.config_factory import StageConfigFactory
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
from vllm_omni.config.stage_config import (
_PIPELINE_REGISTRY,
DeployConfig,
PipelineConfig,
StageConfigFactory,
StageDeployConfig,
StagePipelineConfig,
)
@@ -44,9 +44,9 @@ _TEST_DEPLOY = DeployConfig(
@pytest.fixture()
def mock_stages(monkeypatch):
"""Register a fake pipeline and mock deploy YAML loading."""
monkeypatch.setitem(_PIPELINE_REGISTRY._loaded, _TEST_MODEL, _TEST_PIPELINE)
monkeypatch.setitem(OMNI_PIPELINES, _TEST_MODEL, _TEST_PIPELINE)
monkeypatch.setattr(
"vllm_omni.config.stage_config.load_deploy_config",
"vllm_omni.config.config_factory.load_deploy_config",
lambda _path: _TEST_DEPLOY,
)
return __file__
@@ -462,6 +462,7 @@ def test_explicit_cli_arg_reaches_runtime_overrides(mock_stages):
explicit_kwargs = ns.get_explicit_kwargs_dict()
stages = StageConfigFactory._create_from_registry(
_TEST_MODEL,
_TEST_PIPELINE,
explicit_kwargs,
deploy_config_path=mock_stages,
)
@@ -478,6 +479,7 @@ def test_omitted_default_not_in_runtime_overrides(mock_stages):
explicit_kwargs = ns.get_explicit_kwargs_dict()
stages = StageConfigFactory._create_from_registry(
_TEST_MODEL,
_TEST_PIPELINE,
explicit_kwargs,
deploy_config_path=mock_stages,
)
@@ -499,6 +501,7 @@ def test_config_file_args_reach_runtime_overrides(mock_stages):
explicit_kwargs = ns.get_explicit_kwargs_dict()
stages = StageConfigFactory._create_from_registry(
_TEST_MODEL,
_TEST_PIPELINE,
explicit_kwargs,
deploy_config_path=mock_stages,
)
@@ -516,6 +519,7 @@ def test_per_stage_override_routes_correctly(mock_stages):
explicit_kwargs = ns.get_explicit_kwargs_dict()
stages = StageConfigFactory._create_from_registry(
_TEST_MODEL,
_TEST_PIPELINE,
explicit_kwargs,
deploy_config_path=mock_stages,
)
@@ -537,6 +541,7 @@ def test_explicit_args_omitted_from_yaml(mock_stages):
explicit_kwargs = ns.get_explicit_kwargs_dict()
stages = StageConfigFactory._create_from_registry(
_TEST_MODEL,
_TEST_PIPELINE,
explicit_kwargs,
deploy_config_path=mock_stages,
)
+352 -1
View File
@@ -1,9 +1,15 @@
from types import SimpleNamespace
import numpy as np
import pytest
import torch
from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, GPUARModelRunner
from vllm_omni.outputs import OmniModelRunnerOutput
from vllm_omni.worker.gpu_ar_model_runner import (
ExecuteModelState,
GPUARModelRunner,
OmniAsyncGPUModelRunnerOutput,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -52,6 +58,300 @@ def test_sparse_mm_req_ids_requires_sparse_audio_marker():
assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"]
def test_omni_async_gpu_model_runner_output_builds_lazily_once():
async_output = object.__new__(OmniAsyncGPUModelRunnerOutput)
calls = []
sync_calls = []
def builder():
calls.append("build")
return OmniModelRunnerOutput(req_ids=["r1"], req_id_to_index={"r1": 0})
async_output._model_runner_output = None
async_output._model_runner_output_builder = builder
async_output._invalid_req_indices = []
async_output.sampled_token_ids_cpu = torch.tensor([[7]], dtype=torch.long)
async_output.async_copy_ready_event = SimpleNamespace(synchronize=lambda: sync_calls.append("sync"))
async_output._sampled_token_ids = torch.tensor([[7]], dtype=torch.long)
async_output._logprobs_tensors = None
async_output._logprobs_tensors_cpu = None
async_output._routed_experts = None
async_output._routed_experts_cpu = None
async_output.vocab_size = 10
output = async_output.get_output()
assert calls == ["build"]
assert sync_calls == ["sync"]
assert async_output._model_runner_output_builder is None
assert output.req_ids == ["r1"]
assert output.sampled_token_ids == [[7]]
def test_omni_async_gpu_model_runner_output_reraises_background_exception():
async_output = object.__new__(OmniAsyncGPUModelRunnerOutput)
joined = []
class FakeThread:
def join(self):
joined.append("join")
async_output._background_thread = FakeThread()
async_output._background_exception = RuntimeError("background failed")
with pytest.raises(RuntimeError, match="background failed"):
async_output.get_output()
assert joined == ["join"]
assert async_output._background_thread is None
def _make_async_output_runner(engine_output_type: str = "audio"):
runner = object.__new__(GPUARModelRunner)
model_config = SimpleNamespace(
engine_output_type=engine_output_type,
async_chunk=True,
enable_return_routed_experts=False,
)
runner.vllm_config = SimpleNamespace(model_config=model_config)
runner.model_config = model_config
runner.omni_prefix_cache = None
runner.requests = {"r1": object(), "r2": object()}
runner.supports_mm_inputs = False
runner.routed_experts_initialized = False
runner.model = SimpleNamespace(has_postprocess=False)
runner.model_intermediate_buffer = {}
runner.input_batch = SimpleNamespace(
req_ids=["mutated"],
req_id_to_index={"mutated": 0},
)
return runner
def test_build_omni_output_uses_snapshots_and_connector_after_accumulation(monkeypatch):
runner = _make_async_output_runner()
events = []
monkeypatch.setattr(
GPUARModelRunner,
"_resolve_pooler_payload_req_ids",
lambda self, req_ids: ("audio", req_ids),
)
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: True)
monkeypatch.setattr(
GPUARModelRunner,
"accumulate_full_payload_output",
lambda self, rid, payload, request: events.append(f"accumulate:{rid}"),
)
monkeypatch.setattr(
GPUARModelRunner,
"get_omni_connector_output",
lambda self: events.append("connector") or "connector-output",
)
output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
runner,
scheduler_output=SimpleNamespace(
total_num_scheduled_tokens=3,
num_scheduled_tokens={"r1": 1, "r2": 2},
),
hidden_states=torch.tensor([[1.0], [2.0], [3.0]]),
staged_hidden_states_cpu=None,
multimodal_outputs={"foo": torch.tensor([10.0, 20.0, 30.0])},
req_ids_output_copy=["r1", "r2"],
req_id_to_index_output_copy={"r1": 0, "r2": 1},
valid_sampled_token_ids=[[101], [102]],
logprobs_lists=None,
prompt_logprobs_dict={},
num_nans_in_logits=None,
kv_connector_output=None,
ec_connector_output=None,
cudagraph_stats=None,
kv_extracted_req_ids=["r2"],
seq_len=3,
num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(),
query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
)
assert output.req_ids == ["r1", "r2"]
assert torch.equal(output.multimodal_outputs[0]["hidden"], torch.tensor([[1.0]]))
assert torch.equal(output.multimodal_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]]))
assert output.kv_extracted_req_ids == ["r2"]
assert output.omni_connector_output == "connector-output"
assert events == ["accumulate:r1", "accumulate:r2", "connector"]
def test_build_omni_output_copies_hidden_for_partial_downstream_batch(monkeypatch):
runner = _make_async_output_runner(engine_output_type="latent")
monkeypatch.setattr(
GPUARModelRunner,
"_resolve_pooler_payload_req_ids",
lambda self, req_ids: ("latent", ["r2"]),
)
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
runner,
scheduler_output=SimpleNamespace(
total_num_scheduled_tokens=6,
num_scheduled_tokens={"r1": 1, "r2": 2, "r3": 3},
),
hidden_states=torch.tensor([[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]]),
staged_hidden_states_cpu=None,
multimodal_outputs={},
req_ids_output_copy=["r1", "r2", "r3"],
req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2},
valid_sampled_token_ids=[[], [], []],
logprobs_lists=None,
prompt_logprobs_dict={},
num_nans_in_logits=None,
kv_connector_output=None,
ec_connector_output=None,
cudagraph_stats=None,
kv_extracted_req_ids=None,
seq_len=6,
num_scheduled_tokens_np=np.array([1, 2, 3], dtype=np.int32),
query_start_loc_cpu=torch.tensor([0, 1, 3], dtype=torch.long),
)
assert output.multimodal_outputs is not None
assert output.multimodal_outputs[0] == {}
assert torch.equal(output.multimodal_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]]))
assert output.multimodal_outputs[2] == {}
def test_process_additional_information_uses_snapshot_request_order(monkeypatch):
runner = _make_async_output_runner()
seen = []
class PostprocessModel:
has_postprocess = True
def postprocess(self, hidden_states, **kwargs):
seen.append(hidden_states.clone())
return {}
runner.model = PostprocessModel()
runner.model_intermediate_buffer = {"r1": {}, "r2": {}}
monkeypatch.setattr(
GPUARModelRunner,
"_resolve_pooler_payload_req_ids",
lambda self, req_ids: ("audio", req_ids),
)
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
monkeypatch.setattr(GPUARModelRunner, "_update_intermediate_buffer", lambda *args, **kwargs: None)
GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
runner,
scheduler_output=SimpleNamespace(
total_num_scheduled_tokens=3,
num_scheduled_tokens={"r1": 1, "r2": 2},
),
hidden_states=torch.tensor([[1.0], [2.0], [3.0]]),
staged_hidden_states_cpu=None,
multimodal_outputs={},
req_ids_output_copy=["r1", "r2"],
req_id_to_index_output_copy={"r1": 0, "r2": 1},
valid_sampled_token_ids=[[], []],
logprobs_lists=None,
prompt_logprobs_dict={},
num_nans_in_logits=None,
kv_connector_output=None,
ec_connector_output=None,
cudagraph_stats=None,
kv_extracted_req_ids=None,
seq_len=3,
num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(),
query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
)
assert len(seen) == 2
assert torch.equal(seen[0], torch.tensor([[1.0]]))
assert torch.equal(seen[1], torch.tensor([[2.0], [3.0]]))
def test_async_omni_output_guard_requires_safe_conditions():
runner = _make_async_output_runner()
runner.use_async_scheduling = True
runner.speculative_config = None
runner.model.use_async_omni_output = True
assert GPUARModelRunner._should_use_async_omni_output(runner)
runner.omni_prefix_cache = object()
assert not GPUARModelRunner._should_use_async_omni_output(runner)
runner.omni_prefix_cache = None
runner.model.has_postprocess = True
assert not GPUARModelRunner._should_use_async_omni_output(runner)
runner.model.eager_omni_postprocess_before_async_output = True
assert GPUARModelRunner._should_use_async_omni_output(runner)
def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
runner = _make_async_output_runner(engine_output_type="latent")
runner.model.omni_pooler_payload_include_hidden = False
monkeypatch.setattr(
GPUARModelRunner,
"_resolve_pooler_payload_req_ids",
lambda self, req_ids: ("latent", req_ids),
)
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
runner,
scheduler_output=SimpleNamespace(
total_num_scheduled_tokens=2,
num_scheduled_tokens={"r1": 2},
),
hidden_states=torch.tensor([[1.0], [2.0]]),
staged_hidden_states_cpu=None,
multimodal_outputs={"codes": {"audio": torch.tensor([[7, 8], [9, 10]], dtype=torch.long)}},
req_ids_output_copy=["r1"],
req_id_to_index_output_copy={"r1": 0},
valid_sampled_token_ids=[[101]],
logprobs_lists=None,
prompt_logprobs_dict={},
num_nans_in_logits=None,
kv_connector_output=None,
ec_connector_output=None,
cudagraph_stats=None,
kv_extracted_req_ids=None,
seq_len=2,
num_scheduled_tokens_np=np.array([2], dtype=np.int32),
query_start_loc_cpu=torch.tensor([0], dtype=torch.long),
)
assert output.multimodal_outputs is not None
assert len(output.multimodal_outputs) == 1
assert "hidden" not in output.multimodal_outputs[0]
assert torch.equal(output.multimodal_outputs[0]["codes.audio"], torch.tensor([[7, 8], [9, 10]], dtype=torch.long))
def test_async_snapshot_payload_omits_hidden_when_model_opts_out():
runner = _make_async_output_runner()
runner.model.omni_pooler_payload_include_hidden = False
payload = GPUARModelRunner._build_omni_async_snapshot_payload(
runner,
hidden_states=torch.tensor([[1.0], [2.0]]),
staged_hidden_states_cpu=torch.tensor([[3.0]]),
multimodal_outputs={"codes": {"audio": torch.tensor([[1]], dtype=torch.long)}},
)
assert set(payload.keys()) == {"multimodal_outputs"}
assert payload["multimodal_outputs"]["codes"]["audio"].tolist() == [[1]]
@pytest.mark.parametrize("query_start_loc_attr", ["method", "tensor_attr"])
def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monkeypatch, query_start_loc_attr):
runner = object.__new__(GPUARModelRunner)
@@ -133,3 +433,54 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
output.multimodal_outputs[1]["hidden"],
torch.tensor([[2.0, 20.0], [3.0, 30.0]]),
)
def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch):
"""Tail-only prefix-cache models still need per-step mm passthrough (e.g. codes.audio)."""
runner = _make_async_output_runner(engine_output_type="latent")
runner.omni_prefix_cache = object()
monkeypatch.setattr(
GPUARModelRunner,
"_resolve_pooler_payload_req_ids",
lambda self, req_ids: ("latent", req_ids),
)
monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False)
monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: {"codes.audio"})
monkeypatch.setattr(GPUARModelRunner, "_stage_deferred_prefix_cache_mm_outputs", lambda *args, **kwargs: None)
monkeypatch.setattr(
GPUARModelRunner,
"_prepare_prefix_cache_pooler_payload_sources",
lambda *args, **kwargs: (None, None, None),
)
monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
codes = torch.tensor([[11.0, 12.0], [21.0, 22.0]], dtype=torch.float32)
output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
runner,
scheduler_output=SimpleNamespace(
total_num_scheduled_tokens=2,
num_scheduled_tokens={"r1": 1, "r2": 1},
),
hidden_states=torch.tensor([[1.0], [2.0]]),
staged_hidden_states_cpu=None,
multimodal_outputs={"codes.audio": codes},
req_ids_output_copy=["r1", "r2"],
req_id_to_index_output_copy={"r1": 0, "r2": 1},
valid_sampled_token_ids=[[], []],
logprobs_lists=None,
prompt_logprobs_dict={},
num_nans_in_logits=None,
kv_connector_output=None,
ec_connector_output=None,
cudagraph_stats=None,
kv_extracted_req_ids=None,
seq_len=2,
num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
)
assert torch.equal(output.multimodal_outputs[0]["codes.audio"], codes[0:1])
assert torch.equal(output.multimodal_outputs[1]["codes.audio"], codes[1:2])
Loaded 100 of 279 files, more files were not shown because too many files have changed in this diff. Show more