Compare commits

...
Author SHA1 Message Date
Sy03 dedb045225 [CI] Incorporate upstream Server VAD test quarantine
Merge upstream #7314 for the known Qwen3-Omni Server VAD input_backpressure failure tracked in #7279. Preserve all MOSS-TTS runtime changes and existing validation coverage.

Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-09 19:53:28 +08:00
Sy03 28dd49a541 [Maintenance] Refresh MOSS-TTS optimization on main
Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-09 12:46:57 +08:00
wangyu 09b720b0d2 [CI] Skip Qwen3-Omni Server VAD multi-turn realtime test (#7279) (#7314)
Signed-off-by: wangyu <410167048@qq.com>
2026-09-09 12:23:29 +08:00
xutianle 968f3e379c [Bugfix][CosyVoice3] Resolve hash snapshot pipeline (#6896)
Signed-off-by: xutianle <xutianle@fudan.edu.cn>
2026-09-08 20:33:38 -07:00
Allen WuandTRAE CLI f2e30f46b6 [Bugfix][TTS] Isolate shared Higgs v3 reference encode from request cancellation (#7076)
Signed-off-by: Allen Wu <allenwu2795@gmail.com>
Co-authored-by: TRAE CLI <traecli@bytedance.com>
2026-09-08 20:32:05 -07:00
Jim Ban 358b7d79c1 [Bugfix][Frontend] Enforce image pixel limits for video input references (#6963)
Signed-off-by: BANANASJIM <bananasjim1@gmail.com>
2026-09-08 20:31:54 -07:00
Ziming Wang 43b697541a [Bugfix][Frontend] Clear the P0 multimodal cache through the renderer (#7003)
Signed-off-by: ZenAlexa <zimingwang945@gmail.com>
2026-09-08 20:31:42 -07:00
2bae550c04 [Frontend] Opt-in WebSocket TTS split_granularity and session seed (#7046)
Signed-off-by: rk9595 <rakesh.kariya@somaiya.edu>
Signed-off-by: Rakesh Kariya <rakesh.kariya@somaiya.edu>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-09-08 19:15:59 -07:00
zhengjiaandCursor 81c1c460b3 [Bugfix][Frontend] Accept the msgpack-numpy package's numpy markers on the OpenPI endpoint (#6051)
Signed-off-by: zjli2013 <leezhengjiang@126.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-09-08 19:15:48 -07:00
NancyFyongandNancyFyong e74fe3c6e8 [Model] Add Audio8 TTS Preview 0.6B (DualAR, 44.1 kHz codec) (#6157)
Signed-off-by: NancyFyong <NancyFyong@users.noreply.github.com>
Co-authored-by: NancyFyong <NancyFyong@users.noreply.github.com>
2026-09-08 19:00:25 -07:00
boatman cdf027874c [Feat][OmniVoice]Support Varlen Attn, Request-Batch and Step-Execution (#6408)
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
2026-09-08 18:58:38 -07:00
eval d7678e2bb3 [Bugfix][MiniCPM-o] Fix the audio_embeds input path (#5730)
Signed-off-by: eval-dev <0xe5bca0@gmail.com>
Signed-off-by: eval <74645252+eval-dev@users.noreply.github.com>
2026-09-08 18:41:42 -07:00
Zhichao Zhang 5078c7ad66 [Bugfix][MiMo-Audio] Align code2wav decode with tokenizer device (#6539)
Signed-off-by: chaosansui <zzc15560846421@163.com>
Signed-off-by: Zhichao Zhang <60429419+smartDream-chao@users.noreply.github.com>
2026-09-08 18:38:11 -07:00
Hongsheng Liu f67cc52f74 [Doc] Add AI usage policy for contributions (#7305)
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-09 09:17:24 +08:00
Anjie HouandClaude Opus 5 b6f4df320d [Refactor][Diffusion] Resolve offload topology through one plan resolver (#7209)
Signed-off-by: specture724 <specture724@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-09-09 08:17:38 +08:00
Tianyao Wu b3dd45874a [Bugfix] Build the forced-aligner prompt without a chat template (word timestamps one bin late) (#7240)
Signed-off-by: Tianyao Wu <rayroy31@gmail.com>
2026-09-08 19:27:38 +00:00
Lei KeandKrystalRay 5e4baea64d [Bugfix][NPU] Limit MiniMax H3 modulation grid size (#6794)
Signed-off-by: KrystalRay <keeleiray@gmail.com>
Co-authored-by: KrystalRay <keeleiray@gmail.com>
2026-09-08 09:08:46 -07:00
fdaf4e31bd Add vLLM-Omni library info to Hugging Face Hub requests (#5381)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-08 10:08:38 -06:00
liangmenghuangandClaude Fable 5.1 7a80b79698 [Bugfix][Model] GR00T-N1.7: honor the per-request seed for flow-matching noise (#7253)
Signed-off-by: liangmengh <liangmengh@nvidia.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
2026-09-08 15:53:20 +00:00
andyluo7 de1b0708bd [CI][ROCm] Route LTX2 Ulysses parity to two-GPU lane (#7234)
Signed-off-by: andyluo7 <andy.luo@amd.com>
2026-09-08 22:59:37 +08:00
Hongsheng Liu 7422a85cc7 [Bugfix] Skip HWR store-size scans when no limit is configured (#7131)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-08 22:54:49 +08:00
YancyandAsthenia a31690041d [Bugfix][Examples] Use --profiler-config flag in offline TTS examples (#6763)
Signed-off-by: Asthenia <asthenia0412@gmail.com>
Co-authored-by: Asthenia <asthenia0412@gmail.com>
2026-09-08 13:56:27 +00:00
Honghan Zhuandamy-why-3459 d59d166485 [Bugfix] Preserve additional_config for LLM stages (#7272)
Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-09-08 20:44:49 +08:00
Shuolei Wang 5378b77be4 [Model][JoyAI-VL-Interaction] Add native multi-stage pipeline (#5352)
Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com>
2026-09-08 05:07:08 -07:00
liangmenghuangandClaude Fable 5.1 d7740336ff [Bugfix][Model] GR00T-N1.7: honor processor_config image pipeline (use_albumentations / crop_fraction / letter_box_transform) to match Isaac-GR00T (#7083)
Signed-off-by: liangmengh <liangmengh@nvidia.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
2026-09-08 18:59:01 +08:00
LHXuuu c9c25b33c3 [Frontend] Add server-side VAD for turn-based omni models (#6618)
Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
2026-09-08 03:17:00 -07:00
Sy03 92cecf0296 [Bugfix] Score complete long audio in Seed-TTS WER
Retain audio beyond the default Whisper 30-second window and enable long-form timestamp decoding. Extend the existing transcription regression with a 61-second case while preserving the short-audio path. Clarify the existing test loader and per-turn request copy for static checking.

Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-08 17:40:59 +08:00
Alicia 031612f250 [Bugfix] Fix video benchmark timeout and failure accounting (#7259)
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-09-08 09:40:08 +00:00
andyluo7andwangyu b58ff5cb8b [CI/Build] Wait for duplex reaper recovery in retry test (#7225)
Signed-off-by: andyluo7 <andy.luo@amd.com>
Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com>
2026-09-08 17:10:51 +08:00
tlysanhuo 7c2e30afd3 [Bugfix] Bound omni benchmark per-request timeout to 15 min by default (#7130)
Signed-off-by: tly <2200895168@qq.com>
2026-09-08 02:00:25 -07:00
Chenguang Zheng f3fabe4145 [Bugfix] Keep --disable-log-stats out of stage config resolution (#7237)
Signed-off-by: Chenguang ZHENG <645327136@qq.com>
2026-09-08 16:19:07 +08:00
Sy03 8231f5c978 [Bugfix] Use reference transcripts for MOSS-TTS Local continuation
Build the official assistant audio prefix when Local cloning requests include a reference transcript. Preserve the audio-only fallback, conditioning cache identity, seed and output budget. Cover the prompt contract and fallback in the existing adapter tests.

Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-08 16:00:43 +08:00
Honghan Zhu c086554d1f [Perf]Boogu-Image Skip dense attention masks (#6871)
Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
2026-09-08 00:53:20 -07:00
wuli666 c704aeccae [Bugfix] Fix Step-Audio2 L4 CI KV cache allocation (#6839)
Signed-off-by: wuli666 <421774554@qq.com>
2026-09-08 14:17:33 +08:00
Chenguang Zheng 087ac46318 [Refactor 1/N] Centralize Omni config resolution (#5140)
Signed-off-by: Chenguang ZHENG <645327136@qq.com>
2026-09-08 10:30:26 +08:00
Deep Shahandamy-why-3459 1413eaab7c [Bugfix][Qwen3-Omni] Read code-predictor RoPE theta from rope_parameters (#7228)
Signed-off-by: Deep Shah <deep@socratic.co>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-09-08 10:25:27 +08:00
ZacheryAU b1ec136bd7 [Frontend] Add omni benchmark support for image and video endpoints (#4728) 2026-09-08 09:28:21 +08:00
Yueqian Lin e1c59edec6 [Refactor] fish_speech: vendor DAC codec modules, drop external fish-speech dependency (#4871)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-09-08 09:22:36 +08:00
boatman 043f03d576 [Migrate]Move model-specific helpers out of serving_speech (#6697)
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
2026-09-08 08:59:45 +08:00
Jierui (Jerry) Xu c3498fb869 [Model] Add Boogu-Image Turbo pipeline resolution (#6961)
Signed-off-by: Jerry Xu <xjr2423@gmail.com>
2026-09-08 08:59:23 +08:00
6942a5b600 [1/N] Stream Wan VAE chunks to the media consumer (#7016)
Signed-off-by: specture724 <specture724@gmail.com>
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-08 07:41:13 +08:00
Rahul SteigerandNils Hoffmann 3b335814cc [Bugfix][Cosmos3] Fix distributed Transfer output envelope (#7205)
Signed-off-by: Rahul Steiger <rsteiger@nvidia.com>
Co-authored-by: Nils Hoffmann <nhoffmann@nvidia.com>
2026-09-08 07:34:08 +08:00
Sy03 d530938b5b [Model] Reuse frame-local KV in MOSS-TTS Local
Cache each position once while generating the audio codebooks, preserving frame-local RoPE positions and the full-prefix reference path. Share compiled batch and position shapes to avoid specialization-limit fallback. Add one compact parity regression covering frame reuse.

Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-08 06:37:00 +08:00
Sy03 30d28a9a04 [Test] Consolidate MOSS-TTS batching regression coverage
Combine overlapping row/snapshot checks, cover masked and valid attention rows in one batch, reuse the real codec session in grouping tests, and move output-policy coverage into the existing entrypoint module. Preserve both CUDA synchronization regressions and reduce the added suite from 29 to 15 cases.

Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-08 03:03:32 +08:00
Sy03 55855fee9c [Model] Reduce MOSS-TTS Local intermediate output overhead
Keep internal talker control outputs FINAL_ONLY while codec audio remains incremental. Reapply pipeline sampling constraints after default streaming coercion, log routine coercions at debug level, and cover explicit/default parameters with regression tests. Document matched high-concurrency statistics settings.

Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-08 01:15:46 +08:00
Sy03 dcd6588a36 [Perf][MOSS-TTS] Avoid synchronization in batched index transfers
Construct MTP row offsets and mixed prefill/decode stop indices on the CPU, then enqueue nonblocking device transfers before batched writes. Add CUDA regressions for synchronization, row placement, stop flags, and snapshot ownership.

Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-08 00:39:50 +08:00
cf36f5e474 [Bugfix][MiniCPM-o] Allow ragged audio_feature_lens across a batch (#7071)
Signed-off-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com>
Co-authored-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-09-07 23:51:13 +08:00
Lei KeandKrystalRay 7be014bce6 [Bugfix][NPU] Fix MiniMax-H3 INT8 quantization dispatch (#6876)
Signed-off-by: KrystalRay <keeleiray@gmail.com>
Co-authored-by: KrystalRay <keeleiray@gmail.com>
2026-09-07 14:50:44 +00:00
汪志鹏 9b6a308914 [Perf][Diffusion] Keep LoRA weights resident across activation cycles (#7195)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
2026-09-07 14:27:24 +00:00
WeiQing Chen 30d6a0b4e6 [Bugfix][MiniMax-H3] Stabilize keyframe VAE encoding (#7191)
Signed-off-by: david6666666 <530634352@qq.com>
Signed-off-by: David Chen <530634352@qq.com>
2026-09-07 21:52:54 +08:00
Zeyu Huang | 黃澤宇andAlicia 4b96c57730 [skip ci][Doc] Module design doc for diffusion runtime (#6440)
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-09-07 21:52:24 +08:00
Mu GuanLin 3204a0b2ad [Model][Perf] Optimize LTX Ulysses sequence parallelism (#7079)
Signed-off-by: mglyn <1203789601@qq.com>
2026-09-07 10:48:26 +00:00
Rahul SteigerandRahul Steiger f080f41e03 [Bugfix][Cosmos3] Auto-pad non-divisible Ulysses GEN sequences (#6918)
Signed-off-by: Rahul Steiger <rsteiger@nvl72d188-T17.cm.cluster>
Co-authored-by: Rahul Steiger <rsteiger@nvl72d188-T17.cm.cluster>
2026-09-07 08:58:28 +00:00
FredHuangNV 85b0bad2d3 [Frontend][Model] Add Cosmos3 control uploads to video API (#7027)
Signed-off-by: Fred Huang <frhuang@nvidia.com>
2026-09-07 01:38:56 -07:00
Sy03 251ef4f48f [Model] Optimize MOSS-TTS Local batched execution and streaming codec
Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-07 16:20:49 +08:00
677265b64b [Bugfix][Diffusion] Make Cosmos3 transfer image arrays writable (#6915)
Signed-off-by: Rahul Steiger <rsteiger@nvl72d179-T17.cm.cluster>
Co-authored-by: Rahul Steiger <rsteiger@nvl72d179-T17.cm.cluster>
Co-authored-by: Rahul Steiger <rsteiger@aws-cmh-slurm-1-vscode-04.cm.cluster>
2026-09-07 01:06:16 -07:00
Shenglei FuandClaude 91dfafcaab [Bugfix][Diffusion] Include guidance_scale_2_provided in the request-batch key (#7078)
Signed-off-by: Shenglei Fu <sfu@confluent.io>
Co-authored-by: Claude <noreply@anthropic.com>
2026-09-07 15:58:19 +08:00
summerandAlicia 4b81f692d1 [Bugfix][Diffusion] Canonicalize cache_backend=None to "none" (#7041)
Signed-off-by: summer <128961079+zhang-keliang@users.noreply.github.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-09-07 15:56:51 +08:00
WeiQing Chen b78c31eda6 [Doc] Update vLLM-Omni WeChat QR code (#7178)
Signed-off-by: David Chen <530634352@qq.com>
2026-09-07 06:28:30 +00:00
Lucas Ruan 34aa1d2afa [Doc] Add A100 recipe for Wan2.2-TI2V-5B (#7189)
Signed-off-by: lucasruan1618 <huy40025071@gmail.com>
2026-09-07 14:09:32 +08:00
BruceLoveDecimal 5cd925631c [Feature][LingBot World]stepwise execution (#6844)
Signed-off-by: BruceLoveDecimal <liuqihao970610@gmail.com>
Signed-off-by: liuqihao <liuqihao970610@gmail.com>
2026-09-07 13:38:23 +08:00
NumberWanandAlicia 7f32c326c6 [BugFix][Qwen-Image-Edit] restore txt_seq_lens RoPE width after #3588 regression (#5586)
Signed-off-by: NumberWan <wantszkin2003@gmail.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-09-07 12:10:47 +08:00
linzhenpl07andClaude Sonnet 5 38abc0ba28 [AR-Diffusion] Allocate K and V separately so the compiled path stops cloning the pool (#6463)
Signed-off-by: linzhenpl07 <linzhenpl07@gmail.com>
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-07 11:40:20 +08:00
Zeyu Huang | 黃澤宇 2df95de421 [Core][Diffusion] Refactor prompt updates into a modality-ready interaction framework (#6294)
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
2026-09-07 11:40:04 +08:00
汪志鹏 eeee257913 [Bugfix][Diffusion] Restore the TRTLLM attention default for the MiniMax-H3 modular alias (#7162)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
2026-09-07 07:21:21 +08:00
Xiaobaixia 615f97b9c7 [Feature] Define the runner preprocess phase contract and regression coverage (#7136)
Signed-off-by: xiaobaixia222 <289982752@qq.com>
2026-09-06 17:20:04 +00:00
NATUREandamy-why-3459 ff6e906a25 [Model] Preserve MiniCPM-o Talker full attention until capacity (#6799)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-09-06 22:25:35 +08:00
Hongsheng Liu 715b8b8746 [Model] MiniMax-H3: validate versioned text-conditioning handoff (#6720)
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-06 21:45:10 +08:00
SYLARandhsliuustc0106 14d2d00536 [Bugfix] Fix Hunyuan Image-3 expert mapping contract (#7021)
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
Co-authored-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-06 21:43:37 +08:00
汪志鹏andClaude Opus 5 af74a5a155 [Feature][MiniMax-H3] Support the full LightX2V Turbo LoRA matrix (#7062)
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-09-06 20:23:36 +08:00
439 changed files with 32006 additions and 7199 deletions

No files matched your search

+12 -1
View File
@@ -107,6 +107,11 @@ steps:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
- >-
timeout 20m pytest -s -v
tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py
-m 'core_model and rocm and cards_2'
--run-level "core_model"
- label: "Diffusion GPU Worker Test"
agent_pool: mi300_2
@@ -122,7 +127,13 @@ steps:
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
# Preserve legacy one-GPU tests without cards_1 while excluding every
# registered multi-GPU marker from this one-GPU job.
- >-
timeout 15m pytest -sv tests/diffusion/models/
-m 'core_model and cuda and not
(cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)'
--run-level "core_model"
- label: "Diffusion · Wan22 Test"
agent_pool: mi300_1
+3 -1
View File
@@ -79,6 +79,8 @@ steps:
steps:
- label: "TTS · Qwen3-TTS CustomVoice Test"
source_file_dependencies:
- tests/e2e/features/prefill_phase/
- vllm_omni/worker/gpu_model_runner.py
- tests/helpers/media.py
- tests/e2e/online_serving/test_qwen3_tts_customvoice.py
- tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
@@ -96,7 +98,7 @@ steps:
timeout 40m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m 'advanced_model and cuda' --run-level 'advanced_model'
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py tests/e2e/features/prefill_phase/test_chunked_prefill_tail.py -m 'advanced_model and cuda' --run-level 'advanced_model'
"
mirror_hardwares: l4_1
+2 -2
View File
@@ -76,7 +76,7 @@ steps:
- label: "Diffusion · Other Test"
commands:
- timeout 15m pytest -sv tests/diffusion/*.py -m 'core_model and cuda' --run-level "core_model"
- timeout 15m pytest -sv tests/diffusion/*.py tests/diffusion/ar_diffusion -m 'core_model and cuda' --run-level "core_model"
mirror_hardwares: l4_1
- group: ":card_index_dividers: Tiny Model Tests [Single GPU]"
@@ -93,7 +93,7 @@ steps:
commands:
- |
timeout 15m bash -c "
pytest -sv tests/engine/ tests/model_executor/ -m 'core_model and cuda' --run-level 'core_model'
pytest -sv tests/engine/ tests/model_executor/ tests/worker/test_batched_omni_output.py -m 'core_model and cuda' --run-level 'core_model'
"
mirror_hardwares: l4_1
+2
View File
@@ -19,6 +19,8 @@ steps:
pytest -s -v
tests/diffusion/layers/test_activation.py
tests/diffusion/layers/test_fused_qk_norm_rope_npu.py
tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py
tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py
-m 'core_model and npu and A3'
--run-level 'core_model'
+2
View File
@@ -2,6 +2,8 @@
You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/).
AI-assisted contributions must follow the [AI Usage Policy](docs/contributing/README.md#ai-usage-policy), including disclosure, human review, and validation requirements.
Local `pre-commit` gates (SPDX, forbidden imports, `torch.cuda`, shellcheck,
mypy, test marks, TTS adapter ratchet, markdownlint, Buildkite schema, and
how to extend allowlists) are documented in the
+30 -15
View File
@@ -17,18 +17,18 @@ The main entrypoint is:
vllm serve Qwen/Qwen-Image --omni --port 8099
```
2. Run a minimal benchmark:
1. Run a minimal benchmark:
```bash
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--base-url http://localhost:8099 \
--model Qwen/Qwen-Image \
--task t2i \
--dataset vbench \
--num-prompts 5
--base-url http://localhost:8099 \
--model Qwen/Qwen-Image \
--task t2i \
--dataset vbench \
--num-prompts 5
```
**Notes**
### Notes
- By default, image tasks talk to `http://<host>:<port>/v1/chat/completions`; video tasks talk to `/v1/videos`.
- If you run the server on another host or port, pass `--base-url` accordingly.
@@ -52,14 +52,14 @@ Example (`t2v`):
```bash
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--base-url http://localhost:8099 \
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--task t2v \
--dataset vbench \
--num-prompts 50 \
--width 640 --height 480 \
--num-frames 81 --fps 16 \
--num-inference-steps 40
--base-url http://localhost:8099 \
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--task t2v \
--dataset vbench \
--num-prompts 50 \
--width 640 --height 480 \
--num-frames 81 --fps 16 \
--num-inference-steps 40
```
Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths).
@@ -140,6 +140,21 @@ Traffic / concurrency flags:
- `--request-rate`: Target request rate (requests/second). If set to `inf`, the script sends all requests immediately.
- `--max-concurrency`: Max number of in-flight requests (default: `1`). This can hard-cap the achieved QPS: if it is too small, requests will queue behind the semaphore, and both achieved throughput and observed SLO attainment can be skewed.
### Video job timeout and failures
`--video-job-timeout` controls how long each `/v1/videos` job may be polled,
including time queued on the server (default: `900` seconds). For long video
jobs at high concurrency, use a larger budget, for example
`--video-job-timeout 1800`. This also applies to warmup jobs. The polling budget
starts after job creation; it excludes time waiting for the client concurrency
semaphore. Individual HTTP calls still use the aiohttp session timeout.
An expired job is counted as failed and deleted, which cancels queued or running
server generation. The benchmark continues processing the remaining requests.
The progress bar counts both successful and failed requests. The final report
shows the failure count and the first 10 errors; `--output-file` saves all errors
with their request IDs in `request_errors`.
### Batched warmup note
For batched serving runs, warm the same in-flight shape you plan to measure.
+64 -57
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import asyncio
import base64
import json
@@ -43,6 +46,7 @@ class RequestFuncInput:
video_paths: list[str] | None = None
request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
default_bot_task: str | None = DEFAULT_EDITS_BOT_TASK
video_job_timeout: float = 900.0
@dataclass
@@ -331,64 +335,64 @@ async def async_request_v1_videos(
output = RequestFuncOutput()
output.start_time = time.perf_counter()
files = dict(input.extra_body)
if input.prompt:
files.setdefault("prompt", input.prompt)
if input.width and input.height:
files.setdefault("height", input.height)
files.setdefault("width", input.width)
if input.num_frames:
files.setdefault("num_frames", input.num_frames)
if input.num_inference_steps:
files.setdefault("num_inference_steps", input.num_inference_steps)
if input.seed is not None:
files.setdefault("seed", input.seed)
if input.fps:
files.setdefault("fps", input.fps)
form = aiohttp.FormData()
for k, v in files.items():
form.add_field(k, str(v))
image_file = None
if input.image_paths and input.video_paths:
output.error = "Only one of image_paths or video_paths can be provided"
output.success = False
return output
job_id = None
try:
files = dict(input.extra_body)
if input.prompt:
files.setdefault("prompt", input.prompt)
if input.width and input.height:
files.setdefault("height", input.height)
files.setdefault("width", input.width)
if input.num_frames:
files.setdefault("num_frames", input.num_frames)
if input.num_inference_steps:
files.setdefault("num_inference_steps", input.num_inference_steps)
if input.seed is not None:
files.setdefault("seed", input.seed)
if input.fps:
files.setdefault("fps", input.fps)
if input.image_paths and len(input.image_paths) > 0:
image_path = input.image_paths[0]
image_file = open(image_path, "rb")
form.add_field(
"input_reference",
image_file,
filename=os.path.basename(image_path),
content_type=_guess_mime_type(image_path),
)
elif input.video_paths and len(input.video_paths) > 0:
video_path = input.video_paths[0]
if _is_minimax_h3_mixed_reference(input.model):
with open(video_path, "rb") as reference_file:
video_b64 = base64.b64encode(reference_file.read()).decode("ascii")
form.add_field(
"video_reference",
json.dumps([{"video_url": f"data:video/mp4;base64,{video_b64}"}]),
)
else:
image_file = open(video_path, "rb")
form = aiohttp.FormData()
for k, v in files.items():
form.add_field(k, str(v))
if input.image_paths and input.video_paths:
output.error = "Only one of image_paths or video_paths can be provided"
output.success = False
return output
if input.image_paths and len(input.image_paths) > 0:
image_path = input.image_paths[0]
image_file = open(image_path, "rb")
form.add_field(
"input_reference",
image_file,
filename=os.path.basename(video_path),
content_type=_guess_mime_type(video_path),
filename=os.path.basename(image_path),
content_type=_guess_mime_type(image_path),
)
elif input.video_paths and len(input.video_paths) > 0:
video_path = input.video_paths[0]
if _is_minimax_h3_mixed_reference(input.model):
with open(video_path, "rb") as reference_file:
video_b64 = base64.b64encode(reference_file.read()).decode("ascii")
form.add_field(
"video_reference",
json.dumps([{"video_url": f"data:video/mp4;base64,{video_b64}"}]),
)
else:
image_file = open(video_path, "rb")
form.add_field(
"input_reference",
image_file,
filename=os.path.basename(video_path),
content_type=_guess_mime_type(video_path),
)
job_id = None
job_status = None
poll_json = {}
resp_json = {}
job_status = None
poll_json = {}
resp_json = {}
try:
# invoke a post request (POST /v1/videos)
async with session.post(input.api_url, data=form) as response:
if response.status == 200:
@@ -406,7 +410,7 @@ async def async_request_v1_videos(
# invoke a poll request (GET /v1/videos/{video_id})
poll_interval = 2.0 # Unit(s)
timeout_seconds = 600.0
timeout_seconds = input.video_job_timeout
deadline = time.perf_counter() + timeout_seconds
job_url = f"{input.api_url}/{job_id}"
@@ -422,8 +426,11 @@ async def async_request_v1_videos(
poll_json = await poll_response.json()
job_status = poll_json.get("status")
if time.perf_counter() >= deadline:
output.error = f"Timed out waiting for video job {job_id} to complete."
if job_status not in {"completed", "failed"} and time.perf_counter() >= deadline:
output.error = (
f"Timed out after {timeout_seconds:g}s waiting for video job {job_id} to complete. "
"Increase --video-job-timeout for long-running or queued video jobs."
)
output.success = False
return output
@@ -452,7 +459,7 @@ async def async_request_v1_videos(
elif "peak_memory_mb" in resp_json:
output.peak_memory_mb = resp_json["peak_memory_mb"]
except Exception as e:
output.error = str(e)
output.error = f"{type(e).__name__}: {e}"
output.success = False
finally:
if image_file is not None:
@@ -465,13 +472,13 @@ async def async_request_v1_videos(
except Exception as e:
print(f"Failed to clean up video job {job_id}: {e}")
output.latency = time.perf_counter() - output.start_time
output.latency = time.perf_counter() - output.start_time
if pbar is not None:
pbar.update(1)
if output.success and input.slo_ms is not None:
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
if pbar:
pbar.update(1)
return output
@@ -1,6 +1,6 @@
# adapted from fastvideo
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Benchmark online serving for diffusion models (Image/Video Generation).
@@ -87,6 +87,7 @@ import base64
import glob
import json
import logging
import math
import os
import random
import tempfile
@@ -1116,6 +1117,11 @@ def calculate_metrics(
"duration": total_duration,
"completed_requests": num_success,
"failed_requests": len(error_outputs),
"request_errors": [
{"request_id": req.request_id, "error": out.error}
for req, out in zip(requests_list, outputs)
if not out.success
],
"throughput_qps": num_success / total_duration if total_duration > 0 else 0,
"latency_mean": np.mean(latencies) if latencies else 0,
"latency_median": np.median(latencies) if latencies else 0,
@@ -1261,6 +1267,9 @@ def _default_endpoint_for_task(task: str) -> str:
async def benchmark(args):
if not math.isfinite(args.video_job_timeout) or args.video_job_timeout <= 0:
raise ValueError("--video-job-timeout must be a positive finite number of seconds.")
# Construct base_url if not provided
if args.base_url is None:
args.base_url = f"http://{args.host}:{args.port}"
@@ -1312,6 +1321,8 @@ async def benchmark(args):
print("Loading requests...")
requests_list = dataset.get_requests()
print(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.")
for req in requests_list:
req.video_job_timeout = args.video_job_timeout
if args.return_stage_metrics and args.endpoint in _STAGE_METRICS_ENDPOINTS:
for req in requests_list:
@@ -1376,6 +1387,8 @@ async def benchmark(args):
metrics["model"] = args.model
metrics["dataset"] = args.dataset
metrics["task"] = args.task
if args.endpoint == "/v1/videos":
metrics["video_job_timeout"] = args.video_job_timeout
if args.endpoint == "/v1/images/edits":
metrics["bot_task"] = args.bot_task
@@ -1398,6 +1411,11 @@ async def benchmark(args):
)
)
print("{:<40} {}/{:<15}".format("Successful requests:", metrics["completed_requests"], len(requests_list)))
print("{:<40} {:<15}".format("Failed requests:", metrics["failed_requests"]))
for failure in metrics["request_errors"][:10]:
print(f"Request {failure['request_id']} failed: {failure['error']}")
if metrics["failed_requests"] > 10:
print("Showing first 10 failures; use --output-file to save all request errors.")
# Section 3: Performance Metrics
print(f"{'-' * 50}")
@@ -1481,6 +1499,13 @@ if __name__ == "__main__":
help="Path to local dataset file (optional).",
)
parser.add_argument("--num-prompts", type=int, default=10, help="Number of prompts to benchmark.")
parser.add_argument(
"--video-job-timeout",
type=float,
default=900.0,
help="Maximum polling time in seconds per video job, including server queue time (default: 900). "
"Increase this for long-running video generation or high concurrency.",
)
parser.add_argument(
"--max-concurrency",
type=int,
+19
View File
@@ -85,6 +85,25 @@ vllm bench serve --omni \
--save-result --result-dir ./results
```
#### Set the sample rate for models that are not 24 kHz
Streamed PCM carries no header, so the harness derives `audio_duration` from the
byte count and a **default 24000 Hz** assumption
(`vllm_omni/metrics/definitions.py:DEFAULT_AUDIO_SAMPLE_RATE`). For a model that
emits another rate, export the real one or every duration-derived metric is
wrong by `real_sr / 24000`:
```bash
# Audio8 TTS, Fish Speech S2 Pro, Ming-omni-tts, ... all decode at 44.1 kHz
export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100
```
Without it, a 44.1 kHz model reports `audio_duration` 1.84x too long, so
`audio_rtf` looks 1.84x *better* than reality and `audio_underrun` /
`Streaming continuity OK rate` are measured against inflated chunk budgets.
A quick sanity check: compare `Mean AUDIO_DURATION (s)` against the actual
length of a saved WAV for the same prompt.
#### Streaming continuity (`audio_underrun`)
`audio_underrun` (seconds) is the per-request worst-case buffer deficit
+11
View File
@@ -17,6 +17,17 @@ models:
extra_params:
lang: en
text_normalization: true
# Audio8 TTS Preview: DualAR, 44.1 kHz. trust_remote_code stays false.
# Bench at 44.1 kHz: export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100, otherwise
# the streamed-PCM duration defaults to 24 kHz and RTF looks 1.84x too good.
Audio8/Audio8-TTS-Preview-0.6b:
supported_tasks: [default_voice, voice_clone]
backend: openai-audio-speech
endpoint: /v1/audio/speech
trust_remote_code: false
task_extra_body:
default_voice: {}
voice_clone: {}
# -CustomVoice checkpoints lack speaker_encoder weights, so voice_clone is
# NOT supported (an attempt raises ValueError from _extract_speaker_embedding
+1
View File
@@ -127,6 +127,7 @@ nav:
- Async Chunk: design/feature/async_chunk.md
- Async Diffusion Output: design/feature/async_diffusion_output.md
- Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
- Prefill/Decode Phase Contract: design/feature/preprocess_phase_contract.md
- Automatic Prefix Caching: design/feature/prefix_caching.md
- Realtime AR-Diffusion: design/feature/realtime_ar_diffusion.md
- Communication:
Binary file not shown.

Before

Width:  |  Height:  |  Size: 274 KiB

After

Width:  |  Height:  |  Size: 272 KiB

+70
View File
@@ -436,6 +436,76 @@ This should be seen as an edge case, and if this behavior can be avoided by sett
</details>
### Image / video generation
For `/v1/images/generations`, `/v1/images/edits`, and `/v1/videos`, set `--endpoint` to that path and omit `--backend`. The endpoint path is registered as the bench request adapter, so a separate named backend is not required.
Example (`/v1/images/generations`):
```bash
vllm bench serve --omni \
--endpoint /v1/images/generations \
--dataset-name random \
--model ~/models/Qwen/Qwen-Image \
--tokenizer ~/models/Qwen/Qwen-Image/tokenizer \
--max-concurrency 1 \
--num-warmups 2 \
--num-prompts 10 \
--random-input-len 64 \
--random-output-len 1 \
--ignore-eos \
--percentile-metrics e2el \
--extra-body '{
"num_inference_steps": 20,
"seed": 42,
"true_cfg_scale": 4.0
}'
```
If successful, following output is like:
```text
============ Serving Benchmark Result ============
Successful requests: 3
Failed requests: 0
Maximum request concurrency: 1
Benchmark duration (s): 8.42
Request throughput (req/s): 0.36
Peak concurrent requests: 2.00
-------------------Peak Memory--------------------
Mean PEAK_MEMORY_MB (MB): 58832.00
Median PEAK_MEMORY_MB (MB): 58832.00
P99 PEAK_MEMORY_MB (MB): 58832.00
----------------End-to-end Latency----------------
Mean E2EL (ms): 2798.16
Median E2EL (ms): 2799.83
P99 E2EL (ms): 2800.03
================== Image Result ==================
Total images generated: 3
Image throughput (img/s): 0.36
Average pixels per image: 1048576.00
Mean denoise step latency (ms): 136.59
---------------- Image Generation ----------------
Mean IMAGE_GENERATION (ms): 2731.84
Median IMAGE_GENERATION (ms): 2732.02
P99 IMAGE_GENERATION (ms): 2734.65
==================================================
```
Use the same pattern with `--endpoint /v1/images/edits` or `--endpoint /v1/videos` (and model-specific `--extra-body` as needed). Pure image/video runs omit the Text Result section when there is no generated text.
`/v1/images/edits` defaults to **non-streaming JSON** (`stream=false`) so single-stage edit models are not rejected by the server. For multi-stage pipelines that need SSE (AR TTFT / image chunks), pass `"stream": true` in `--extra-body`.
`/v1/videos` is an async job API: the client creates a job, then polls until `completed`/`failed`. Measured **e2el** therefore includes client poll sleep and any overshoot after the job actually finishes. Tune polling via `--extra-body`:
- `poll_interval_s` (default `2.0`): sleep between status polls
- `poll_timeout_s` (default `21600`, i.e. 6 hours): give up waiting for the job
**VIDEO_RTF** prefers server-reported generation time when available:
- `video_rtf = video_generation_time_ms / 1000 / video_duration`, where `video_generation_time_ms` comes from response `stage_durations` (or `inference_time_s` as fallback)
- If generation time is missing, falls back to `e2el / video_duration` (this fallback *does* include poll overhead)
### Multi-Stage Benchmark
<details class="admonition abstract" markdown="1">
+22
View File
@@ -262,6 +262,28 @@ or serving, explain the required L2/L3 environment and provide the exact command
Repository skills guide the agent, but they do not replace contributor judgment, the accepted issue or RFC, required test evidence, or maintainer review. Review all generated changes before submitting them.
## AI Usage Policy
AI tools are welcome in vLLM-Omni contributions. Contributors, including maintainers, remain responsible for everything they submit. These requirements apply to pull requests, issues, discussions, and review comments:
- **Disclose AI assistance.** Name the tools used and describe their role in the submission, including assistance with code, tests, documentation, or writing.
- **Understand your changes.** Review every generated change and be able to explain its behavior, design choices, and interaction with vLLM-Omni without relying on AI to answer for you.
- **Validate before submitting.** Run relevant checks and report actual results. For model, kernel, or performance changes, include the hardware, model, configuration, and commands needed to reproduce accuracy or benchmark evidence. Clearly identify checks you could not run; generated claims are not test results.
- **Keep communication useful.** Verify technical claims, search for existing reports, and edit AI-assisted text for accuracy and brevity. Submit only content you have personally reviewed, and take responsibility for responding to review feedback.
- **Label generated media.** Model-generated images, video, and audio are valid examples or test artifacts for this multimodal project. Identify them as generated output and provide the model, inputs, and reproduction settings. Do not present them as ground-truth references or independent evidence of correctness.
For example, a PR description could include:
```text
AI assistance: Used Codex to draft the regression test and PR description.
I reviewed the changes and verified the test reproduces the reported failure.
Validation: <commands, results, hardware, and any checks not run>.
```
Maintainer review time is limited. Contributions that are unverified, misleading, or beyond the author's understanding may be closed. Repeated violations may lead to contribution restrictions. New contributors are welcome to ask for help and submit small, well-understood changes.
This policy is adapted from [Ghostty's AI Usage Policy](https://github.com/ghostty-org/ghostty/blob/main/AI_POLICY.md) for vLLM-Omni's contribution and multimodal validation workflows. The [coding-agent workflow](#using-repository-skills-with-coding-agents) provides practical guidance for applying it.
## Issues
If you encounter a bug or have a feature request, please search existing issues first to see if it has already been reported. If not, please file a new issue, providing as much relevant information as possible.
@@ -779,6 +779,31 @@ turning them on. For Qwen-Image-style serving examples, document
`--step-execution` as the feature gate and `--max-num-seqs N` as the
companion batching knob.
### Streaming Output and Mid-Stream Interactions
To add streaming generation with optional mid-stream prompt updates to a new diffusion pipeline (often a video generation pipeline):
0. **Model should internally support** chunked video emission and bounded VAE decode of a chunk.
1. **Implement stepwise execution** so interaction RPCs can be consumed at chunk
boundaries (see [Step Execution](#step-execution) above).
2. **Mix in** ``InteractionMixin`` on the pipeline class.
3. **Implement** ``peek_chunk_media(StepRequestState) -> ChunkMediaSpec`` to offer information
about the next chunk to generate.
This is called whenever one supported interaction modality needs frame-level calculation,
e.g., per-frame camera trajectory. If only chunk-level calculation is involved, peeking is skipped.
4. **Optionally implement** ``prepare_next_chunk(StepRequestState)`` if the
pipeline must rebuild something (e.g., per-chunk state) after an interaction is applied.
This function is called at chunk boundary.
5. **Register** supported modalities and handlers in ``vllm_omni/diffusion/interaction/registry.py``
(keyed by ``od_config.model_class_name`` → modality key as in ``req["multi_modal_data"]``).
Serve with ``--diffusion-streaming-output`` (auto-enables step execution). See
[Diffusion Execution Modes](../../user_guide/diffusion/execution_modes.md#streaming-output)
and the example script at `examples/online_serving/streaming_video_generation/`.
**Interaction payload contract:** Although ``OmniInteractionPrompt`` type annocation requires ``event_id``,
the WebSocket API payload may omit ``event_id`` and let the API layer assigns one.
### Cache Acceleration
#### TeaCache
@@ -222,6 +222,16 @@ __all__ = ["Qwen3OmniMoeForConditionalGeneration"]
## Key Components
### Preprocess phase metadata
Models with `has_preprocess=True`, including out-of-tree plugins, can rely on
the runner's stable `_omni_is_prefill`, `_omni_prompt_len`, and
`_omni_num_computed_tokens` fields. Use request progress to distinguish prompt
prefill from decode: a one-token span can still be a prefill tail. See the
[runner-to-model phase contract](../../design/feature/preprocess_phase_contract.md)
for types, batched-hook semantics, MTP routing guarantees, and the correctness
workaround for vLLM-Omni 0.20.x and earlier.
### 1. Model Interfaces
Your model should implement the appropriate interfaces:
+15 -10
View File
@@ -22,11 +22,15 @@ A pipeline opts in by implementing `SupportsARDiffusionPipeline`. Its immutable
- named fixed-length cross-attention KV, model-specific scratch geometry, and
the pipeline's requested session-capacity upper bound.
The runner binds an `ARDiffusionKVState` only around one `forward()` call. The
pipeline may use it during that context but must not retain it. The session
survives subsequent requests with the same `session_id`. Request reset, explicit
close, LRU eviction, and forward exceptions all release runner-owned KV and call
the pipeline's reset/close lifecycle hook. Cross-attention KV is allocated lazily
The runner binds an `ARDiffusionKVState` only around one runner invocation
(`execute_model` or `execute_stepwise`). The pipeline may use it during that
context but must not retain it. Scratch and committed pages stay on the
runner-owned session object, so unbinding does not drop an in-flight stepwise
chunk. Tick-mode sessions survive subsequent requests with the same
`session_id`. On the stepwise path the session id is the request id and is
released when that request finishes or fails. Request reset, explicit close,
LRU eviction, and forward exceptions all release runner-owned KV and call the
pipeline's reset/close lifecycle hook. Cross-attention KV is allocated lazily
per session and logical KV branch, and is released on the same path. A pipeline
publishes one named cache transactionally by yielding exactly one `(k, v)` pair
per layer:
@@ -48,11 +52,12 @@ successful forward. The runner also exposes `reset_session()` and
## Current execution limits
The experimental runner supports request-mode execution with
`max_num_seqs=1`. It rejects step execution and request-batch execution at load
time, and its inherited batch/step entry points also fail explicitly. Supporting
those paths requires a batch-aware state-binding contract rather than silently
bypassing `bind_ar_diffusion_state()`.
The experimental runner supports `max_num_seqs=1`. Request-batch execution is
still rejected. Step execution is allowed only when the loaded pipeline
implements `SupportsStepExecution`; other AR-Diffusion pipelines still fail at
load if `step_execution=True`. Capable pipelines bind KV around each
`execute_stepwise` invocation instead of inheriting the generic runner entry
point unwrapped.
The implementation may retain multiple sessions per model instance. It selects
the largest feasible capacity up to the pipeline's `session_capacity`, based on
+16 -7
View File
@@ -323,18 +323,24 @@ inputs, fields, and ownership rules are described below.
```mermaid
flowchart TB
layer1["Layer 1 · Authoring inputs<br/>PipelineConfig + DeployConfig"]
layer2["Layer 2 · Resolve once<br/>StageConfigFactory.create_from_model()<br/>VllmOmniConfig.from_pipeline_config()"]
layer3["Layer 3 · Transport-safe control plane<br/>VllmOmniConfig"]
layer2["Layer 2 · Production resolution boundary<br/>resolve_omni_config()<br/>(StageConfigFactory + VllmOmniConfig internally)"]
layer3["Layer 3 · Startup hand-off<br/>OmniConfigResolution<br/>PipelineConfig + temporary OmegaConf stage bridge"]
layer4["Layer 4 · Runtime launch planning<br/>StageRuntime"]
layer5["Layer 5 · Engine materialization<br/>VllmConfig / OmniDiffusionConfig"]
layer1 --> layer2 --> layer3 --> layer4 --> layer5
```
The single resolution boundary is implemented by
`StageConfigFactory.create_from_model()` and
`VllmOmniConfig.from_pipeline_config()` in
The production resolution boundary is `resolve_omni_config()` in
[`vllm_omni/config`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/config).
It delegates typed construction to `StageConfigFactory.create_from_model()` and
`VllmOmniConfig.from_pipeline_config()`, then returns an
`OmniConfigResolution` consumed by both `AsyncOmniEngine` and headless startup.
Until `StageRuntime` consumes typed stage configs directly, this envelope carries
the effective `PipelineConfig` alongside OmegaConf-compatible `stage_configs` as
a temporary runtime bridge; both views describe the same resolved topology,
including injected stages.
The legacy `stage_args` YAML path has been removed. Model topology now resolves
through `PipelineConfig`, with runtime overrides supplied by `DeployConfig`.
@@ -354,9 +360,12 @@ The important ownership rules are:
deploy-time defaults; it does not redefine the model graph.
3. CLI and Python overrides are applied at the resolution boundary, with
per-stage overrides taking precedence over global values where supported.
4. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan`
4. `OmniConfigResolution` is the sole startup hand-off. Its `pipeline_config`
and temporary `stage_configs` compatibility view must describe the same
topology.
5. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan`
is runtime-private state, not a user configuration object.
5. `VllmConfig` and the enriched `OmniDiffusionConfig` are materialized in the
6. `VllmConfig` and the enriched `OmniDiffusionConfig` are materialized in the
process that owns the corresponding engine.
## Main features
+24 -9
View File
@@ -53,6 +53,11 @@ must use distinct hook names and remove only hooks they own.
## Discovery and topology
`resolve_offload_plan(pipeline, config)` is the single topology entry point.
Every backend calls it once when `enable()` starts and then consumes the
resolved components, block stacks, resident heads, and staged-residency flags
it returns. Backends do not re-read model declarations.
Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
- `_dit_modules`;
@@ -61,16 +66,26 @@ Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
- `_resident_modules`.
Dotted paths are supported. Legacy pipelines may use the fallback scan of
well-known attribute names, but new integrations should declare components
explicitly.
well-known attribute names; it now warns once per pipeline class, and new
integrations must declare components explicitly.
Both layerwise backends first consume pipeline `OffloadPlan` metadata. A
plan's `block_attrs` maps each DiT path to its ordered block containers, while
`encoder_block_attrs` declares streamable encoder stacks. DiTs absent from the
plan fall back to `_layerwise_offload_blocks_attrs` (including the deprecated
singular-name compatibility path). Discovery metadata describes structure
only; the backend remains responsible for transfer, synchronization, and
storage ownership.
Block topology comes from the pipeline `OffloadPlan` first: `block_attrs` maps
each DiT path to its ordered block containers, `encoder_block_attrs` declares
streamable encoder stacks, `on_demand_component_paths` marks pipeline-managed
residency, `resident_dit_paths` marks the DiTs that may hold resident layers,
and `encoder_dlo_weight_replication` marks the encoders whose loader-produced
weights are safe for AllGather. DiTs absent from the plan fall back to
`_layerwise_offload_blocks_attrs` (including the deprecated singular-name
compatibility path).
The resolver is pure: it moves no tensor, installs no hook, writes no module
attribute, and reads no process group — multi-rank facts come from
`OffloadConfig`. It owns topology validation, so a configuration the model
cannot serve fails before the first component is placed. Explicit component
selection turns every declaration problem into an error, while the
compatibility topology (no component selector) keeps the historical
warn-and-skip behavior. Metadata still describes structure only; the backend
remains responsible for transfer, synchronization, and storage ownership.
## Cross-strategy invariants
@@ -0,0 +1,172 @@
# Runner-to-model prefill/decode phase contract
## Overview
Autoregressive models with a custom `preprocess` hook need to distinguish
prompt processing from decoding. A one-token span can be either a complete
one-token prompt, the last chunk of a longer prompt, or a decode step. Using
span length as the phase signal can replace a prompt embedding with a generated
audio embedding and advance teacher-forced text too early without an exception.
[Issue #4382](https://github.com/vllm-project/vllm-omni/issues/4382) requests a
stable contract for the metadata and routing guards introduced by
[PR #3662](https://github.com/vllm-project/vllm-omni/pull/3662).
The shared runner already implements those guards. This design documents the
existing interface and adds runner-level regression coverage rather than
introducing a second phase detector or changing scheduling policy.
## Scope and objectives
- Make the three existing phase keys safe to depend on in out-of-tree models.
- Preserve prompt embeddings and model state on a one-token prefill tail.
- Keep normal and batched decode routing consistent in mixed batches.
- Explain the correctness workaround required on older releases.
This contract does not enable chunked prefill for models whose other state
handling does not support it. It does not change the scheduler, introduce new
kwargs, or standardize other `_omni_*` fields. The implementation adds
documentation and regression coverage; no hot-path allocation is added.
## Design
### Ownership and data flow
```text
Scheduler / current input batch
prompt_token_ids + num_computed_tokens_cpu (before this step)
|
OmniGPUModelRunner._preprocess
|
per-request phase metadata
/ \
prefill / other spans single-token decode
preprocess preprocess or preprocess_decode_batch
| |
prompt embedding runner-managed talker_mtp
\ /
model forward
```
The producer is `OmniGPUModelRunner._preprocess` in
`vllm_omni/worker/gpu_model_runner.py`. `GPUARModelRunner` and
`GPUGenerationModelRunner` inherit it. The current `OmniNPUModelRunner` also
inherits this method; device-specific execution still needs its own hardware
validation. A plugin runner that overrides preprocessing is responsible for
preserving this interface.
### Stable model-facing fields
The following names, types, and meanings are a supported interface for in-tree
and out-of-tree model hooks despite the leading underscore. Models should treat
them as read-only runner-owned values. Changes to their meanings or removal
require an explicit compatibility/migration plan.
| Key | Python type | Meaning |
| --- | --- | --- |
| `_omni_prompt_len` | `int` | Length of the current request state's `prompt_token_ids`. This is the stage's prompt length, not the current chunk size or raw text length. |
| `_omni_num_computed_tokens` | `int` | The row's `input_batch.num_computed_tokens_cpu` value **before** processing the scheduled span. It includes already-computed/cached progress, not just tokens computed in the current call. |
| `_omni_is_prefill` | `bool` | Exactly `_omni_num_computed_tokens < _omni_prompt_len`. Equality and greater progress mean decode. |
The runner refreshes the fields after current request state has been gathered,
immediately before selecting the per-row hook. Existing values in intermediate
model state cannot override the current phase. Row indices follow the current
input batch (including any reordering), not request arrival order. A resumed
or streaming request is classified from its **current** prompt/progress values;
the phase is not latched for the lifetime of the request.
The normal hook receives these fields as keyword arguments:
```python
def preprocess(self, input_ids, input_embeds, **info):
if info["_omni_is_prefill"]:
offset = info["_omni_num_computed_tokens"]
# Select the scheduled slice of the model's prompt embeddings.
...
else:
# Advance decode-only text/audio state.
...
```
For `preprocess_decode_batch(input_ids=..., req_infos=...)`, each `req_infos`
entry contains the same fields, aligned with its corresponding input row.
Every such row has `_omni_is_prefill == False`. The optional earlier
`preprocess_batch` hook is a separate whole-batch preparation hook; these freshly
computed per-row fields are not promised at that earlier boundary.
### Routing guarantee
With `has_preprocess=True`, the batched decode hook is eligible only when it is
callable, the runner has `talker_mtp`, the scheduled span is one token, and
`_omni_is_prefill` is false. Otherwise the normal `preprocess` hook runs.
Runner-managed MTP is likewise restricted to single-token decode rows. A prefill
row, including a one-token tail, is never added to that MTP batch: the runner
does not consume its `mtp_inputs`, generate codes for it, or overwrite its
preprocessed embedding with an MTP embedding. Decode rows in the same step
continue to receive their MTP outputs at their own token offsets. Models may
still implement their own model-specific work inside `preprocess`.
| Prompt length | Computed before step | Scheduled span | Phase | Runner MTP eligible |
| --- | --- | --- | --- | --- |
| 5 | 0 | 4 | Prefill | No |
| 5 | 4 | 1 | Prefill tail | No |
| 5 | 5 | 1 | Decode | Yes |
| 5 | 6 | 1 | Decode | Yes |
| 1 | 0 | 1 | Prefill | No |
| 5 | 5 | 2 | Decode | No (not a single-token span) |
### Compatibility
!!! warning "vLLM-Omni 0.20.x and earlier"
AR talker models with custom preprocessing that infer phase from span length
must use `enable_chunked_prefill: false` for correctness on these releases.
They lack the phase metadata and the runner's prefill-tail MTP protection.
A model-only heuristic cannot prevent the older runner from overwriting a
one-token prompt embedding. Disabling chunking avoids split tails; it does
not make span length a reliable phase detector for every possible prompt.
The mechanism landed in #3662 (v0.21.0rc2 / v0.22.0 and later); this document
formalizes that existing behavior. Prefer upgrading the matching vLLM and
vLLM-Omni release pair and consuming the explicit phase keys. Do not silently
fall back to `input_ids.numel() > 1` when the model requires chunking correctness.
For a custom or backported runtime, verify both metadata and MTP guards.
## Correctness and testing plan
Runner tests call the production `_preprocess` and `_talker_mtp_forward` methods
with CPU buffers and a deterministic MTP model. They inspect kwargs at the model
boundary, resulting embeddings, and per-request generated codes. Coverage must
include a multi-step prompt split with a one-token tail followed by decode,
mixed prefill/decode batches in different orders, normal/batched decode hooks,
adjacent decode requests in the same hook call with exact per-request codes,
single-token prompts, cached progress, and decode spans longer than one token.
No test should pass merely because MTP was disabled for every row.
Real-weight validation uses Qwen3-TTS CustomVoice through the full talker and
code2wav pipeline. Record the actual scheduled prompt progress, assert that at
least one one-token prefill tail was observed, check its hook metadata and MTP
exclusion, and require finite non-empty output audio. Compare to an unchunked
run with the same real input, greedy sampling, and `VLLM_BATCH_INVARIANT=1`.
Greedy sampling alone does not guarantee identical codes across different
prefill shapes: floating-point reduction differences can change an argmax and
diverge in subsequent autoregressive steps. Invariant kernels keep this exact
comparison focused on the phase contract. The two engines must be initialized
separately because their chunking settings and token budgets differ. GPU tests belong under
`tests/e2e/features/`; existing model tests remain useful adjacent regression
coverage. Functional audio checks do not establish throughput or all-model
semantic parity; this change makes no performance claim.
## Open questions and discussions
There are no new field names or execution modes to choose. Future runner
implementations, including Model Runner V2, should preserve these semantics or
provide an explicit migration path. Hardware validation beyond the tested
platform remains separate from the shared Python routing tests.
## References
- [Issue #4382](https://github.com/vllm-project/vllm-omni/issues/4382)
- [Original metadata and guards, #3662](https://github.com/vllm-project/vllm-omni/pull/3662)
- [Qwen2.5-Omni consumer fix, #5019](https://github.com/vllm-project/vllm-omni/pull/5019)
- [Adding an Omni model](../../contributing/model/adding_omni_model.md)
- [Test writing guide](../../contributing/ci/test_writing_guide.md)
+4 -4
View File
@@ -66,7 +66,8 @@ The checkpoint keeps these verified contracts:
- scheduler data-plane append over a resumable request;
- stale epoch/turn/response fencing;
- `/v1/duplex` and OpenAI Realtime projection;
- optional OpenAI `server_vad` turn detection backed by per-session Silero VAD;
- optional OpenAI `server_vad` turn detection backed by a process-shared Silero
model and per-session VAD state;
- existing JoyVL behavior.
The checkpoint does not claim:
@@ -652,7 +653,7 @@ EOS decisions advance the model conversation.
Explicit Realtime clients may still use `input_audio_buffer.commit` to create a
conversation item. The translator validates that wire input before producing a
commit carrying `realtime_item_id`. Native auto-response may already have
commit carrying `item_id`. Native auto-response may already have
streamed those PCM samples into the runtime, so the runner accepts a validated
commit even when no runtime-side chunk remains. A truly empty explicit buffer
continues to return `input_audio_buffer_empty`.
@@ -677,8 +678,7 @@ multi-turn runtime evidence.
During auto-response overlap, `preserve_realtime_input` distinguishes "do not
append this silent chunk to the native buffer" from "clear the open Realtime
item". Silent overlap no longer discards earlier user PCM. A separate explicit
`server_vad` policy is documented in the Realtime contract above.
item". Silent overlap no longer discards earlier user PCM.
The first chunk of one overlapping input item also reserves its target model
turn. A later Realtime commit uses that reserved identity even if response EOS
+1
View File
@@ -24,6 +24,7 @@ implementation contract; it is not, by itself, a general support claim.
- [Async Chunk](feature/async_chunk.md)
- [Async Diffusion Output](feature/async_diffusion_output.md)
- [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
- [Runner-to-model Prefill/Decode Phase Contract](feature/preprocess_phase_contract.md)
- [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
- [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md)
@@ -8,22 +8,44 @@ owners:
- "@SamitHuang"
- "@fhfuih"
primary_code_paths:
- vllm_omni/diffusion/diffusion_engine.py
- vllm_omni/diffusion/request.py
- vllm_omni/diffusion/data.py
- vllm_omni/diffusion/sched/**
- vllm_omni/diffusion/executor/**
- vllm_omni/diffusion/worker/**
related_code_paths:
- vllm_omni/diffusion/stage_diffusion_client.py
- vllm_omni/diffusion/stage_diffusion_proc.py
- vllm_omni/diffusion/inline_stage_diffusion_client.py
- vllm_omni/diffusion/ipc.py
- vllm_omni/diffusion/hooks/**
- vllm_omni/diffusion/postprocess/**
depends_on:
- ../engine_orchestration.md
- ../stage_runtime.md
- ../input_output_modality_contracts.md
- diffusion_model_integration.md
validation_paths:
- tests/diffusion/diffusion_backend/**
- tests/diffusion/hooks/**
- tests/diffusion/test_diffusion_engine.py
- tests/diffusion/test_diffusion_engine_cleanup.py
- tests/diffusion/test_diffusion_engine_dummy_run.py
- tests/diffusion/test_diffusion_engine_rpc_routing.py
- tests/diffusion/test_diffusion_scheduler.py
- tests/diffusion/test_diffusion_model_runner.py
- tests/diffusion/test_diffusion_worker.py
- tests/diffusion/test_multiproc_engine_concurrency.py
- tests/diffusion/test_uniproc_executor.py
- tests/diffusion/test_result_pump.py
- tests/diffusion/test_async_output_worker.py
- tests/diffusion/test_diffusion_ipc.py
- tests/diffusion/test_diffusion_streaming_output.py
- tests/diffusion/test_stage_diffusion_proc.py
- tests/diffusion/test_inline_stage_diffusion_client.py
upstream_refs:
- diffusers.DiffusionPipeline
last_reviewed: 2026-07-16
last_reviewed: 2026-09-02
last_verified_commit: e51fe6ec1b9a9a0e14bb1fdb296d61b6593b93c6
---
# Diffusion runtime
@@ -31,6 +53,380 @@ last_reviewed: 2026-07-16
The diffusion runtime owns request admission, scheduling, execution, progress,
output, cancellation, and cleanup inside a diffusion stage.
It is easiest to think of the runtime as a small control loop:
1. the **scheduler** chooses ready requests;
2. the **executor** sends that choice to the worker layer;
3. each **worker** manages its device and delegates model work to a **runner**;
4. the runner calls the model pipeline and returns per-request results; and
5. the engine feeds those results back to the scheduler and output streams.
!!! note
This page describes the runtime plumbing. Model loading and pipeline
contracts belong to [Diffusion model integration](diffusion_model_integration.md).
Batch compatibility and step-state details belong to
[Diffusion continuous batching](continuous_batching.md).
## Goals and non-goals
The runtime aims to keep request policy separate from device execution. It
also keeps request identity stable from admission through the final output,
including cancellation and failures.
This module does **not** choose the next Omni stage, place stage replicas, or
define a model's denoising algorithm:
- cross-stage routing belongs to [Engine orchestration](../engine_orchestration.md);
- process placement and replica lifecycle belong to
[Stage runtime](../stage_runtime.md); and
- model-specific work belongs to
[Diffusion model integration](diffusion_model_integration.md).
## Runtime at a glance
The stage owns `DiffusionEngine` and a `DiffusionExecutor` backend. The inline
stage client keeps that stack in the caller process; a process-backed stage
runs it inside `StageDiffusionProc`.
The executor then chooses how the worker runs:
- **`uni`** (default when `num_gpus == 1`): `UniProcDiffusionExecutor` builds
one in-process worker. There is no worker subprocess and no shared-memory
RPC.
- **`mp`** (default when `num_gpus > 1`, or when set explicitly):
`MultiprocDiffusionExecutor` starts one `WorkerProc` per device.
Either way, each worker owns its device, distributed state,
`DiffusionWorker`, `DiffusionModelRunner`, and pipeline instance.
```mermaid
flowchart TB
client["Stage client"]
subgraph stage["Diffusion stage process"]
stageProc["StageDiffusionProc<br/>stage transport and lifecycle"]
engine["DiffusionEngine<br/>admission and control loop"]
scheduler["RequestScheduler or StepScheduler<br/>request state and policy"]
executor["DiffusionExecutor<br/>uni or mp backend"]
streams["Per-request output streams"]
stageProc --> engine
engine --> scheduler
scheduler -->|"DiffusionSchedulerOutput"| engine
engine --> executor
engine --> streams
end
subgraph workers["Worker layer"]
workerProc["WorkerProc<br/>mp IPC loop only"]
worker["DiffusionWorker<br/>device and distributed setup"]
runner["DiffusionModelRunner<br/>model state and execution"]
pipeline["Diffusion pipeline<br/>model-specific computation"]
workerProc --> worker
worker --> runner --> pipeline
end
client -->|"process-backed: ZMQ"| stageProc
client -.->|"inline: direct call"| engine
executor -->|"mp: broadcast / result queues"| workerProc
executor -.->|"uni: in-process call"| worker
streams -->|"process-backed"| stageProc
stageProc -->|"ZMQ results"| client
streams -.->|"inline"| client
```
!!! info "Stage process vs worker process"
A deployed stage may already run in a `StageDiffusionProc` child process.
That is separate from the optional `mp` worker processes. With `uni`, the
worker lives in the same process as the engine.
## Component responsibilities
| Component | Owns | Does not own |
| --- | --- | --- |
| `DiffusionEngine` | Admission, the busy loop, scheduler/executor coordination, RPC ordering, output streams, warmup, abort routing | Batch policy details, device setup, model computation |
| `BaseScheduler` | Request states, waiting/running sets, capacity, compatibility, terminal transitions | IPC, worker calls, output formatting |
| `RequestScheduler` | Complete-request waves and optional admission delay for request-level batching | Denoising-step progress |
| `StepScheduler` | Per-request denoising progress and step-wise completion | Pipeline tensors and model state |
| `DiffusionExecutor` | Execution backend contract, worker RPC, health, shutdown | Admission and request-state transitions |
| `UniProcDiffusionExecutor` | Single in-process worker for `num_gpus == 1`; no IPC or async output pump | Multi-GPU execution |
| `MultiprocDiffusionExecutor` | Worker processes, shared-memory message queues, result dispatch, worker monitoring | In-process single-GPU path |
| `WorkerProc` | One `mp` worker process's IPC loop and reply rules | Scheduling policy; unused by `uni` |
| `DiffusionWorker` | Device/distributed setup, LoRA activation, profiling, sleep/wake, runner delegation | Request admission |
| `DiffusionModelRunner` | Pipeline loading, request-local model state, cache/compile setup, request or step execution | Queueing and cross-stage routing |
The boundary between policy and execution is
`DiffusionSchedulerOutput`. It contains newly admitted request payloads,
IDs for requests whose runner state is already cached, finished IDs,
optional KV-prefetch work, and per-request `diffusion_kv_metadata` on
`NewRequestData` envelopes when paged Diffusion KV is enabled. Executors and
workers consume this output; they do not decide what should run next.
`BaseScheduler.initialize()` sets `max_num_running_reqs` from
`od_config.max_num_seqs` (default 1). The engine may override that capacity
for distributed layerwise offload with AllGather DP concurrency.
## One scheduling tick
Async requests share one engine busy loop. Worker calls and control RPCs pass
through that loop, so they cannot race on the executor transport.
```mermaid
sequenceDiagram
participant Caller
participant Engine as DiffusionEngine
participant Scheduler
participant Executor
participant Worker
participant Runner
Caller->>Engine: add request
Engine->>Scheduler: add_request(request)
Engine->>Scheduler: schedule()
Scheduler-->>Engine: DiffusionSchedulerOutput
Engine->>Executor: execute_batch() or execute_step()
Executor->>Worker: run worker method
Worker->>Runner: execute model or one step
Runner-->>Worker: RunnerOutput(s)
Worker-->>Executor: result
Executor-->>Engine: BaseRunnerOutput
Engine->>Scheduler: update_from_output(...)
Scheduler-->>Engine: finished request IDs
Engine-->>Caller: chunk or final output
```
The engine catches request execution errors and turns them into per-request
error outputs. A dead worker group is different: the executor marks itself
failed, health checks raise `EngineDeadError`, and the owning stage client
handles the stage-level failure.
## Execution modes
The engine resolves one mode at startup and binds the matching scheduler and
executor call.
| Mode | Scheduler | Executor call | Runner path |
| --- | --- | --- | --- |
| Request batch | `RequestScheduler` | `execute_batch()` | `execute_model()` for one request, or `execute_model_batch()` for a fused batch |
| Step batch | `StepScheduler` | `execute_step()` | `execute_stepwise()` |
**Request mode** runs a complete pipeline forward for each scheduled wave. A
single-request wave is the conservative path. A multi-request wave usually
uses fused `execute_model_batch()` and requires the pipeline to declare
request-batch support.
**Distributed layerwise offload with AllGather (DLO DP concurrency)** is a
separate multiproc dispatch path. It activates when
`data_parallel_size > 1`, `enable_distributed_layerwise_offload` is set, and
`dlo_use_allgather` is true. The engine then sets `dp_concurrent = True` and
raises `scheduler.max_num_running_reqs` to `dp_size`, overriding the
`max_num_seqs` value from `initialize()`. The scheduler still emits a
multi-request `DiffusionSchedulerOutput`; the multiproc executor routes that
wave through `execute_request()` instead of fused `execute_model_batch()`.
Optional admission coalescing can fill a `dp_size` wave: when
`request_batch_max_wait_ms > 0`, `RequestScheduler` may wait briefly before
the first schedule of a wave (under `dp_concurrent`, the stable window is
`min(0.3s, wait/2)`). With the default `request_batch_max_wait_ms == 0`,
there is no wait.
Before dispatch, `MultiprocDiffusionExecutor.execute_request()` rejects waves
whose requests differ in sampling-parameter compatibility or `extra_args`
(AllGather requires every DP rank to follow the same forward schedule). Each
worker rank then picks one envelope from the wave:
`req[dp_rank % len(req)]`, so every rank enters the same collective while
computing different requests.
**Step mode** keeps `StepRequestState` in the runner. New requests run
`prepare_encode()` once; every tick runs `denoise_step()` and
`step_scheduler()`; completed requests run `post_decode()`. The scheduler keeps
only lifecycle and progress metadata—it does not hold model tensors.
**Step mode is also the path for streaming diffusion output**. Enabling
`streaming_output` turns on step execution if needed. Chunk-capable pipelines
can emit intermediate outputs through the same request stream; final-only step
pipelines still emit only the final result. If the pipeline does not implement
step execution, initialization fails.
!!! tip
For user-facing flags and examples, see
[Diffusion execution modes](../../../user_guide/diffusion/execution_modes.md).
That guide explains how to select a mode; this page explains what happens
after the choice is made.
## Request state and identity
The scheduler is the source of truth for an admitted request. Its normal state
flow is:
```mermaid
stateDiagram-v2
[*] --> WAITING: add_request
WAITING --> RUNNING: schedule
WAITING --> FINISHED_ABORTED: abort
RUNNING --> FINISHED_ABORTED: abort
RUNNING --> FINISHED_COMPLETED: successful output
RUNNING --> FINISHED_ERROR: failed or missing output
FINISHED_COMPLETED --> [*]: deliver and remove state
FINISHED_ABORTED --> [*]: deliver and remove state
FINISHED_ERROR --> [*]: deliver and remove state
```
`PREEMPTED` and `BaseScheduler.preempt_request()` exist on the scheduler API,
but the engine loop does not call them today—only scheduler unit tests do.
`request_id` ties together scheduler state, runner state, executor results, and
the output queue. Batch positions are temporary and must never replace request
identity when results are mapped back.
Only compatible requests enter the same running wave. The scheduler compares a
mode-specific sampling-parameter key and stops at the first incompatible
waiting request. This is intentionally conservative and can cause
head-of-line blocking.
## Executor and IPC
`DiffusionExecutor` is the backend interface. Built-in backends are selected by
`distributed_executor_backend`:
| Backend | When chosen | Worker layout |
| --- | --- | --- |
| `uni` | Default for `num_gpus == 1`, or set explicitly | One in-process worker |
| `mp` | Default for `num_gpus > 1`, or set explicitly | One worker process per device |
Ray and external-launcher diffusion backends are not implemented. A custom
`DiffusionExecutor` subclass or import path is also accepted.
### Uniproc backend
`UniProcDiffusionExecutor` constructs `WorkerWrapperBase` in the engine process
and calls worker methods directly. It avoids a second model load, MessageQueue
rings, ZMQ IPC sockets, and `/dev/shm` tensor packing. RPC timeouts are
accepted for interface parity but are not enforced: a hung worker blocks the
calling thread. Sticky accelerator faults mark the executor dead through
failure callbacks; ordinary request errors stay per-request.
### Multiproc backend
`MultiprocDiffusionExecutor`:
1. creates a broadcast queue shared by all workers;
2. spawns one worker process per configured device;
3. waits until every worker reports ready;
4. sends generation and control calls as RPC messages;
5. applies rank-aware reply rules so only expected ranks respond; and
6. monitors worker process sentinels and fails the executor if a worker dies.
In request mode on this backend, each async output produces two correlated
messages keyed by `async_output_id`:
- `COMPUTE_DONE` — the forward finished and the device can start the next
request.
- `OUTPUT_READY` — background D2H/SHM packing finished and the final output
is ready.
The worker queues background packing work before it enqueues `COMPUTE_DONE`,
but either message may reach the executor first. The result pump and
`execute_batch()` must handle both orderings. Step mode keeps the synchronous
result path and does not start those pumps.
See [Async diffusion output](../../feature/async_diffusion_output.md) for the
multiproc request-mode timeline.
## Worker and runner boundary
On the `mp` path, `WorkerProc` receives messages and calls methods through
`WorkerWrapperBase`. On the `uni` path, the executor owns that wrapper
directly. In both cases, `DiffusionWorker` handles the parts tied to a device:
distributed initialization, model-runner construction, LoRA activation,
profiling, and memory sleep/wake. It delegates actual model work to
`DiffusionModelRunner`.
The runner owns long-lived model-side state:
- the loaded pipeline and compilation setup;
- cache and offload integration;
- random-generator setup;
- request batches for complete forwards; and
- `StepRequestState` plus `InputBatch` for step execution.
This split keeps infrastructure out of model pipelines. A pipeline receives a
request batch or step state and performs model work; it should not inspect
engine queues or change scheduler state.
## Lifecycle and cleanup
### Startup
`DiffusionEngine.make_engine()` resolves the engine class, constructs the
executor and scheduler, then normally runs a small dummy request to warm up the
model.
Warmup is skipped only when distributed layerwise offload with AllGather
is enabled (via `enable_distributed_layerwise_offload` and `dlo_use_allgather`)
and `max(data_parallel_size, sequence_parallel_size) > 1`, because
the dummy run sends one request while AllGather needs every shard rank to
enter the same collective. If initialization or warmup fails, the runtime
closes scheduler state and worker resources before returning the error.
### Cancellation
`abort()` places request IDs on an abort queue. The busy loop marks those
requests `FINISHED_ABORTED`; finalization removes scheduler state and emits an
aborted output when a consumer still exists. The runner removes cached
step-mode state when it receives a scheduler output that reports the finished
request ID.
On the multiproc request-mode async path, abort or consumer drop while output
is still materializing should release the associated async-output bookkeeping
so late results are not retained after the request ends
*(pending #6253/#6439/#6580)*. Today, an unconsumed `OUTPUT_READY` is cached
by design until `wait_output_ready()` or executor teardown. Dropping an
output-stream consumer removes that consumer's queue. It does not take
ownership of scheduler cleanup.
### Shutdown and worker failure
`DiffusionEngine.close()` is idempotent. It stops the busy loop, wakes pending
streams with an error, closes scheduler state, and shuts down the executor.
- **`mp`**: the executor asks workers to stop, waits for them, then terminates
workers that miss the grace period. Shutdown marks unfinished RPC and
async-output futures as failed, then clears them. Completed async outputs
cached for later `wait_output_ready()` calls are dropped with the executor
object.
- **`uni`**: the executor shuts down the in-process worker, drops the final
model reference, and empties the accelerator cache so a later engine can reuse
the device.
!!! warning
Do not continue using a worker group after a fatal collective timeout or
unexpected worker exit. Distributed state may be incomplete. The
multiprocess executor deliberately fails closed so the stage can restart.
The uniproc executor fails closed when the accelerator context itself is
poisoned.
## Extension boundaries
- A custom engine may set `default_diffusion_model_runner_cls`; an explicit
`diffusion_model_runner_cls` configuration still wins.
- Tests and custom engine integrations may inject a `BaseScheduler` subclass.
`SchedulerInterface` remains only as a deprecated compatibility name.
- `distributed_executor_backend` may be `"uni"`, `"mp"`, a custom `DiffusionExecutor`
subclass, or an import path. A backend must preserve scheduler output,
request identity, health, and cleanup contracts. Use `"mp"` when you need
process isolation or multi-GPU; `"uni"` is the single-GPU default.
- Worker extensions go through `WorkerWrapperBase`. They add worker methods;
they do not become a second scheduler or request lifecycle.
These are advanced Python integration points, not stable end-user CLI
customization surfaces.
## Candidate invariants
### DIFF-RUNTIME-INV-001: One lifecycle owner
@@ -41,19 +437,50 @@ until completion, cancellation, or failure.
### DIFF-RUNTIME-INV-002: Execution follows scheduler output
**Rule:** Executors and workers MUST execute scheduler decisions without
admitting, reordering, or forwarding requests independently.
admitting, reordering, or forwarding requests independently. Distributed
layerwise offload with AllGather is an exception at dispatch time: the
scheduler may schedule a multi-request wave, but each DP rank executes one
envelope from that wave so every rank enters the same collective schedule.
### DIFF-RUNTIME-INV-003: Terminal cleanup is complete
**Rule:** Every terminal path MUST release request state, temporary tensors,
hooks, and runtime-owned resources.
hooks, and runtime-owned resources. Async-output bookkeeping on abort or
consumer drop before consumption is also required *(pending #6253/#6439/#6580)*.
### DIFF-RUNTIME-INV-004: Optional features use runtime hooks
**Rule:** Cache, profiling, offload, and parallel features SHOULD integrate at
defined hooks instead of creating another request lifecycle.
### DIFF-RUNTIME-INV-005: Results keep request identity
**Rule:** Batched and asynchronous results MUST be mapped by stable request ID,
not by assumed completion order.
## Safe-change guide
Test success, cancellation, failure, shutdown, multiple requests, and resource
cleanup.
Test the smallest affected slice, then cover its neighboring boundary:
| Change area | Minimum evidence |
| --- | --- |
| Admission or state transitions | `test_diffusion_scheduler.py`, including duplicate IDs, compatibility, abort, and missing output |
| Engine loop or output delivery | `test_diffusion_engine.py`, `test_diffusion_engine_cleanup.py`, `test_diffusion_engine_rpc_routing.py` |
| Executor or IPC | `test_multiproc_engine_concurrency.py`, `test_uniproc_executor.py`, `test_result_pump.py`, `test_diffusion_ipc.py`, `test_async_output_worker.py` |
| Worker or runner | `test_diffusion_worker.py`, `test_diffusion_model_runner.py` |
| Stage boundary | `test_stage_diffusion_proc.py`, `test_inline_stage_diffusion_client.py` |
| Warmup or streaming | `test_diffusion_engine_dummy_run.py`, `test_diffusion_streaming_output.py` |
Always exercise success, cancellation, per-request failure, fatal worker
failure, repeated shutdown, one request, and multiple compatible requests.
For step-mode changes, also test partial progress and cleanup of runner state.
## Related documents
- [Diffusion module overview](index.md)
- [Diffusion model integration](diffusion_model_integration.md)
- [Diffusion continuous batching](continuous_batching.md)
- [Diffusion execution modes](../../../user_guide/diffusion/execution_modes.md)
- [Async diffusion output](../../feature/async_diffusion_output.md)
- [Engine orchestration](../engine_orchestration.md)
- [Stage runtime](../stage_runtime.md)
+1
View File
@@ -124,6 +124,7 @@ def extract_omni_serve_subparser_init():
"make_arg_parser": lambda parser: parser, # no-op for doc
"_ensure_vllm_platform": lambda: None, # no-op for doc
"_nonneg_finite_float": float,
"_parse_stage_overrides": __import__("json").loads,
"VLLM_SUBCMD_PARSER_EPILOG": "",
"logger": logger,
"DummySubparsers": DummySubparsers,
+2
View File
@@ -64,6 +64,7 @@ th {
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `LongCatVideoAvatarPipeline` | LongCat-Video-Avatar-1.5 A2V/AI2V (native single-speaker and multi-speaker AI2V/AVC) | `meituan-longcat/LongCat-Video-Avatar-1.5` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md) |
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, `Boogu/Boogu-Image-0.1-Edit-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) |
| `BooguImageTurboPipeline` | Boogu-Image Turbo distilled T2I | `Boogu/Boogu-Image-0.1-Turbo` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) |
| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — |
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
@@ -96,6 +97,7 @@ th {
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
| `Audio8TTSSlowARForConditionalGeneration` | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✅︎ | | | | — |
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only), SenseNova-U1.5 | `SenseNova/SenseNova-U1-8B-MoT`, `sensenova/SenseNova-U1.5-8B-MoT` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.5.md) |
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
+7 -1
View File
@@ -64,7 +64,7 @@ connect
<- msgpack(policy_server_config)
infer
-> msgpack({"endpoint": "infer", "session_id": "...", ...observation})
-> msgpack({"endpoint": "infer", "session_id": "...", "seed": <optional int>, ...observation})
<- msgpack(ndarray | dict[str, ndarray])
reset
@@ -92,6 +92,12 @@ frames are not observation messages.
- The policy pipeline owns observation transforms and persistent model state,
normally keyed by `session_id`; the API layer forwards the raw observation
dictionary.
- An optional integer `seed` in an inference message becomes the engine
request's `sampling_params.seed`. A policy whose sampling consumes the
request generator (GR00T-N1.7) returns the same action chunk for the same
observation and seed. If omitted, the engine assigns a random per-request
seed, so unseeded inferences are independent of each other. The API layer
consumes `seed`; it is not forwarded inside the observation.
## Limits and Errors
+24 -17
View File
@@ -1,17 +1,17 @@
# Realtime Duplex API
vLLM-Omni serves full-duplex speech models over a WebSocket endpoint,
`/v1/realtime?duplex=1`, that speaks the OpenAI Realtime vocabulary plus a
`/v1/realtime`, that speaks the OpenAI Realtime vocabulary plus a
set of duplex extensions: the model decides when to listen and when to speak,
the user can talk over the assistant, playback progress is reported back so
history stays honest, and a dropped connection can resume the same session.
This page covers how to run a duplex deployment, how to drive it from Python
with `vllm_omni.clients.duplex.DuplexClient`, and the complete wire contract.
The endpoint is served for models that ship a duplex plugin — currently
MiniCPM-o 4.5, PersonaPlex, and Nemotron VoiceChat — and only for deploy
configurations that declare `session_mode: duplex`. The runtime architecture
is described in [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md).
The handler is enabled by either a model-native `session_mode: duplex`
deployment or an explicit `duplex_session` configuration for turn-based
Server VAD. The runtime architecture is described in
[Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md).
## Quick Start
@@ -27,7 +27,7 @@ vllm-omni serve openbmb/MiniCPM-o-4_5 \
`vllm_omni/deploy/minicpmo_4_5.yaml` declares `session_mode: duplex` and
`duplex_session.max_sessions: 4`, which mounts the WebSocket routes
`ws://<host>:8099/v1/realtime?duplex=1` (this page) and `ws://<host>:8099/v1/duplex`
`ws://<host>:8099/v1/realtime` (this page) and `ws://<host>:8099/v1/duplex`
(the native dialect for raw-protocol clients) next to the usual HTTP API.
PersonaPlex and Nemotron VoiceChat use `vllm_omni/deploy/personaplex.yaml` and
`vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml` in the same way.
@@ -240,7 +240,7 @@ on the same helpers.
## API Reference
This section is the complete wire contract of `/v1/realtime?duplex=1`: how
This section is the complete wire contract of `/v1/realtime`: how
the dialect relates to the OpenAI Realtime protocol, which surfaces each
model's capability flags gate, a catalogue of every client-to-server and
server-to-client event, and a JSON example for each. `DuplexClient` speaks
@@ -253,10 +253,11 @@ its protocol tests.
### Endpoint and transport
The endpoint is `ws(s)://<host>/v1/realtime?duplex=1`; the `duplex` query
parameter (`1`, `true`, or `on`) selects the duplex session handler, and the
route is mounted only for deployments whose deploy configuration declares
`session_mode: duplex`. Optional query parameters are `model`, `session_id`,
The endpoint is `ws(s)://<host>/v1/realtime`. Bare connections select the
session handler when the deployment declares either `session_mode: duplex`
or `duplex_session`. `?duplex=1` remains a compatibility alias, while
`?duplex=0` selects the legacy Realtime handler. Optional query parameters
are `model`, `session_id`,
`autostart` (`0` means resume-only), `resume`, and `native_duplex`
(`minicpmo45_native_duplex` is accepted as a deprecated alias and folded
into the canonical name). Every message is one JSON object per WebSocket
@@ -342,9 +343,14 @@ Semantic divergences hidden behind shared names:
before any commit; a commit may end in `response.listen` and no
response, and the model may open a response with no commit at all.
OpenAI: commit ⇒ item, `response.create` ⇒ exactly one response.
- `turn_detection.interrupt_response=false` is rejected;
`create_response` is ignored; `semantic_vad` is unsupported; the VAD
runs per session (Silero) and implies `overlap_policy=barge_in_on_speech`.
- With `turn_detection={"type":"server_vad"}`, omitted `interrupt_response`
defaults to `false` on turn-based models such as Qwen3-Omni, enabling
endpointing and automatic responses without interruption. Explicit `true`
is rejected there. Model-native duplex keeps the `true` default and rejects
`false`; its VAD implies `overlap_policy=barge_in_on_speech`.
This capability-dependent default differs from OpenAI's `true` default;
the effective configuration is returned in `session.created` / `session.updated`.
`semantic_vad` is unsupported.
- `rate_limits.updated` is always an empty list (compatibility only).
- Cancellation never reuses a `response_id`; truncation is driven by
`playback.ack` as well as `conversation.item.truncate`.
@@ -1201,9 +1207,10 @@ Realtime projector renames and fans out these events before they reach a
barge-in, or audio truncation; Nemotron VoiceChat does not support barge-in
or audio truncation; camera frames are consumed only by MiniCPM-o 4.5; tool
calls are produced only by Nemotron VoiceChat.
- `turn_detection` supports only `server_vad` with `interrupt_response=true`;
`semantic_vad`, `interrupt_response=false`, and `create_response` are not
supported. Server VAD requires a model-native session.
- `turn_detection` supports `server_vad` or `null`, not `semantic_vad`.
Turn-based models support endpointing without interruption, including
`create_response=false` for explicit response creation. Model-native duplex
requires `interrupt_response=true` and `create_response=true`.
- `input_audio_transcription` and `input_audio_noise_reduction` are accepted
and echoed but no separate transcription or noise-reduction stage runs;
transcripts come from the model.
+31 -12
View File
@@ -279,10 +279,15 @@ curl -X POST http://localhost:8091/v1/audio/voices \
## Streaming Text Input (WebSocket)
The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally and generates audio per sentence as boundaries are detected.
The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally.
By default (`split_granularity=none`) it buffers until `input.done` and
synthesizes that flush as **one** TTS request, which keeps long-form timbre
stable. Set `split_granularity` to `sentence` or `clause` to emit a request
at each detected boundary (lower time-to-first-audio for STT/LLM pipelines).
> Note: text input is always streamed incrementally. Audio output remains sentence-scoped:
> use `stream_audio=false` for one binary frame per sentence, or `stream_audio=true` for one or more PCM chunks per sentence.
> Note: `stream_audio` only changes how **audio bytes** are framed (one WAV/PCM
> payload vs chunked PCM). Text segmentation is controlled separately by
> `split_granularity`.
### WebSocket Protocol
@@ -314,14 +319,16 @@ upstream LLM) pays the WebSocket handshake once instead of once per utterance.
- The session config is sticky. Send `input.text` again straight after
`session.done` to reuse it, or send another `session.config` first to change
voice, format, or reference audio. A `session.config` sent while text is
still buffered is rejected so no pending input is silently dropped.
- An utterance is the flush unit, not a linguistic one: it is whatever text was
buffered when `input.done` arrived, of any length, synthesized as one request.
`utterance_index` counts those flushes across the connection, so it tells you
which `input.done` a frame belongs to. `sentence_index` counts within one
flush and so pairs with `total_sentences`, which means every utterance reports
`sentence_index: 0` of `total_sentences: 1` (or `0` for an empty buffer).
voice, format, or reference audio. A `session.config` sent in the middle of
an utterance is rejected so no pending input is silently dropped and so a
split utterance cannot end up half in one voice and half in another.
- An utterance is the flush unit, not a linguistic one: it is one `input.done`
cycle. `utterance_index` counts those flushes across the connection, so it
tells you which `input.done` a frame belongs to. `sentence_index` counts the
TTS requests inside one flush and so pairs with `total_sentences`: with the
default `split_granularity=none` that is always `sentence_index: 0` of
`total_sentences: 1` (or `0` for an empty buffer), while `sentence` or
`clause` counts the linguistic units actually synthesized.
- End the connection with `session.close`, or by closing the socket. An idle
connection is still closed after the server's idle timeout, which now also
applies to the gap between utterances.
@@ -332,7 +339,15 @@ All REST API parameters are supported, plus:
| Parameter | Type | Default | Description |
| ----------- | ------ | --------- | ------------- |
| `stream_audio` | bool | false | Stream one or more PCM chunks for the buffered input over WebSocket |
| `stream_audio` | bool | false | Stream one or more PCM chunks for each TTS request over WebSocket |
| `split_granularity` | string | `"none"` | `"none"`: one request per `input.done`. `"sentence"`: split on `.!?` plus CJK `。!?…`, Indic danda `।॥`, and Arabic `؟`. `"clause"`: also split on `,;,;،؛`. |
| `seed` | integer | null | Forwarded to the speech engine for this session |
ASCII punctuation only ends a unit when whitespace or `input.done` follows it,
and decimals (`3.14`), thousands separators (`1,000`), abbreviations (`Dr.`,
`e.g.`) and initials (`J. R.`) are not treated as boundaries. A punctuation run
and any closing quote or bracket stay with the unit they close, so `Wait...`
and `He said "Hello."` are one request each.
```bash
DELETE /v1/audio/voices/{name}
@@ -752,6 +767,10 @@ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` ne
| ------- | ------------- |
| `k2-fsa/OmniVoice` | Pure-diffusion TTS. Supports voice cloning via `ref_audio` (with optional `ref_text`); no built-in voice presets. |
OmniVoice uses packed variable-length attention for batched generator execution. The attention operator accepts FP16 and BF16 inputs, so its query,
key, and value tensors are evaluated in BF16 even when the stage is configured with `dtype: float32`; the attention output is converted back to the model's
hidden-state dtype before the output projection. Consequently, a float32 stage configuration does not imply FP32 attention arithmetic.
### VoxCPM2
| Model | Description |
+19
View File
@@ -68,6 +68,8 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
| Parameter | Type | Default | Description |
| ----------- | ------ | --------- | ------------- |
| `input_reference` | file | null | Uploaded reference image or video for image-to-video/video-to-video requests |
| `control_reference` | file | null | Optional uploaded image/video control, up to 512 MiB, for models that declare control-upload support |
| `control_type` | string | null | Model control name associated with `control_reference`; currently Cosmos3 supports `edge`, `blur`, `depth`, `seg`, and `wsm` |
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` or `video_reference` |
| `video_reference` | string | null | JSON-encoded reference video payload; do not combine with `input_reference` or `image_reference` |
| `audio_reference` | string | null | JSON-encoded audio reference for speech-to-video: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
@@ -177,6 +179,23 @@ corresponding addition. `negative_metadata_mode` accepts `same`, `inverse`, or
`none` and defaults to `same` for transfer. See the Cosmos3 recipe for complete
examples.
A client that cannot place the control on the server filesystem can upload one
control with `control_reference` and identify it with `control_type`. The API
streams the upload to request-scoped storage, supplies its path to the model,
and removes it after synchronous or asynchronous generation completes. Other
options for the selected control can remain in `extra_params`; do not also set
`control` or `control_path` there. Uploads larger than 512 MiB are rejected.
```bash
curl -s http://localhost:8091/v1/videos/sync \
-F "prompt=Preserve the scene while following the world-state control" \
-F "input_reference=@input.mp4;type=video/mp4" \
-F "control_reference=@wsm.mp4;type=video/mp4" \
-F "control_type=wsm" \
-F 'extra_params={"wsm":{"control_weight":1.0}}' \
-o output.mp4
```
HTTP redirects for `image_reference.image_url` follow vLLM's
`VLLM_MEDIA_URL_ALLOW_REDIRECTS` setting. Before starting the server, set it to
`1` (the default) to allow redirects or `0` to reject them. A redirect target
+6 -4
View File
@@ -8,7 +8,7 @@ are exposed.
## Choose a Mode
| Goal | CLI configuration |
|---|---|
| --- | --- |
| Serial request execution | `--max-num-seqs 1` |
| Fused request-level batching | `--max-num-seqs N` |
| Single-request step execution | `--step-execution --max-num-seqs 1` |
@@ -177,8 +177,10 @@ wait
The scheduler may batch these requests when their sampling parameters are
compatible. See the
[Image Generation API](../../serving/image_generation_api.md) for response
formats and additional client examples. Streaming models can use
model-specific streaming endpoints documented by their serving guide.
formats and additional client examples.
Streaming models can use model-specific streaming endpoints documented at
[Streaming Video Output API](../../serving/streaming_video_output_api.md).
## Python API
@@ -241,7 +243,7 @@ For step execution, set `step_execution: true` and remove
## CLI Reference
| Flag | Default | Effect |
|---|---:|---|
| --- | ---: | --- |
| `--step-execution` | disabled | Select step-wise scheduling |
| `--max-num-seqs` | `1` for diffusion stages | Set request- or step-scheduler capacity |
| `--request-batch-max-wait-ms` | `0` | Wait for burst coalescing in request mode |
+8 -2
View File
@@ -249,8 +249,14 @@ MiniMax-H3 supports two few-step mechanisms that must not be conflated:
- **Checkpoint-pinned schedule**: a merged release writes `base_schedule` into
`model_index.json`. Requests must pass `num_inference_steps` as the interval
count (for example `4` for `[1.0, 0.7, 0.4, 0.15, 0.0]`).
- **Runtime Turbo LoRA**: LightX2V Turbo artifacts request five sigma points and
enforce `flow_shift=6`, `audio_flow_shift=3`.
- **Runtime Turbo LoRA**: each LightX2V Turbo artifact carries its own contract,
read from its filename. A four-step artifact requests five sigma points and an
eight-step one requests nine; the 768p retrains enforce `flow_shift=6` and the
544p artifacts `flow_shift=12`. `audio_flow_shift=3` across the family. A
request that does not match the loaded artifact is rejected by name. Alpha
comes from the artifact's metadata, or 8 when it declares none. `ref2v`
artifacts require `--task-type ref2va`; `fl2v` artifacts serve T2VA and FL2VA.
The ComfyUI exports use a fused-QKV layout and are not supported.
- **Runtime native LoRA**: FlashGen-style artifacts declare
`key_format=minimax-h3-native` and embed `base_schedule` in safetensors
metadata. When active, the adapter schedule overrides the base checkpoint and
+3 -3
View File
@@ -162,9 +162,9 @@ The following tables show which models support each feature:
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
| **Wan2.2-S2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2.5** | ❌ | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ (Ulysses only) | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ (Ulysses only) | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2.5** | ❌ | ❓ (one-stage) | ✅* (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ |
| **SANA-Video-2B T2V I2V** | ❌ | ❌ | ✅ (`--usp`, frame-sharded) | ✅ | ✅ (TP=2 only) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
@@ -126,11 +126,7 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
## Fish Speech S2 Pro
4B dual-AR text-to-speech model from FishAudio with the DAC codec at 44.1 kHz.
### Prerequisites
```bash
pip install fish-speech
```
No extra packages are required; the DAC codec is vendored in vLLM-Omni.
### Quick start
```bash
@@ -155,13 +155,8 @@ The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-tex
## Fish Speech S2 Pro
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
### Prerequisites
```bash
pip install fish-speech
```
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec, which is vendored in
vLLM-Omni (no extra packages required).
### Launch
@@ -30,7 +30,6 @@ from typing import Any
import numpy as np
_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
_NUM_INFERENCE_STEPS = 4
_MAX_SEQUENCE_LENGTH = 512
_MAX_PIXEL_AREA = 480 * 832
_MAX_RAW_FRAMES = 117
@@ -209,7 +208,6 @@ def build_request(
"height": args.height,
"width": args.width,
"num_frames": args.num_frames,
"num_inference_steps": _NUM_INFERENCE_STEPS,
"max_sequence_length": _MAX_SEQUENCE_LENGTH,
"seed": args.seed,
"fps": args.fps,
@@ -155,7 +155,6 @@ async def run(argv: Sequence[str] | None = None) -> Path:
height=args.height,
width=args.width,
num_frames=9,
num_inference_steps=4,
max_sequence_length=512,
seed=args.seed,
output_type="latent",
@@ -5,10 +5,12 @@ This example shows how to use vLLM-Omni for running offline inference
with the correct prompt format on Qwen2.5-Omni
"""
import argparse
import functools
import json
import os
import time
from typing import NamedTuple
from typing import Any, NamedTuple
import numpy as np
import soundfile as sf
@@ -26,6 +28,20 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
SEED = 42
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"{flag_name} must be valid JSON: {e}") from e
if not isinstance(config, dict):
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")
class QueryResult(NamedTuple):
inputs: dict
limit_mm_per_prompt: dict[str, int]
@@ -378,11 +394,10 @@ def main(args):
for i, prompt in enumerate(prompts):
prompt["modalities"] = output_modalities
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
if profiler_enabled and hasattr(omni, "start_profile"):
profiler_enabled = args.profiler_config is not None
if profiler_enabled:
print("[Profiler] Starting profiling...")
omni.start_profile(stages=[0])
elif profiler_enabled:
print("[Warn] VLLM_TORCH_PROFILER_DIR is set, but current engine does not support profiler controls.")
omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
# Determine output directory: prefer --output-dir; fallback to --output-wav
@@ -418,7 +433,7 @@ def main(args):
print(f"Request ID: {request_id}, Saved audio to {output_wav}")
processed_count += 1
if profiler_enabled and hasattr(omni, "stop_profile") and processed_count >= total_requests:
if profiler_enabled and processed_count >= total_requests:
print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
# Stop the profiler while workers are still alive
omni.stop_profile()
@@ -476,6 +491,12 @@ def parse_args():
default=300,
help="Timeout for initializing stages in seconds (default: 300)",
)
parser.add_argument(
"--profiler-config",
type=parse_profiler_config,
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
parser.add_argument(
"--shm-threshold-bytes",
type=int,
@@ -14,6 +14,7 @@ list of supported architectures across all modalities, see
| Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
|---|---|---|---|---|---|---|
| Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | dual-AR | ✓ | ✓ | 11 languages | 44.1 kHz |
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
| Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz |
@@ -44,6 +45,57 @@ python examples/offline_inference/text_to_speech/<model>/end2end.py \
---
## Audio8 TTS Preview
0.6B DualAR text-to-speech model (Audio8) with its own 44.1 kHz neural audio
codec: a slow AR transformer predicts one semantic token per codec frame, and a
4-layer fast AR predicts that frame's remaining 9 codebooks.
### Prerequisites
No extra packages: the codec architecture ships in tree and its weights
(`codec.pth`) come with the checkpoint.
### Quick start
```bash
python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \
--text "Welcome to Audio8 TTS, a compact multilingual text to speech model."
```
### Voice cloning
```bash
python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \
--text "Welcome to Audio8 TTS." \
--ref-audio /path/to/reference.wav \
--ref-text "The exact transcript of the reference recording."
```
The reference transcript must match what is actually spoken in the clip;
a mismatch degrades both stability and speaker similarity.
### Streaming
```bash
python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \
--text "Welcome to Audio8 TTS." --streaming
```
### Notes
- Output: 44.1 kHz mono WAV; the codec emits ~21.5 frames/s (2048 samples per frame).
- Context is 2048 packed text+audio positions, so `max_model_len` is 2048.
- Greedy decoding (`temperature: 0`) is degenerate for this checkpoint — it
ends the utterance almost immediately. Keep the deploy defaults
(`temperature 0.7`, `top_k 50`, `top_p 0.9`), which also match
`generation_config.json`.
- Repetition-Aware Sampling is applied to the semantic token: a repeat inside
the last 10 frames is re-drawn from a flatter distribution instead of being
masked out.
- `--num-prompts 4` smoke-tests concurrency (stage 0 runs `max_num_seqs: 4`).
### Measured behaviour (1x H20, shared GPU)
`--streaming` on a 58-character English sentence: TTFA ~0.74 s (4-frame first
chunk), 7 chunks, 1.55 s wall for 3.16 s of audio (RTF ~0.49). Treat these as
smoke-test magnitudes, not a benchmark.
---
## CosyVoice3
2-stage TTS pipeline (`talker` + `code2wav`) at 24 kHz.
@@ -130,11 +182,7 @@ Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24
## Fish Speech S2 Pro
4B dual-AR text-to-speech model from FishAudio with the DAC codec at 44.1 kHz.
### Prerequisites
```bash
pip install fish-speech
```
No extra packages are required; the DAC codec is vendored in vLLM-Omni.
### Quick start
```bash
@@ -0,0 +1,237 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Offline inference demo for Audio8 TTS Preview 0.6B via vLLM Omni.
Usage:
# Text-only synthesis
python end2end.py --text "Welcome to Audio8 TTS."
# Zero-shot voice cloning (the transcript must match the reference audio)
python end2end.py --text "Welcome to Audio8 TTS." \
--ref-audio reference.wav --ref-text "The exact transcript of the reference recording."
# Streaming (per-chunk TTFA / inter-chunk timings)
python end2end.py --text "Welcome to Audio8 TTS." --streaming
"""
import asyncio
import os
import time
import numpy as np
import soundfile as sf
import torch
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
from vllm_omni import AsyncOmni, Omni
from vllm_omni.model_executor.models.audio8_tts.codec_utils import estimate_reference_code_frames
from vllm_omni.model_executor.models.audio8_tts.prompt_utils import (
build_text_only_prompt_ids,
estimate_voice_clone_prompt_len,
normalize_text,
)
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
DEFAULT_MODEL = "Audio8/Audio8-TTS-Preview-0.6b"
def build_prompt(
text: str,
ref_audio_path: str | None = None,
ref_text: str | None = None,
model_name: str = DEFAULT_MODEL,
) -> dict:
"""Build an engine prompt, using the same protocol as online serving."""
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
if ref_audio_path is None and ref_text is None:
prompt_ids, normalized_text = build_text_only_prompt_ids(tokenizer, text)
return {
"prompt_token_ids": prompt_ids,
"additional_information": {"text": [normalized_text]},
}
if not ref_audio_path or not ref_text:
raise ValueError("Audio8 TTS voice cloning requires both --ref-audio and --ref-text")
normalized_text = normalize_text(text)
normalized_ref_text = normalize_text(ref_text, add_default_speaker=True)
ref_audio_wav, ref_audio_sr = sf.read(ref_audio_path, dtype="float32", always_2d=False)
if ref_audio_wav.ndim > 1:
ref_audio_wav = ref_audio_wav.mean(axis=-1)
ref_frames = estimate_reference_code_frames(int(ref_audio_wav.shape[0]), int(ref_audio_sr))
placeholder_len = estimate_voice_clone_prompt_len(tokenizer, normalized_text, normalized_ref_text, ref_frames)
# The clone prompt embeds the reference audio's own codec codes, so the
# model builds it in preprocess(); we only reserve the exact length here.
return {
"prompt_token_ids": [1] * placeholder_len,
"additional_information": {
"text": normalized_text,
"ref_text": normalized_ref_text,
"ref_audio_wav": torch.from_numpy(np.asarray(ref_audio_wav, dtype=np.float32)),
"ref_audio_sr": int(ref_audio_sr),
"audio8_structured_voice_clone": True,
},
}
def extract_audio(multimodal_output: dict) -> tuple[torch.Tensor, int]:
"""Pull the waveform and sample rate out of a multimodal output.
The output processor concatenates the per-step delta tensors under
``model_outputs``; ``audio`` is the older alias. Never use ``a or b`` on
these values -- truthiness on a tensor raises.
"""
audio = multimodal_output.get("model_outputs")
if audio is None:
audio = multimodal_output.get("audio")
if audio is None:
raise ValueError(f"No audio in multimodal_output: {sorted(multimodal_output)}")
if isinstance(audio, list):
audio = torch.cat([torch.as_tensor(chunk).reshape(-1) for chunk in audio], dim=0)
audio = torch.as_tensor(audio).reshape(-1)
sr_raw = multimodal_output.get("sr")
if isinstance(sr_raw, list):
sr_raw = sr_raw[-1] if sr_raw else None
if sr_raw is None:
raise ValueError("Missing sample rate in multimodal_output")
sample_rate = int(sr_raw.item()) if hasattr(sr_raw, "item") else int(sr_raw)
return audio, sample_rate
def _save_wav(output_dir: str, request_id: str, audio: torch.Tensor, sample_rate: int) -> None:
out_wav = os.path.join(output_dir, f"output_{request_id}.wav")
sf.write(out_wav, audio.float().cpu().numpy(), samplerate=sample_rate, format="WAV")
duration = audio.numel() / sample_rate
# print(), not logging: vLLM reconfigures logging on engine start and
# disables loggers created before that, which would swallow these lines.
print(f"Request {request_id}: saved {out_wav} (sr={sample_rate}, {duration:.2f}s)")
def _omni_kwargs(args, model_name: str) -> dict:
kwargs = {
"model": model_name,
"log_stats": args.log_stats,
"stage_init_timeout": args.stage_init_timeout,
}
if args.deploy_config:
kwargs["deploy_config"] = args.deploy_config
return kwargs
def main(args) -> None:
os.makedirs(args.output_dir, exist_ok=True)
model_name = args.model or DEFAULT_MODEL
prompts = [
build_prompt(text, args.ref_audio, args.ref_text, model_name) for text in ([args.text] * args.num_prompts)
]
omni = Omni(**_omni_kwargs(args, model_name))
t_start = time.perf_counter()
total_audio_seconds = 0.0
for stage_outputs in omni.generate(prompts):
request_output = stage_outputs
if request_output is None or not request_output.outputs:
continue
audio, sample_rate = extract_audio(request_output.outputs[0].multimodal_output)
assert audio.numel() > 0, "Audio8 TTS produced an empty waveform"
total_audio_seconds += audio.numel() / sample_rate
_save_wav(args.output_dir, request_output.request_id, audio, sample_rate)
elapsed = time.perf_counter() - t_start
rtf = elapsed / max(total_audio_seconds, 1e-6)
print(f"Total {elapsed * 1000:.1f} ms for {total_audio_seconds:.2f} s of audio (RTF={rtf:.3f})")
async def main_streaming(args) -> None:
os.makedirs(args.output_dir, exist_ok=True)
model_name = args.model or DEFAULT_MODEL
prompt = build_prompt(args.text, args.ref_audio, args.ref_text, model_name)
omni = AsyncOmni(**_omni_kwargs(args, model_name))
request_id = "0"
chunks: list[torch.Tensor] = []
sample_rate = None
t_start = time.perf_counter()
t_prev = t_start
chunk_index = 0
async for stage_output in omni.generate(prompt, request_id=request_id):
multimodal_output = stage_output.outputs[0].multimodal_output
if not multimodal_output:
continue
audio = multimodal_output.get("model_outputs")
if audio is None:
audio = multimodal_output.get("audio")
if audio is not None:
if isinstance(audio, list):
chunks.extend(torch.as_tensor(chunk).reshape(-1) for chunk in audio)
else:
chunks.append(torch.as_tensor(audio).reshape(-1))
sr_raw = multimodal_output.get("sr")
if isinstance(sr_raw, list):
sr_raw = sr_raw[-1] if sr_raw else None
if sr_raw is not None:
sample_rate = int(sr_raw.item()) if hasattr(sr_raw, "item") else int(sr_raw)
if stage_output.finished:
break
now = time.perf_counter()
if chunk_index == 0:
print(f"chunk 0: TTFA={(now - t_start) * 1000:.1f} ms")
else:
print(f"chunk {chunk_index}: inter_chunk={(now - t_prev) * 1000:.1f} ms")
t_prev = now
chunk_index += 1
if not chunks or sample_rate is None:
raise RuntimeError("Audio8 TTS streaming produced no audio")
_save_wav(args.output_dir, request_id, torch.cat(chunks, dim=0), sample_rate)
print(f"Streaming done: total={(time.perf_counter() - t_start) * 1000:.1f} ms chunks={chunk_index}")
def parse_args():
parser = TrackingArgumentParser(description="Audio8 TTS Preview offline inference with vLLM Omni")
parser.add_argument("--model", type=str, default=DEFAULT_MODEL, help="Model path or HuggingFace repo ID.")
parser.add_argument(
"--text",
type=str,
default="Welcome to Audio8 TTS, a compact multilingual text to speech model.",
help="Text to synthesize.",
)
parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.")
parser.add_argument(
"--ref-text",
type=str,
default=None,
help="Transcript of --ref-audio; required for voice cloning and must match the recording.",
)
parser.add_argument("--output-dir", type=str, default="output_audio", help="Directory for output WAV files.")
parser.add_argument(
"--num-prompts",
type=int,
default=1,
help="Duplicate the prompt N times to smoke-test concurrent requests.",
)
parser.add_argument(
"--deploy-config",
type=str,
default=None,
help="Override the deploy config path; defaults to vllm_omni/deploy/audio8_tts.yaml via model_type.",
)
parser.add_argument("--streaming", action="store_true", help="Use AsyncOmni and report per-chunk latency.")
parser.add_argument("--log-stats", action="store_true", help="Enable engine stats logging.")
parser.add_argument("--stage-init-timeout", type=int, default=600, help="Per-stage init timeout in seconds.")
return parser.parse_args()
if __name__ == "__main__":
parsed_args = parse_args()
if parsed_args.streaming:
asyncio.run(main_streaming(parsed_args))
else:
main(parsed_args)
@@ -1,9 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import functools
import json
import os
import urllib.request
from pathlib import Path
from typing import Any
import numpy as np
import soundfile as sf
@@ -29,6 +32,20 @@ def _default_ref_audio() -> str:
return str(dest)
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"{flag_name} must be valid JSON: {e}") from e
if not isinstance(config, dict):
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 run_e2e():
parser = argparse.ArgumentParser()
# ""FunAudioLLM/Fun-CosyVoice3-0.5B-2512
@@ -64,6 +81,12 @@ def run_e2e():
required=True,
help="Path to tokenizer directory (e.g., <model_path>/CosyVoice-BlankEN).",
)
parser.add_argument(
"--profiler-config",
type=parse_profiler_config,
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
args = parser.parse_args()
# Ensure tokenizer directory exists
if not os.path.exists(args.tokenizer):
@@ -79,6 +102,7 @@ def run_e2e():
deploy_config=args.deploy_config,
tokenizer=args.tokenizer,
log_stats=True,
profiler_config=args.profiler_config,
)
sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
@@ -149,8 +173,8 @@ def run_e2e():
sampling_params_list = [gpt_sampling, s2mel_sampling]
# Start profiling (requires VLLM_TORCH_PROFILER_DIR env var)
if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
profiler_enabled = args.profiler_config is not None
if profiler_enabled:
print("Starting profiler...")
omni.start_profile()
@@ -158,7 +182,7 @@ def run_e2e():
outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list[:2]))
# Stop profiling and get results
if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
if profiler_enabled:
print("Stopping profiler...")
profile_results = omni.stop_profile()
print(f"Profile traces saved to: {profile_results}")
+21 -3
View File
@@ -264,7 +264,7 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene
#### Realtime WebSocket client (`openai_realtime_client.py`)
[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://<host>:<port>/v1/realtime`**, streams a local WAV as **PCM16 mono @ 16 kHz** in fixed-size chunks (OpenAI-style `input_audio_buffer.append` / `commit`), and receives **`response.audio.delta`** (incremental PCM for the reply) plus **`transcription.*`** events. By default it concatenates audio deltas and writes **`--output-wav`** (model output is typically **24 kHz**). Optional **`--delta-dump-dir`** saves each delta as `delta_000001.wav`, … for debugging.
[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://<host>:<port>/v1/realtime`**, streams a local WAV as **PCM16 mono @ 16 kHz** in fixed-size chunks, and receives **`response.audio.*`**, **`response.output_text.*`**, and transcript events. By default it uses explicit `input_audio_buffer.commit`; pass **`--server-vad`** to let the server detect speech endpoints and commit the turn automatically. The client concatenates audio deltas and writes **`--output-wav`** (model output is typically **24 kHz**). Optional **`--delta-dump-dir`** saves each delta as `delta_000001.wav`, … for debugging.
Streaming input works well for translation-style use cases; if the Thinker runs while input is still incomplete, consider limiting **`max_tokens`** in your session / server defaults to avoid over-generation.
@@ -282,6 +282,7 @@ python openai_realtime_client.py \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct \
--input-wav /path/to/input_16k_mono.wav \
--output-wav realtime_output.wav \
--server-vad \
--delta-dump-dir ./rt_delta_wavs
```
@@ -297,15 +298,32 @@ python openai_realtime_client.py \
| `--chunk-ms` | `200` | Size of each uploaded audio chunk (milliseconds of audio) |
| `--send-delay-ms` | `0` | Delay between chunk sends (simulate realtime upload) |
| `--delta-dump-dir` | *(optional)* | Directory to write per-`response.audio.delta` WAV files |
| `--server-vad` | disabled | Enable server-side Silero VAD instead of sending explicit commits |
| `--num-requests` | `1` | Number of sequential sessions (see `--concurrency`) |
| `--concurrency` | `1` | Max concurrent WebSocket sessions when `--num-requests` > 1 |
Ensure the server is running, for example:
Server VAD uses the pinned Silero v6.2 ONNX artifact from `istupakov/silero-vad-onnx` and never downloads model files while processing a session. Make the artifact available in the Hugging Face cache before startup, or configure a local artifact in the deployment YAML:
```yaml
duplex_session:
server_vad_model_path: /models/silero_vad.onnx
```
Start from the bundled Qwen3-Omni deployment YAML, add the fields above at the top level, and serve that complete configuration:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
cp vllm_omni/deploy/qwen3_omni_moe.yaml /path/to/qwen3_omni_server_vad.yaml
# Edit /path/to/qwen3_omni_server_vad.yaml and add duplex_session.
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
--omni \
--port 8091 \
--deploy-config /path/to/qwen3_omni_server_vad.yaml
```
Keep Qwen's default `session_mode: turn`; `duplex_session` enables the Realtime handler without changing scheduler
semantics. Bare `/v1/realtime` selects that handler. The client uses `?duplex=0` only for the existing non-Server-VAD
wire flow; `?duplex=1` remains a supported compatibility alias.
The Python client supports the following command-line arguments:
- `--query-type` (or `-q`): Query type (default: `use_video`). Options: `text`, `use_audio`, `use_image`, `use_video`
@@ -1,9 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Realtime client for vLLM-Omni /v1/realtime (audio + text events).
This client:
1) Reads a local WAV file (must be mono, 16-bit PCM, 16kHz),
2) Streams PCM16 chunks to /v1/realtime with OpenAI-style events,
3) Receives response.audio.* and transcription.* events,
3) Receives response.audio.*, response.output_text.*, and transcript events,
4) Saves synthesized audio to an output WAV file and optional text file.
By default each ``response.audio.delta`` is treated as an **incremental PCM**
@@ -18,6 +21,7 @@ Usage:
--model Qwen/Qwen3-Omni-30B-A3B-Instruct \
--input-wav input_16k_mono.wav \
--output-wav realtime_output.wav \
--server-vad \
--delta-dump-dir ./rt_delta_wavs
Dependencies:
@@ -32,6 +36,7 @@ import base64
import json
import wave
from pathlib import Path
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
try:
import websockets
@@ -70,6 +75,14 @@ def _write_wav_pcm16(path: Path, pcm16_bytes: bytes, sample_rate_hz: int) -> Non
wf.writeframes(pcm16_bytes)
def _with_realtime_route(url: str, *, server_vad: bool) -> str:
parts = urlsplit(url)
query = [(key, value) for key, value in parse_qsl(parts.query, keep_blank_values=True) if key != "duplex"]
if not server_vad:
query.append(("duplex", "0"))
return urlunsplit(parts._replace(query=urlencode(query)))
async def run_client(
url: str,
model: str,
@@ -79,6 +92,7 @@ async def run_client(
chunk_ms: int,
send_delay_ms: int,
delta_dump_dir: Path | None,
server_vad: bool,
request_idx: int = 1,
total_requests: int = 1,
) -> None:
@@ -92,23 +106,44 @@ async def run_client(
delta_index = 0
text_chunks: list[str] = []
final_text: str = ""
final_audio_transcript: str = ""
if delta_dump_dir is not None:
delta_dump_dir.mkdir(parents=True, exist_ok=True)
# Keep the CLI mode authoritative even if the input URL already contains
# an explicit duplex query parameter.
url = _with_realtime_route(url, server_vad=server_vad)
async with websockets.connect(url, max_size=64 * 1024 * 1024) as ws:
# 1) Validate model.
await ws.send(
json.dumps(
{
"type": "session.update",
if server_vad:
session_update = {
"type": "session.update",
"session": {
"model": model,
}
)
)
"audio": {
"input": {
"format": {"type": "audio/pcm", "rate": 16_000},
"turn_detection": {
"type": "server_vad",
"threshold": 0.5,
"prefix_padding_ms": 300,
"silence_duration_ms": 500,
"create_response": True,
"interrupt_response": False,
},
}
},
},
}
else:
# Preserve the existing non-duplex wire format.
session_update = {"type": "session.update", "model": model}
await ws.send(json.dumps(session_update))
# 2) Start generation once (non-final commit).
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False}))
# 2) Preserve the legacy explicit-commit flow unless Server VAD owns endpointing.
if not server_vad:
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False}))
# 3) Stream audio chunks.
for i in range(0, len(pcm16), chunk_bytes):
@@ -124,8 +159,20 @@ async def run_client(
if send_delay_ms > 0:
await asyncio.sleep(send_delay_ms / 1000.0)
# 4) Final commit closes input stream.
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
# 4) Finalize the turn explicitly, or append silence for Server VAD endpointing.
if server_vad:
trailing_silence = bytes(16_000 * 2)
for i in range(0, len(trailing_silence), chunk_bytes):
await ws.send(
json.dumps(
{
"type": "input_audio_buffer.append",
"audio": base64.b64encode(trailing_silence[i : i + chunk_bytes]).decode("utf-8"),
}
)
)
else:
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
# 5) Receive server events until audio done.
while True:
@@ -144,7 +191,7 @@ async def run_client(
sr = event.get("sample_rate_hz")
if isinstance(sr, int) and sr > 0:
output_sample_rate = sr
audio_b64 = event.get("audio", "")
audio_b64 = event.get("delta" if server_vad else "audio", "")
if audio_b64:
pcm_delta = base64.b64decode(audio_b64)
incremental_pcm_parts.append(pcm_delta)
@@ -158,14 +205,14 @@ async def run_client(
)
continue
if event_type == "transcription.delta":
if event_type in {"response.output_text.delta", "transcription.delta"}:
delta = event.get("delta", "")
if delta:
text_chunks.append(delta)
print(delta, end="", flush=True)
continue
if event_type == "transcription.done":
if event_type in {"response.output_text.done", "transcription.done"}:
final_text = event.get("text", "") or "".join(text_chunks)
usage = event.get("usage")
final_text_with_tag = f"Final transcription: {final_text}"
@@ -176,7 +223,16 @@ async def run_client(
print(f"{log_prefix}text usage: {usage}")
continue
if event_type == "response.audio_transcript.done":
final_audio_transcript = event.get("transcript", "")
continue
if event_type == "response.audio.done":
if not server_vad:
break
continue
if event_type == "response.done" and server_vad:
break
if event_type == "error":
@@ -191,7 +247,7 @@ async def run_client(
print(f"{log_prefix}Saved realtime audio to: {output_wav} (incremental chunks joined)")
if output_text is not None:
text_to_save = final_text if final_text else "".join(text_chunks)
text_to_save = final_text or "".join(text_chunks) or final_audio_transcript
output_text.parent.mkdir(parents=True, exist_ok=True)
output_text.write_text(text_to_save, encoding="utf-8")
print(f"{log_prefix}Saved realtime text to: {output_text}")
@@ -213,6 +269,7 @@ async def run_clients_concurrent(
chunk_ms: int,
send_delay_ms: int,
delta_dump_dir: Path | None,
server_vad: bool,
num_requests: int,
concurrency: int,
) -> None:
@@ -235,6 +292,7 @@ async def run_clients_concurrent(
chunk_ms=chunk_ms,
send_delay_ms=send_delay_ms,
delta_dump_dir=per_delta_dir,
server_vad=server_vad,
request_idx=index,
total_requests=num_requests,
)
@@ -283,6 +341,11 @@ def main() -> None:
default=None,
help="If set, each response.audio.delta is saved as delta_NNNNNN.wav under this directory",
)
parser.add_argument(
"--server-vad",
action="store_true",
help="Let the server detect speech endpoints instead of sending input_audio_buffer.commit",
)
parser.add_argument("--num-requests", type=int, default=1, help="Total number of requests to send")
parser.add_argument(
"--concurrency",
@@ -309,6 +372,7 @@ def main() -> None:
chunk_ms=args.chunk_ms,
send_delay_ms=args.send_delay_ms,
delta_dump_dir=args.delta_dump_dir,
server_vad=args.server_vad,
)
)
else:
@@ -322,6 +386,7 @@ def main() -> None:
chunk_ms=args.chunk_ms,
send_delay_ms=args.send_delay_ms,
delta_dump_dir=args.delta_dump_dir,
server_vad=args.server_vad,
num_requests=args.num_requests,
concurrency=concurrency,
)
@@ -1,7 +1,7 @@
# 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.
It covers text-to-video, and image-conditioned models through `--image-reference`.
## Start The Server
@@ -20,7 +20,7 @@ The `--diffusion-streaming-output` CLI flag is forwarded as `streaming_output=Tr
| 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. |
| 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. Image-conditioned models take a first frame as `image_reference`: `{"image_url":"https://..."}` or a `data:` URL. |
| Client to server | `session.interaction` | JSON text: `{"type":"session.interaction","interaction":{"event_id":"xxx","event":{"prompt":"..."},"transition_chunks":3}}` | Updates the active prompt midway through generation. `event_id` is optional and `transition_chunks` defaults to the model setting. |
| 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_metadata` | JSON text: `{"type":"video.chunk_metadata","request_id":"...","kind":"media","transport_chunk_index":0,"generation_chunk_index":0,"num_frames":9,"byte_length":1234,"started_event_ids":[],"active_event_ids":[],"completed_event_ids":[]}` | Precedes each binary frame and describes the immediately following payload. |
@@ -62,7 +62,7 @@ python streaming_video_client.py \
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.
Schedule midway prompt updates with `--prompt-updates`. Each entry uses `"at"` as seconds on the client clock after the server sends `video.start`:
Schedule midway prompt updates with `--prompt-updates`. None are sent unless you ask for them, and only pipelines that implement midway prompt updates accept them (LingBot-World, for one, does not, and rejects the whole session). Each entry uses `"at"` as seconds on the client clock after the server sends `video.start`:
```bash
python streaming_video_client.py \
@@ -73,6 +73,19 @@ python streaming_video_client.py \
]'
```
The Helios demo continues the default prompt into an underwater storm:
```bash
python streaming_video_client.py \
--model BestWishYsh/Helios-Distilled \
--prompt-updates '[
{"at": 4,
"prompt": "An underwater tornado appears and affects the ocean floor in a dramatic and chaotic scene. The water is murky, swirling violently, carrying debris and marine life into the vortex. The tropical fish on the scene all swim in panic, trying to avoid the powerful currents. The camera remains stationary, capturing the intensity of the underwater tornado as it disrupts the serene ocean floor. Close-up shot emphasizing the turbulent motion and destruction."},
{"at": 11,
"prompt": "The swirling underwater vortex now seizes a heavy, encrusted treasure chest, its lid flapping open as it is smashed onto the ocean floor. Gold coins and silver trinkets spill out, glittering briefly in the murky water before being swept instantly into the violent funnel. The heavy wooden box tumbles end over end, colliding with floating rocks and adding to the debris field. Swirling sediment and bubbles surround the spilling fortune, highlighting the chaotic power of the storm as it ravages the seabed. Close-up shot emphasizing the turbulent motion and destruction."}
]'
```
## Run The Gradio Demo
```bash
@@ -106,3 +119,31 @@ Disable it in the CLI example with `--no-helios-distilled-preset`, or override/e
python streaming_video_client.py \
--extra-params '{"pyramid_num_inference_steps_list":[2, 2, 2]}'
```
### LingBot-World 2.0
`robbyant/lingbot-world-v2-14b-causal-fast-diffusers` is image-conditioned and runs on the
AR-Diffusion engine, so it needs the stepwise deploy config and a first frame. See
[the recipe](../../../recipes/Robbyant/LingBot-World-2.0.md) for the constraints on
`width`/`height` and `num_frames`.
```bash
vllm serve robbyant/lingbot-world-v2-14b-causal-fast-diffusers \
--omni \
--deploy-config vllm_omni/deploy/lingbot_world_v2_stepwise.yaml \
--port 8000
```
```bash
python streaming_video_client.py \
--model robbyant/lingbot-world-v2-14b-causal-fast-diffusers \
--prompt "The camera moves slowly forward through the scene." \
--image-reference /path/to/first_frame.png \
--width 832 --height 480 --num-frames 33 --fps 16 --seed 42 \
--extra-params '{"camera_action_script":[[["w"],["w"],["w"]],[["a"],[],[]],[[],[],[]]]}' \
--output lingbot_world_v2_stream.mp4
```
`camera_action_script` carries one three-latent-frame WASD action list per generated
chunk. Mid-session `session.interaction` only updates the prompt, so a LingBot rollout
follows the camera script it started with.
@@ -1,11 +1,12 @@
#!/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`.
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""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.
``--image-reference`` sends a first frame for image-conditioned models such as
LingBot-World 2.0; text-only models need no image.
Requirements:
pip install av websockets
@@ -13,9 +14,11 @@ Requirements:
import argparse
import asyncio
import base64
import contextlib
import io
import json
import mimetypes
import os
import time
from dataclasses import dataclass
@@ -53,6 +56,18 @@ class ScheduledPromptUpdate:
transition_chunks: int
def _image_reference(value: str) -> dict[str, Any]:
"""Accept an http(s) URL as-is, or inline a local image as a data URL."""
if value.startswith(("http://", "https://", "data:")):
return {"image_url": value}
path = Path(value)
if not path.is_file():
raise argparse.ArgumentTypeError(f"--image-reference is neither a URL nor a readable file: {value}")
media_type = mimetypes.guess_type(path.name)[0] or "image/png"
encoded = base64.b64encode(path.read_bytes()).decode()
return {"image_url": f"data:{media_type};base64,{encoded}"}
def _parse_extra_params(value: str) -> dict[str, Any]:
try:
parsed = json.loads(value)
@@ -116,6 +131,7 @@ def build_session_start(args: argparse.Namespace) -> dict[str, Any]:
"prompt": args.prompt,
}
_maybe_set(payload, "image_reference", args.image_reference)
_maybe_set(payload, "negative_prompt", args.negative_prompt)
_maybe_set(payload, "width", args.width)
_maybe_set(payload, "height", args.height)
@@ -362,6 +378,12 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--fps", type=int, default=16, help="Frames per second")
parser.add_argument("--num-frames", type=int, default=429, help="Number of generated frames")
parser.add_argument(
"--image-reference",
type=_image_reference,
default=None,
help="First frame for image-conditioned models: a local image path or an http(s) URL.",
)
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")
@@ -383,36 +405,16 @@ def parse_args() -> argparse.Namespace:
default=None,
help="JSON object merged into extra_params; overrides preset keys on conflict.",
)
# Midway prompt updates are pipeline-specific: a model that does not implement
# them (LingBot-World, for one) rejects the whole session on the first one, so
# nothing is scheduled unless the caller asks for it.
parser.add_argument(
"--prompt-updates",
type=_parse_prompt_updates,
default=json.dumps(
[
{
"at": 4,
"prompt": (
"An underwater tornado appears and affects the ocean floor in a dramatic and chaotic scene. "
"The water is murky, swirling violently, carrying debris and marine life into the vortex. "
"The tropical fish on the scene all swim in panic, trying to avoid the powerful currents. "
"The camera remains stationary, capturing the intensity of the underwater tornado as it disrupts the serene ocean floor. "
"Close-up shot emphasizing the turbulent motion and destruction."
),
},
{
"at": 11,
"prompt": (
"The swirling underwater vortex now seizes a heavy, encrusted treasure chest, its lid flapping open as it is smashed onto the ocean floor. "
"Gold coins and silver trinkets spill out, glittering briefly in the murky water before being swept instantly into the violent funnel. "
"The heavy wooden box tumbles end over end, colliding with floating rocks and adding to the debris field. "
"Swirling sediment and bubbles surround the spilling fortune, highlighting the chaotic power of the storm as it ravages the seabed. "
"Close-up shot emphasizing the turbulent motion and destruction."
),
# "transition_chunks": 3,
},
]
),
default=[],
help=(
"JSON array of scheduled prompt updates. Each object requires "
"JSON array of scheduled prompt updates; none are sent by default. "
"Each object requires "
"'at' (seconds after video.start) and 'prompt'. Optional "
f"'transition_chunks' defaults to {DEFAULT_TRANSITION_CHUNKS}."
),
@@ -14,6 +14,7 @@ For the full list of supported architectures across all modalities, see
| Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
| --- | --- | --- | --- | --- | --- |
| Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | uploaded audio voice only; no presets | ✓ |
| 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`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — |
@@ -253,16 +254,152 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \
---
## Fish Speech S2 Pro
## Audio8 TTS Preview
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
0.6B DualAR TTS at 44.1 kHz, 11 languages, zero-shot voice cloning.
### Prerequisites
None beyond the base install: the neural audio codec is implemented in tree and
its weights (`codec.pth`) ship with the checkpoint.
### Launch
```bash
pip install fish-speech
vllm serve Audio8/Audio8-TTS-Preview-0.6b --omni --port 8092
# or:
./audio8_tts/run_server.sh
```
The deploy config auto-loads from `vllm_omni/deploy/audio8_tts.yaml` (HF
`model_type` is `arktts`). Do **not** pass `--trust-remote-code`: vllm-omni
registers its own `arktts` config, and transformers would otherwise prefer the
checkpoint's remote code and bypass it.
### Text-only synthesis
```bash
curl -X POST http://localhost:8092/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Welcome to Audio8 TTS.",
"response_format": "wav"
}' --output output.wav
```
### Voice cloning
```bash
curl -X POST http://localhost:8092/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Welcome to Audio8 TTS.",
"ref_audio": "https://example.com/reference.wav",
"ref_text": "The exact transcript of the reference recording."
}' --output cloned.wav
```
`ref_audio` also accepts a `data:audio/wav;base64,...` URL. `ref_text` is
mandatory whenever `ref_audio` is present and must match the recording.
Uploading a voice via `POST /v1/audio/voices` lets the server reuse the encoded
reference codes across requests (`voice: "<name>"`).
### Streaming
```bash
python audio8_tts/speech_client.py --text "Welcome to Audio8 TTS." --stream --output out.pcm
ffplay -f s16le -ar 44100 -ac 1 out.pcm
```
### Gradio demo
```bash
./audio8_tts/run_gradio_demo.sh # server + demo
python audio8_tts/gradio_demo.py --api-base http://localhost:8092 # demo only
```
### Benchmark
```bash
export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100 # required: harness defaults to 24 kHz
vllm bench serve --omni --host 127.0.0.1 --port 8092 \
--model Audio8/Audio8-TTS-Preview-0.6b \
--backend openai-audio-speech --endpoint /v1/audio/speech \
--dataset-name seed-tts-text --dataset-path benchmarks/build_dataset/seed_tts_smoke \
--seed-tts-locale en --num-prompts 20 --num-warmups 2 \
--max-concurrency 1 --request-rate inf \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun
```
Measured on 1x H20 **shared with an unrelated job holding ~80% SM**, so treat
these as a lower bound. 20 prompts from `seed_tts_smoke/en`, mean audio 4.5 s,
deploy defaults:
| | HF reference (batch=1) | vllm-omni c=1 | c=4 | c=8 |
| --- | --- | --- | --- | --- |
| RTF (mean) | 1.008 | **0.19** | 0.30 | 0.54 |
| Throughput (req/s) | 0.23 | 1.18 | 2.84 | **2.93** |
| E2E latency (mean, ms) | 4375 | 850 | 1349 | 2361 |
| Time to first packet (mean, ms) | n/a (no streaming) | 63 | 105 | 1124 |
| Underrun p99 / continuity OK | n/a | 0 s / 100% | 0 s / 100% | 0 s / 100% |
~5.3x lower per-request RTF than the upstream reference implementation at
concurrency 1, ~12.8x aggregate throughput at concurrency 8.
#### Stage-0 `max_num_seqs`: first-packet vs throughput
This subsection is operator tuning guidance, not a claim this PR optimizes
anything: the shipped deploy default is unchanged (`max_num_seqs: 4`). The
numbers below are a single run on the same contended H20 as above (one unrelated
job holding ~80% SM, no repeats), so read them as directional, not as measured
speedups. Same setup, varying stage 0's `max_num_seqs` (stage 1 stays at 1):
| stage-0 `max_num_seqs` | client concurrency | req/s | RTF | TTFP mean / p99 (ms) | underrun p99 (s) |
| --- | --- | --- | --- | --- | --- |
| 4 (default) | 1 | 1.18 | 0.19 | 63 / 67 | 0.00 |
| 4 | 4 | 2.84 | 0.30 | 105 / 148 | 0.00 |
| 4 | 8 | 2.93 | 0.54 | 1124 / 1642 | 0.00 |
| 8 | 1 | 1.16 | 0.19 | 64 / 70 | 0.00 |
| 8 | 4 | 2.88 | 0.30 | 107 / 148 | 0.00 |
| 8 | 8 | **3.72** | 0.40 | **247** / 733 | 0.22 |
| 8 | 16 | 3.91 | 0.65 | 1366 / 2530 | 0.29 |
Reading of this:
- Below the limit (`concurrency <= max_num_seqs`) the setting is irrelevant, as
expected — c=1 and c=4 are identical for both.
- At c=8, raising the limit to 8 admits every request instead of queueing half
of them: in this single run, roughly +27% throughput and ~4.5x lower
first-packet latency.
- Throughput saturates around 3.7-3.9 req/s (~17 s of audio per second). Past
c=8 you only buy queueing delay.
- The cost of admitting more streams is a 0.22-0.29 s worst-case buffer deficit
in 1 of 20 requests (`Streaming continuity OK rate` 95%), because stage 1 runs
`max_num_seqs: 1` and the codec window round-robins across streams. Client-side
prebuffering (the bundled Gradio demo holds 2 chunks) absorbs that; raising
stage 1's `max_num_seqs` is the untested knob if you need both.
The shipped default stays at 4 because every measured point there had zero
underrun. Raise it to 8 if first-packet latency under load matters more than a
rare sub-300 ms gap.
### Notes
- Output: 44.1 kHz mono; ~21.5 codec frames per second.
- No built-in speaker presets. Omit `voice` for a random timbre, or clone one.
- `max_new_tokens` caps generated codec frames (1 frame ~= 46 ms).
- Context is 2048 packed text+audio positions (`max_model_len: 2048`).
- Do not send `temperature: 0`: greedy decoding is degenerate for this
checkpoint and truncates the utterance. The deploy defaults
(0.7 / top_k 50 / top_p 0.9) mirror `generation_config.json`.
---
## Fish Speech S2 Pro
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec, which is vendored in
vLLM-Omni (no extra packages required).
### Kvcache attention fast path
Fish Speech S2 Pro uses a Triton decode-only kvcache attention fast path by
@@ -425,6 +562,46 @@ python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.
---
## MOSS-TTS Local Transformer v1.5
For a single H200, the optional
[`moss_tts_local_h200.yaml`](../../../vllm_omni/deploy/moss_tts_local_h200.yaml)
deployment places both the talker and codec on logical GPU 0:
```bash
CUDA_VISIBLE_DEVICES=0 VLLM_OMNI_EVENT_DRIVEN_ORCH=1 \
vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \
--omni --trust-remote-code \
--deploy-config vllm_omni/deploy/moss_tts_local_h200.yaml \
--disable-log-stats
```
Run from the repository root and select an available physical GPU with
`CUDA_VISIBLE_DEVICES`. The command enables event-driven orchestration and
disables detailed per-request statistics logging to reduce CPU overhead at
high concurrency. Remove `--disable-log-stats` when those statistics are needed;
keep these settings identical when comparing performance.
This preset configures 256 request slots per stage,
a 32 GiB talker KV cache, talker CUDA Graph buckets through 512 scheduled tokens,
and codec batch buckets through 256. It requires more than 80 GiB of GPU memory;
request capacity also depends on input and generated lengths. The default
[`moss_tts_local.yaml`](../../../vllm_omni/deploy/moss_tts_local.yaml) remains
available for smaller deployments.
The preset enables the optional codec backend with
`hf_overrides.codec_attention_backend: triton` on stage 1. It preserves the
streaming ring-cache mask and uses BF16 attention with 64-dimensional heads;
other attention shapes use PyTorch SDPA. Set the backend to `sdpa` to use the
default attention implementation. Codec terminal tails share execution only
when they already map to the same padded CUDA Graph; returned audio retains
each request's actual length.
Voice cloning requests use `ref_audio` and `ref_text`. For streaming output,
set `stream: true`, `stream_format: "audio"`, and `response_format: "pcm"`;
the native PCM format is 48 kHz, stereo, signed 16-bit little-endian. Compute
audio throughput as `PCM bytes / (48000 * 2 * 2) / elapsed seconds`, and compare
the same requests, concurrency, and reference-cache state.
## MOSS-TTS-Nano
Single-stage 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec at 48 kHz mono. Every request must include `ref_audio`; there are no built-in speaker presets.
@@ -668,13 +845,21 @@ Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and
### Streaming WebSocket
The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
The `/v1/audio/speech/stream` endpoint accepts text incrementally and, by default, synthesizes the buffered text as one continuous request on `input.done`:
```bash
python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
```
For per-sentence audio (including Indic danda `।`) before `input.done`:
```bash
python qwen3_tts/streaming_speech_client.py \
--text "नमस्ते। कैसे हो?" \
--split-granularity sentence
```
`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection:
```bash
@@ -0,0 +1,258 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Gradio demo for Audio8 TTS Preview via the /v1/audio/speech API.
Supports:
- text-to-speech synthesis in 11 languages
- zero-shot voice cloning from reference audio (upload, microphone, or URL)
- streaming (progressive PCM chunks) and non-streaming modes
Usage:
# Start the server first (see run_server.sh), then:
python gradio_demo.py --api-base http://localhost:8092
"""
import argparse
import base64
import io
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
import httpx
import numpy as np
import soundfile as sf
#: Audio8 TTS decodes at 44.1 kHz; raw PCM must be interpreted at that rate.
PCM_SAMPLE_RATE = 44100
DEFAULT_API_BASE = "http://localhost:8092"
def encode_audio_to_base64(audio_data: tuple) -> str:
"""Encode a Gradio ``(sample_rate, numpy_array)`` input as a data URL."""
sample_rate, audio_np = audio_data
if audio_np.dtype != np.int16:
if audio_np.dtype in (np.float32, np.float64):
audio_np = (np.clip(audio_np, -1.0, 1.0) * 32767).astype(np.int16)
else:
audio_np = audio_np.astype(np.int16)
buf = io.BytesIO()
sf.write(buf, audio_np, sample_rate, format="WAV")
return f"data:audio/wav;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}"
def build_payload(
text: str,
ref_audio: tuple | None,
ref_audio_url: str,
ref_text: str,
response_format: str = "wav",
stream: bool = False,
) -> dict:
"""Build the ``/v1/audio/speech`` payload.
Audio8 TTS has no built-in speakers: omit ``voice`` for a random timbre, or
supply ``ref_audio`` + ``ref_text`` to clone a voice.
"""
if not text or not text.strip():
raise gr.Error("Please enter text to synthesize.")
payload: dict = {
"input": text.strip(),
"response_format": "pcm" if stream else response_format,
"stream": stream,
}
if stream:
payload["stream_format"] = "audio"
ref_url = ref_audio_url.strip() if ref_audio_url else ""
if ref_audio is not None:
payload["ref_audio"] = encode_audio_to_base64(ref_audio)
elif ref_url:
payload["ref_audio"] = ref_url
if "ref_audio" in payload:
if not ref_text or not ref_text.strip():
raise gr.Error("Voice cloning requires the transcript of the reference audio.")
payload["ref_text"] = ref_text.strip()
return payload
def generate_speech(api_base: str, text: str, ref_audio, ref_audio_url, ref_text, response_format):
"""Non-streaming request; returns the whole waveform."""
payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=False)
try:
with httpx.Client(timeout=300.0) as client:
resp = client.post(
f"{api_base}/v1/audio/speech",
json=payload,
headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
)
except httpx.TimeoutException:
raise gr.Error("Request timed out. The server may be busy.") from None
except httpx.ConnectError:
raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") from None
if resp.status_code != 200:
raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
if "application/json" in resp.headers.get("content-type", ""):
raise gr.Error(f"Server error: {resp.text}")
try:
if response_format == "pcm":
audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0
return (PCM_SAMPLE_RATE, audio_np)
audio_np, sample_rate = sf.read(io.BytesIO(resp.content))
if audio_np.ndim > 1:
audio_np = audio_np[:, 0]
return (sample_rate, audio_np.astype(np.float32))
except Exception as e:
raise gr.Error(f"Failed to decode audio: {e}") from None
def generate_speech_stream(api_base, text, ref_audio, ref_audio_url, ref_text, response_format, chunk_seconds=0.5):
"""Streaming request; yields PCM chunks for progressive playback."""
payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=True)
min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds)
prebuffer_chunks = 2
pending: list[np.ndarray] = []
pending_len = 0
chunks_yielded = 0
prebuffer: list[np.ndarray] = []
leftover = b""
try:
with (
httpx.Client(timeout=300.0) as client,
client.stream(
"POST",
f"{api_base}/v1/audio/speech",
json=payload,
headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
) as resp,
):
if resp.status_code != 200:
resp.read()
raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
for chunk in resp.iter_bytes():
if not chunk:
continue
raw = leftover + chunk
# s16le: never split a sample across chunk boundaries.
usable = len(raw) - (len(raw) % 2)
leftover = raw[usable:]
if usable == 0:
continue
samples = np.frombuffer(raw[:usable], dtype=np.int16).copy().astype(np.float32) / 32767.0
pending.append(samples)
pending_len += len(samples)
if pending_len < min_chunk_samples:
continue
audio_chunk = np.concatenate(pending)
pending.clear()
pending_len = 0
if chunks_yielded < prebuffer_chunks:
prebuffer.append(audio_chunk)
chunks_yielded += 1
if chunks_yielded == prebuffer_chunks:
yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer))
prebuffer.clear()
else:
yield (PCM_SAMPLE_RATE, audio_chunk)
remaining = prebuffer + pending
if remaining:
yield (PCM_SAMPLE_RATE, np.concatenate(remaining))
except httpx.TimeoutException:
raise gr.Error("Request timed out. The server may be busy.") from None
except httpx.ConnectError:
raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") from None
def on_stream_change(stream: bool):
"""Streaming only makes sense with raw PCM."""
if stream:
return gr.update(value="pcm", interactive=False)
return gr.update(value="wav", interactive=True)
def build_interface(api_base: str, stream_chunk_seconds: float = 0.5):
with gr.Blocks(title="Audio8 TTS Preview Demo") as demo:
gr.Markdown("# Audio8 TTS Preview 0.6B - Text to Speech")
gr.Markdown(f"**Server:** `{api_base}` | **Model:** Audio8/Audio8-TTS-Preview-0.6b | **Output:** 44.1 kHz")
with gr.Row():
with gr.Column(scale=3):
text_input = gr.Textbox(label="Text to Synthesize", placeholder="Enter text here...", lines=4)
with gr.Accordion("Voice Cloning (optional)", open=False):
gr.Markdown(
"Upload or link a short reference recording (5-30 s) and paste its **exact** transcript. "
"A transcript that does not match the audio degrades both stability and speaker similarity."
)
ref_audio = gr.Audio(label="Reference Audio", type="numpy", sources=["upload", "microphone"])
ref_audio_url = gr.Textbox(
label="Reference Audio URL (alternative to upload)",
placeholder="https://example.com/reference.wav",
lines=1,
)
ref_text = gr.Textbox(
label="Reference Audio Transcript (required for cloning)",
placeholder="Exact transcript of the reference audio...",
lines=2,
)
with gr.Row():
response_format = gr.Dropdown(
choices=["wav", "mp3", "flac", "pcm"], value="wav", label="Audio Format", scale=1
)
stream_checkbox = gr.Checkbox(
label="Stream output", value=False, info="Progressive PCM streaming", scale=1
)
generate_btn = gr.Button("Generate Speech", variant="primary", size="lg")
with gr.Column(scale=2):
audio_output = gr.Audio(label="Generated Audio", interactive=False, streaming=True, autoplay=True)
gr.Markdown(
"### About\n"
"- **Audio8 TTS Preview** — 0.6B DualAR model (slow AR + fast AR codebook decoder)\n"
"- **Languages**: Cantonese, Chinese, Dutch, English, French, German, "
"Italian, Japanese, Korean, Polish, Spanish\n"
"- **Voice cloning**: zero-shot from a reference clip + transcript\n"
"- **44.1 kHz** output, ~21.5 codec frames per second"
)
stream_checkbox.change(fn=on_stream_change, inputs=[stream_checkbox], outputs=[response_format])
all_inputs = [text_input, ref_audio, ref_audio_url, ref_text, response_format]
def dispatch(stream_enabled, *args):
if stream_enabled:
yield from generate_speech_stream(api_base, *args, chunk_seconds=stream_chunk_seconds)
else:
yield generate_speech(api_base, *args)
generate_btn.click(fn=dispatch, inputs=[stream_checkbox] + all_inputs, outputs=[audio_output])
demo.queue()
return demo
def main() -> None:
parser = argparse.ArgumentParser(description="Gradio demo for Audio8 TTS Preview")
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})")
parser.add_argument("--host", default="0.0.0.0", help="Gradio host")
parser.add_argument("--port", type=int, default=7861, help="Gradio port")
parser.add_argument("--share", action="store_true", help="Create a public share link")
parser.add_argument("--stream-chunk-seconds", type=float, default=0.5, help="Seconds of audio per yielded chunk")
args = parser.parse_args()
print(f"Connecting to vLLM server at: {args.api_base}")
build_interface(args.api_base, stream_chunk_seconds=args.stream_chunk_seconds).launch(
server_name=args.host, server_port=args.port, share=args.share
)
if __name__ == "__main__":
main()
@@ -0,0 +1,44 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# Launch the Audio8 TTS Preview server + Gradio demo together.
#
# Usage:
# ./run_gradio_demo.sh
# CUDA_VISIBLE_DEVICES=0 PORT=8092 GRADIO_PORT=7861 ./run_gradio_demo.sh
set -e
MODEL="${MODEL:-Audio8/Audio8-TTS-Preview-0.6b}"
PORT="${PORT:-8092}"
GRADIO_PORT="${GRADIO_PORT:-7861}"
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
echo "Starting Audio8 TTS Preview server (port $PORT)..."
vllm serve "$MODEL" \
--omni \
--host 0.0.0.0 \
--port "$PORT" &
SERVER_PID=$!
cleanup() {
echo "Stopping server (PID $SERVER_PID)..."
kill $SERVER_PID 2>/dev/null
wait $SERVER_PID 2>/dev/null
}
trap cleanup EXIT
echo "Waiting for server to start..."
for _ in $(seq 1 120); do
if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then
echo "Server ready."
break
fi
sleep 2
done
echo "Starting Gradio demo (port $GRADIO_PORT)..."
python "$SCRIPT_DIR/gradio_demo.py" \
--api-base "http://localhost:$PORT" \
--port "$GRADIO_PORT"
@@ -0,0 +1,24 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# Launch a vLLM-Omni server for Audio8 TTS Preview 0.6B.
#
# Usage:
# ./run_server.sh
# CUDA_VISIBLE_DEVICES=0 PORT=8092 ./run_server.sh
set -e
MODEL="${MODEL:-Audio8/Audio8-TTS-Preview-0.6b}"
PORT="${PORT:-8092}"
echo "Starting Audio8 TTS Preview server with model: $MODEL"
# --trust-remote-code is deliberately NOT passed: vllm-omni registers its own
# `arktts` config, and transformers would otherwise prefer the checkpoint's
# remote code and bypass it.
vllm serve "$MODEL" \
--omni \
--host 0.0.0.0 \
--port "$PORT"
@@ -0,0 +1,131 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Client for Audio8 TTS Preview via the /v1/audio/speech endpoint.
Examples:
# Basic TTS
python speech_client.py --text "Welcome to Audio8 TTS."
# Zero-shot voice cloning (the transcript must match the reference audio)
python speech_client.py --text "Welcome to Audio8 TTS." \
--ref-audio ref.wav --ref-text "The exact transcript of the reference recording."
# Streaming PCM output
python speech_client.py --text "Welcome to Audio8 TTS." --stream --output output.pcm
"""
import argparse
import base64
import os
import httpx
DEFAULT_API_BASE = "http://localhost:8092"
DEFAULT_API_KEY = "EMPTY"
DEFAULT_MODEL = "Audio8/Audio8-TTS-Preview-0.6b"
#: The codec runs at 44.1 kHz, so raw PCM must be played back at that rate.
PCM_SAMPLE_RATE = 44100
def encode_audio_to_base64(audio_path: str) -> str:
"""Encode a local audio file as a base64 data URL."""
if not os.path.exists(audio_path):
raise FileNotFoundError(f"Audio file not found: {audio_path}")
ext = audio_path.lower().rsplit(".", 1)[-1]
mime_type = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac", "ogg": "audio/ogg"}.get(
ext, "audio/wav"
)
with open(audio_path, "rb") as handle:
audio_b64 = base64.b64encode(handle.read()).decode("utf-8")
return f"data:{mime_type};base64,{audio_b64}"
def build_payload(args) -> dict:
payload = {
"model": args.model,
"input": args.text,
"response_format": args.response_format,
}
if args.ref_audio:
if args.ref_audio.startswith(("http://", "https://")):
payload["ref_audio"] = args.ref_audio
else:
payload["ref_audio"] = encode_audio_to_base64(args.ref_audio)
if args.ref_text:
payload["ref_text"] = args.ref_text
if args.max_new_tokens is not None:
payload["max_new_tokens"] = args.max_new_tokens
if args.stream:
payload["stream"] = True
payload["stream_format"] = "audio"
payload["response_format"] = "pcm"
return payload
def run_tts(args) -> None:
payload = build_payload(args)
api_url = f"{args.api_base}/v1/audio/speech"
headers = {"Content-Type": "application/json", "Authorization": f"Bearer {args.api_key}"}
print(f"Model: {args.model}")
print(f"Text: {args.text}")
if args.ref_audio:
print(f"Voice cloning: ref_audio={args.ref_audio!r} ref_text={args.ref_text!r}")
if args.stream:
output_path = args.output or "output.pcm"
with (
httpx.Client(timeout=300.0) as client,
client.stream("POST", api_url, json=payload, headers=headers) as resp,
):
if resp.status_code != 200:
print(f"Error {resp.status_code}: {resp.read().decode()}")
return
total_bytes = 0
with open(output_path, "wb") as handle:
for chunk in resp.iter_bytes():
handle.write(chunk)
total_bytes += len(chunk)
seconds = total_bytes / 2 / PCM_SAMPLE_RATE
print(f"Streamed {total_bytes} bytes (~{seconds:.2f}s of s16le @ {PCM_SAMPLE_RATE} Hz) to {output_path}")
print(f"Play with: ffplay -f s16le -ar {PCM_SAMPLE_RATE} -ac 1 {output_path}")
return
with httpx.Client(timeout=300.0) as client:
response = client.post(api_url, json=payload, headers=headers)
if response.status_code != 200:
print(f"Error {response.status_code}: {response.text}")
return
if response.headers.get("content-type", "").startswith("application/json"):
print(f"Error: {response.text}")
return
output_path = args.output or f"output.{args.response_format}"
with open(output_path, "wb") as handle:
handle.write(response.content)
print(f"Audio saved to: {output_path} ({len(response.content)} bytes)")
def main() -> None:
parser = argparse.ArgumentParser(description="Audio8 TTS Preview client")
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name")
parser.add_argument("--text", required=True, help="Text to synthesize")
parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning (path or URL)")
parser.add_argument("--ref-text", default=None, help="Transcript of the reference audio")
parser.add_argument("--max-new-tokens", type=int, default=None, help="Cap the number of generated codec frames")
parser.add_argument("--stream", action="store_true", help="Stream PCM instead of returning a whole file")
parser.add_argument(
"--response-format",
default="wav",
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
help="Audio format (default: wav)",
)
parser.add_argument("--output", "-o", default=None, help="Output file path")
run_tts(parser.parse_args())
if __name__ == "__main__":
main()
@@ -24,6 +24,11 @@ Usage:
--text "Hello world. How are you? I am fine." \
--simulate-stt --stt-delay 0.1
# Opt into per-sentence synthesis (Indic danda, CJK, Latin)
python streaming_speech_client.py \
--text "नमस्ते। कैसे हो?" \
--split-granularity sentence
# Receive JSON sidecar chunks with word-level timestamps
python streaming_speech_client.py \
--text "Hello world. How are you?" \
@@ -73,8 +78,9 @@ except ImportError:
def frame_basename(msg: dict) -> str:
"""Name a file after the utterance and sentence a frame belongs to.
Every utterance is sentence 0, so the utterance index is what keeps the
files of one connection from overwriting each other.
Every utterance uses `sentence_index` for units inside that flush.
With default `split_granularity=none` that index stays 0; sentence mode
uses 0, 1, ... so files from one connection do not overwrite each other.
"""
return f"utterance_{msg['utterance_index']:03d}_sentence_{msg['sentence_index']:03d}"
@@ -368,6 +374,13 @@ def main():
)
parser.add_argument("--speed", type=float, default=1.0, help="Playback speed (0.25-4.0)")
parser.add_argument("--max-new-tokens", type=int, default=None, help="Max tokens")
parser.add_argument("--seed", type=int, default=None, help="Sampling seed forwarded to the server")
parser.add_argument(
"--split-granularity",
default=None,
choices=["none", "sentence", "clause"],
help="Text split mode (default: server none = one request per input.done)",
)
# Base task options
parser.add_argument("--ref-audio", default=None, help="Reference audio")
@@ -409,6 +422,8 @@ def main():
"max_new_tokens",
"ref_audio",
"ref_text",
"seed",
"split_granularity",
]:
val = getattr(args, key.replace("-", "_"), None)
if val is not None:
+7 -5
View File
@@ -53,7 +53,9 @@ dev = [
"soundfile==0.14.0",
"imageio[ffmpeg]==2.37.4",
"opencv-python==5.0.0.93",
"mooncake-transfer-engine-cuda13==0.3.11.post1",
# 0.3.12 includes acknowledged TCP writes (Mooncake #2850): completion
# must mean the receive buffer is ready before zero-copy consumers read it.
"mooncake-transfer-engine-cuda13==0.3.12.post1",
"av==18.0.0", # for ComfyUI tests
"openpyxl==3.1.5", # for nightly CI
"pyttsx3==2.99",
@@ -75,7 +77,6 @@ dev = [
# Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo,
# whose top-level stepaudio2 package conflicts with this distribution.
"step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml)
"silero-vad==6.1", # Realtime server-VAD duplex E2E
]
demo = [
@@ -85,9 +86,6 @@ demo = [
"cn2an>=0.5.22", # Chinese number normalization for TTS (e.g. "1次" -> "一次")
]
server-vad = [
"silero-vad==6.1",
]
forced-aligner = [
"qwen-asr>=0.0.6", # faithful Japanese/Korean word segmentation for the TTS forced aligner
@@ -376,3 +374,7 @@ ba = "ba"
SER = "SER"
# NumPy ndarray.flags.writeable (canonical flag name, not "writable").
writeable = "writeable"
# "convnet" abbreviates convolutional network. It appears in the vendored Fish
# Speech DAC codec identifiers (convnet_type, transconvnet_type) and is not a
# typo of "convent".
convnet = "convnet"
+49 -14
View File
@@ -1,12 +1,13 @@
# Boogu-Image
> Text-to-image and image-editing online serving
(Boogu-Image-0.1-Base / -Edit / -Edit-Turbo)
(Boogu-Image-0.1-Base / -Turbo / -Edit / -Edit-Turbo)
## Summary
- Vendor: Boogu
- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image),
`Boogu/Boogu-Image-0.1-Turbo` (four-step text-to-image),
`Boogu/Boogu-Image-0.1-Edit` (image editing), and
`Boogu/Boogu-Image-0.1-Edit-Turbo` (four-step image editing)
- Task: Text-to-image generation and text-guided image editing (TI2I)
@@ -16,8 +17,9 @@
## When to use this recipe
Use this recipe when you want a known-good starting point for serving
`Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, or
`Boogu/Boogu-Image-0.1-Edit-Turbo` with vLLM-Omni's native pipeline (no
`Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Turbo`,
`Boogu/Boogu-Image-0.1-Edit`, or `Boogu/Boogu-Image-0.1-Edit-Turbo` with
vLLM-Omni's native pipeline (no
`--diffusion-load-format diffusers`, and the upstream `boogu` package is not
required).
@@ -25,12 +27,14 @@ Boogu-Image-0.1 is an Apache-2.0 unified image generation and editing model
family. The Base text-to-image checkpoint pairs a Qwen3-VL multimodal encoder
with a Diffusion Transformer (DiT) and a flow-match Euler scheduler with
time-shift. It handles photorealistic generation and Chinese/English text
rendering. The Edit and Edit-Turbo checkpoints use the same native
`BooguImagePipeline` for image editing.
rendering. The Edit checkpoint uses the native `BooguImagePipeline`, while the
distilled Turbo and Edit-Turbo checkpoints run the four-step DMD path through
`BooguImageTurboPipeline`.
## References
- Upstream model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Base>
- Turbo model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Turbo>
- Edit model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Edit>
- Edit-Turbo 1K hotfix: <https://huggingface.co/Boogu/Boogu-Image-0.1-Edit-Turbo/tree/hotfix-1k-20260708>
- Project page: <https://boogu.org>
@@ -186,6 +190,19 @@ ID passed to `vllm serve`. For example, use
corresponding Edit model.
Base text-to-image:
## Fast text-to-image (Boogu-Image-0.1-Turbo)
Turbo is the distilled text-to-image checkpoint. Its `model_index.json`
declares `BooguImageTurboPipeline`, which runs the four-step DMD path instead
of the scheduler-driven path used by Base.
### Command
```bash
vllm serve Boogu/Boogu-Image-0.1-Turbo --omni --port 8091
```
### Verification
```bash
curl -X POST http://localhost:8091/v1/images/generations \
@@ -230,6 +247,23 @@ curl -X POST http://localhost:8091/v1/images/edits \
checkpoint-declared FP8 configuration. With the regular checkpoints used for
online FP8, the MLLM remains in BF16. The VAE and scheduler are outside this
quantization routing.
"model": "Boogu/Boogu-Image-0.1-Turbo",
"prompt": "A mountain lake at sunset, photorealistic, cinematic lighting",
"size": "1024x1024",
"seed": 42
}' | jq -r '.data[0].b64_json' | base64 -d > output-turbo.png
```
### Notes
- **Required settings:** `guidance_scale=1.0` and `guidance_scale_2=1.0`; any
other value is rejected. `num_inference_steps` defaults to 4. All three are
pipeline defaults, so the request above omits them.
- **Resolution:** upstream constrains Turbo to 1K, so `1024x1024` is the
supported working resolution. The shared pipeline still clamps at 2K; larger
sizes are not validated for this checkpoint.
- **Known limitations:** the same single-GPU limitations as Base apply; CPU
offload, Cache-DiT, and multi-GPU parallelism are not yet validated.
## Image editing (Boogu-Image-0.1-Edit)
@@ -354,16 +388,18 @@ curl -s http://localhost:8091/v1/chat/completions \
## Fast image editing (Boogu-Image-0.1-Edit-Turbo)
Edit-Turbo is the distilled image-editing checkpoint. Its `model_index.json`
declares `BooguImagePipeline`, so it uses the same native TI2I path as Edit.
The upstream 1K hotfix is recommended for stable results and is pinned below
to avoid silently loading the older checkpoint from the repository's default
revision.
declares `BooguImagePipeline`, so `--model-class-name BooguImageTurboPipeline`
is required to reach the four-step DMD path; without it the checkpoint is
served on the dense Edit path. The upstream 1K hotfix is recommended for
stable results and is pinned below to avoid silently loading the older
checkpoint from the repository's default revision.
### Command
```bash
vllm serve Boogu/Boogu-Image-0.1-Edit-Turbo \
--omni \
--model-class-name BooguImageTurboPipeline \
--revision hotfix-1k-20260708 \
--port 8091
```
@@ -407,11 +443,10 @@ curl -s http://localhost:8091/v1/chat/completions \
- **Pinned revision:** use `hotfix-1k-20260708`; upstream recommends the 1K
hotfix over the 1.5K variant for more stable results.
- **Recommended settings:** `num_inference_steps=4` and
`guidance_scale=1.0`. Edit-Turbo is guidance-distilled, meaning the guidance
behavior is already baked into the distilled checkpoint. A scale of `1.0`
avoids applying additional classifier-free guidance; unlike the regular Edit
checkpoint, it should not be increased to `5.0`.
- **Required settings:** `guidance_scale=1.0` and `guidance_scale_2=1.0`; the
Turbo pipeline rejects any other value, so the `guidance_scale_2=2.0` shown
for the regular Edit checkpoint does not apply here. `num_inference_steps`
defaults to 4.
- **Reference images:** the same single-reference and `align_res` behavior as
the regular Edit checkpoint applies.
- **Known limitations:** the same single-GPU limitations as Base and Edit apply;
+48
View File
@@ -219,3 +219,51 @@ the audio-track caveat.
5 chunks, `key_frames=0` = summarize all chunk frames) are aligned to the JoyVL
reference adapter so per-tick behavior matches the released model; the framework only
supplies the serving structure.
## Native multi-stage pipeline (experimental)
An experimental native speech path is available alongside the existing Day-0 serving path:
```text
one request with image/video and text
-> JoyAI complete action text
-> Qwen3-TTS Talker
-> Code2Wav
-> 24 kHz audio
```
It runs behind one `vllm serve --omni` endpoint:
```bash
vllm serve jdopensource/JoyAI-VL-Interaction-Preview --omni \
--deploy-config vllm_omni/deploy/joyai_vl_interaction.yaml \
--port 8091
```
With `modalities: ["text", "audio"]`, its behavior is:
| JoyAI action | Native output |
|---|---|
| `</response> <text>` | action text + speech for `<text>` |
| `</silence>` | action text + empty audio output (zero samples) |
| `</response> <note> </delegation> <question>` | action text + speech for `<note>` only |
Native speech supports the Qwen3-TTS `CustomVoice` task. The request-level `voice` and
`language` fields select the speaker and language; when omitted, they default to
`Vivian` and `Auto`. Values supplied as `tts_speaker` and `tts_language` inside
`additional_information` take priority.
This experimental native multi-stage pipeline is request-scoped, stateless, and
all-sync. It currently does not include:
- continuous frame sessions, persistent standing instructions, or session memory;
- audio input / ASR, full-duplex interaction, or barge-in;
- background Agent execution for delegated questions;
- async-chunk support from Talker to Code2Wav;
- integration with the existing WebUI.
For actions with spoken text, the Talker starts only after the complete JoyAI action is
available, and Code2Wav starts only after the Talker completes. Silence actions skip both
TTS stages and return an empty audio output with zero samples. All three stages are placed
on GPU 0 by the provided
[`deploy config`](../../vllm_omni/deploy/joyai_vl_interaction.yaml).
+23
View File
@@ -110,6 +110,29 @@ vllm serve Lightricks/LTX-2.5-Diffusers \
Use the T2V or I2V request below after the server is ready.
### Multi-GPU low-latency path
For low-latency 1920x1088 distilled two-stage serving, match the Ulysses and
VAE patch-parallel degrees to the GPU count. For example, on four GPUs:
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3 \
vllm serve Lightricks/LTX-2.5-Diffusers \
--omni \
--model-class-name LTX2DistilledTwoStagePipeline \
--usp 4 \
--vae-patch-parallel-size 4 \
--vae-use-tiling \
--host 0.0.0.0 \
--port 8000 \
--stage-init-timeout 900
```
Strict Ulysses only: `((F - 1) / 8 + 1) * (H / 32) * (W / 32)` and TP-local
SP attention head counts must be divisible by `ulysses_degree`. Use each
phase's `H, W` (half/full resolution for two-stage); non-divisible shapes fail.
Each GPU must still fit the resident weights. Warm up with the production shape.
## Offline inference
Choose values from the pipeline table. For example, the distilled two-stage
+8 -3
View File
@@ -101,6 +101,12 @@ sigmas and input latents.
Ordinary two-stage uses layer-fused LoRA for an unquantized BF16 Transformer
and automatically switches to dynamic LoRA when quantization is enabled.
## Sequence Parallelism
Strict Ulysses only: `((F - 1) / 8 + 1) * (H / 32) * (W / 32)` and TP-local
SP attention head counts must be divisible by `ulysses_degree`. Use each
phase's `H, W` (half/full resolution for two-stage); non-divisible shapes fail.
## Serving
Start either one-stage checkpoint:
@@ -335,9 +341,8 @@ bundled offline CLI do not currently expose `sigmas`.
`--max-num-seqs 1`.
- `--cfg-parallel-size` shards the complete LTX guidance plan, including STG,
modality guidance, and rescale-compatible prediction gathering.
- Sequence parallelism may pad audio latents. Pure Ulysses masks the padding;
Ring cannot, so audio length must be SP-divisible. Use `ring_degree=1` or a
divisible request shape.
- Ulysses masks padded audio latents; video sequences must satisfy the
[sequence-parallel divisibility constraints](#sequence-parallelism).
- Cache-DiT is one-stage only; multi-stage configurations are rejected.
## Operational Notes
+13 -5
View File
@@ -68,16 +68,24 @@ The `/v1/videos` request schema and `extra_params.task` values (`t2va`,
## Turbo LoRA
MiniMax-H3 Turbo uses five sigma points, `flow_shift=6`, and
`audio_flow_shift=3`. Start Turbo deployments with the dedicated defaults so
requests that omit sampling controls do not inherit the 50-step base schedule:
The deploy config below carries the four-step 768p contract -- five sigma
points, `flow_shift=6`, `audio_flow_shift=3` -- so requests that omit sampling
controls do not inherit the 50-step base schedule:
```bash
vllm-omni serve MiniMaxAI/MiniMax-H3 \
--omni \
--lora-path /path/to/MiniMax-H3-Turbo \
--lora-path /path/to/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors \
--deploy-config vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml
```
`--lora-path` must name one artifact, or a directory holding exactly one; the
Turbo repository holds several and a directory of them is rejected as
ambiguous. Serving a different artifact means overriding those defaults with
that artifact's own contract -- see the
[Turbo LoRA table](MiniMax-H3.md#turbo-lora) for every published row.
The base deployment intentionally retains 50 inference steps for non-LoRA
quality. Turbo LoRA supports T2VA and FL2VA requests, not Ref2VA.
quality. This disaggregated topology is validated on the `fl2v` artifacts,
which serve T2VA and FL2VA; the `ref2v` artifacts need a `--task-type ref2va`
server.
+70 -12
View File
@@ -869,13 +869,51 @@ balanced switch order.
### Turbo LoRA
Only the native Diffusers 4-step FL2VA/T2VA v1.0 artifact is supported:
The eight Diffusers-layout LightX2V Turbo artifacts are supported. The
filename records the contract, so the server reads the step count, task family
and flow shift from it and validates each request against the artifact that is
loaded. It does not rewrite request or deploy-config sampling values: the
request must carry that artifact's own settings, listed here, or it is
rejected.
```text
lightx2v/Minimax-h3-Turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors
```
| Artifact | Task | Forwards | `num_inference_steps` | `flow_shift` | declared `alpha` |
| --- | --- | ---: | ---: | ---: | ---: |
| `minimax_h3_fl2v_turbo_4step_v0.1.safetensors` | T2VA / FL2VA | 4 | 5 | 12 | none -> 8 |
| `minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors` | T2VA / FL2VA | 4 | 5 | 6 | 128 |
| `minimax_h3_fl2v_turbo_4step_v1.1_768p_bf16.safetensors` | T2VA / FL2VA | 4 | 5 | 6 | 128 |
| `minimax_h3_fl2v_turbo_4step_v1.2_768p_bf16.safetensors` | T2VA / FL2VA | 4 | 5 | 6 | 8 |
| `minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors` | T2VA / FL2VA | 8 | 9 | 12 | 8 |
| `minimax_h3_fl2v_turbo_8step_v1.0_768p_bf16.safetensors` | T2VA / FL2VA | 8 | 9 | 6 | 8 |
| `minimax_h3_ref2v_turbo_4step_v0.1_bf16.safetensors` | Ref2VA | 4 | 5 | 12 | 8 |
| `minimax_h3_ref2v_turbo_8step_v1.0_768p_bf16.safetensors` | Ref2VA | 8 | 9 | 6 | 8 |
Download only that file:
`audio_flow_shift` is `3.0` across the family. Each row is the complete
published filename; use it verbatim as `TURBO_FILE` below.
Alpha needs no manual compensation: the server reads it from the artifact's
metadata, falling back to 8 with a warning for
`minimax_h3_fl2v_turbo_4step_v0.1`, the one artifact that declares none. The
request-level `scale` is a further multiplier on top of it.
> [!NOTE]
> Every artifact is rank 128 and the delta is applied at `scale * alpha /
> rank`, so the `alpha=8` rows drive at 1/16 the strength of the `alpha=128`
> rows. 8 is the default of LightX2V's reference script, which never reads the
> metadata; its documented `v0.1` command does not override that default.
Every artifact except `minimax_h3_fl2v_turbo_4step_v0.1` also ships a
`_comfyui_` export of the same weights. Those fuse Q/K/V into one projection
and are **not** supported; downloading one is refused by name. Take the
Diffusers file. The filename is the contract, so do not rename an artifact
either -- a renamed file is rejected rather than served on a guess.
FL2VA artifacts serve `t2va` and `fl2va` on any FL2VA or combined server.
Ref2VA artifacts require
`--task-type ref2va`: a combined server serves `ref2va` from a second DiT that
the adapter cannot bind to, so loading one there is refused rather than silently
running an undistilled model on the few-step schedule.
Download the artifact you want:
```bash
export TURBO_DIR=/path/to/minimax-h3-turbo
@@ -884,10 +922,19 @@ hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}"
export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}"
```
`--lora-path` accepts one artifact, or a directory holding exactly one.
> [!IMPORTANT]
> This changes earlier behaviour. `--lora-path /path/to/minimax-h3-turbo`
> pointing at a full clone of the Turbo repository used to select the v1.0 768p
> file implicitly; a directory holding several recognized artifacts is now
> rejected as ambiguous. Name the artifact instead:
> `--lora-path /path/to/minimax-h3-turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors`.
Start from a non-offloaded or DLO FL2VA server command and add
`--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`.
`--lora-path` preloads the adapter; each request still activates it and uses
the published sampling settings:
`--lora-path` preloads the adapter; each request still activates it and must
carry that artifact's sampling settings:
```bash
-F 'num_inference_steps=5' \
@@ -896,14 +943,25 @@ the published sampling settings:
-F "lora={\"name\":\"h3-turbo-v1.0\",\"path\":\"${TURBO_LORA}\",\"scale\":1.0}"
```
For FL2VA, change `task` and add `input_reference` as shown above. The 8-step,
ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is
dynamic-only and does not support prefusion or LoRA composition. DLO is
Switching to another FL2VA artifact means repointing `TURBO_FILE`, which moves
both `--lora-path` and the request's `lora.path`, and carrying that row's
`num_inference_steps` and `flow_shift`: `9` and `6` for `8step_v1.0_768p`, `9`
and `12` for the 544p `8step_v1.0`. A request that does not match the loaded
artifact is rejected, so a mismatch cannot silently degrade output.
The two `ref2v` rows are not served by this FL2VA command. Start a
`--task-type ref2va` server, take a request from
[Ref2VA](#3-ref2va-image-only-imageaudio-or-mixed-references) and override
`num_inference_steps` and `flow_shift` with that row's values; those examples
already send `audio_flow_shift=3.0`.
For FL2VA, change `task` and add `input_reference` as shown above. This
integration is dynamic-only and does not support prefusion or LoRA composition. DLO is
supported by keeping the request-switchable LoRA A/B buffers resident on the
accelerator while DLO streams only the base blocks; budget for this additional
fixed HBM usage. Model-level and standard layerwise offload remain unsupported.
The five requested sigma points produce the four denoiser evaluations expected
by the Turbo artifact.
The requested sigma points always number one more than the artifact's denoiser
evaluations.
### FlashGen native LoRA
+10 -2
View File
@@ -16,7 +16,7 @@ Use this recipe for 24 kHz multilingual TTS with voice cloning (20 languages inc
Chinese and English). Choose a variant based on your latency and quality requirements:
| Model | Params | Use case |
|---|---|---|
| --- | --- | --- |
| MOSS-TTS | 8B | General TTS, highest quality |
| MOSS-TTS-v1.5 | 8B | General TTS upgrade of 1.0: 31 languages, steadier cloning, `[pause Xs]` markers (set `language` for best results); same `MossTTSDelay` API |
| MOSS-TTS-Realtime | 1.7B | Lowest latency (TTFB ~180 ms), streaming-first |
@@ -24,9 +24,17 @@ Chinese and English). Choose a variant based on your latency and quality require
| MOSS-SoundEffect | 8B | Sound effect synthesis from text description |
| MOSS-VoiceGenerator | 1.7B | Zero-shot voice design |
All variants share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and
The variants above share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and
output 24 kHz mono audio.
MOSS-TTS-Local-Transformer-v1.5 uses MOSS-Audio-Tokenizer-v2 and outputs 48 kHz
stereo audio. For Local voice cloning through `/v1/audio/speech`, provide an
accurate `ref_text` transcript alongside `ref_audio`. The adapter uses the
reference transcript and audio as a continuation prefix, then generates only
the requested `input` speech. Without a nonblank `ref_text`, Local uses
audio-reference generation. Reference transcripts must match the reference
audio; they are not style instructions.
## References
- Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)
+1
View File
@@ -85,6 +85,7 @@ recipes/
| [`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-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 1x datacenter Blackwell (B200/B300) |
| [`Wan-AI/Wan2.2-TI2V-5B-A100.md`](./Wan-AI/Wan2.2-TI2V-5B-A100.md) | Unified 720p text-to-video and image-to-video serving (Wan2.2 5B) | 1x A100-SXM4-80GB |
| [`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) |
+111 -10
View File
@@ -7,13 +7,31 @@
- Vendor: Robbyant
- Model: `robbyant/lingbot-world-v2-14b-causal-fast-diffusers`
- Task: image-conditioned interactive world generation
- Modes: offline trajectory replay and in-process realtime AR-Diffusion ticks
- Modes: offline trajectory replay, stepwise one-request streaming, and in-process realtime AR-Diffusion ticks (deprecated)
- Hardware validated: NVIDIA H200 and B200
- Maintainer: Community
The checkpoint is separately licensed under CC BY-NC-SA and restricted to
non-commercial use. The vLLM-Omni integration code remains Apache-2.0.
## Which path to use
There are three ways to drive this model, each documented in its own section
below.
- **Realtime stepwise** — see [Streaming video serving](#streaming-video-serving).
The suggested path: one `WS /v1/realtime/video` session produces the whole
rollout, one video chunk per AR block. It cannot take realtime camera
interaction yet; that is being added, and this path is the one that will
remain.
- **Realtime tick** *(deprecated)* — see
[Realtime in-process generation](#realtime-in-process-generation-deprecated).
The only path that takes realtime camera interaction today, so it stays until
the stepwise path can, but it takes no new work and will be removed.
- **Offline** — see [Offline generation](#offline-generation). Replays a fixed
pose/intrinsics trajectory in one request and writes an MP4. Use it when the
camera path is known up front and streaming is not needed.
## Offline generation
The offline path consumes one source image and an action directory containing
@@ -34,7 +52,12 @@ counts must be `9 + 12k`, up to 117 frames. Camera files may contain a
longer bounded source trajectory (the official example contains 269 frames);
the runtime consumes only the prefix needed by the request.
## Realtime in-process generation
## Realtime in-process generation (deprecated)
Use [Streaming video serving](#streaming-video-serving) instead unless you need
mid-session camera interaction, which only this path has. It still works and
emits no runtime warning; removal is tracked as B4 of the LingBot World 2.0
roadmap ([#6672](https://github.com/vllm-project/vllm-omni/issues/6672)).
The realtime example keeps AR-Diffusion KV and model-owned state across
requests. Each JSONL line describes the prompt and/or three latent-frame camera
@@ -61,12 +84,81 @@ exercises the same `ARDiffusionSessionManager -> ARDiffusionOmniTickConsumer
-> AsyncOmni -> ARDiffusionEngine` path used by a future HTTP or WebSocket
transport.
This PR intentionally does not define a public realtime HTTP/WebSocket schema.
An online serving client should be added together with the public transport
API rather than exposing LingBot-specific event fields from the generic model
runtime. The structured camera-interaction frontend and transport work is
tracked separately in
[vllm-project/vllm-omni#5527](https://github.com/vllm-project/vllm-omni/pull/5527).
Realtime serving now goes through the generic `WS /v1/realtime/video`
transport, which exposes no LingBot-specific event fields; the tick loop above
stays available only for the older control plane.
## Streaming video serving
The stepwise path keeps AR-Diffusion paged KV but issues **one** request for
the whole rollout: `prepare_encode` runs once, then every AR block is four DMD
steps followed by one streamed chunk. Identity metadata uses
`session_id = request_id` with contiguous `chunk_index` values from zero.
Camera motion for this landing is request-scoped: one three-frame action list
per chunk, fixed when the rollout starts.
Serve it with the AR-Diffusion deploy config, which selects the AR-Diffusion
engine and enables streamed step execution:
```bash
vllm serve robbyant/lingbot-world-v2-14b-causal-fast-diffusers \
--omni \
--deploy-config vllm_omni/deploy/lingbot_world_v2_stepwise.yaml \
--port 8000
```
`--deploy-config` is required: the `lingbot_world` pipeline deliberately
registers no default deploy config, so offline replay and the tick example
keep their request-mode topology when no deploy config is given, and the
stepwise serving topology is only ever an explicit choice.
Clients then use the generic WebSocket protocol documented in
[`docs/serving/video_stream_api.md`](../../docs/serving/video_stream_api.md):
`session.start` begins one rollout and each AR block arrives as a binary video
chunk. The model is image-conditioned, so the first frame is required and rides
on `image_reference` as an `http(s)` or `data:` URL; per-chunk camera actions
ride on `extra_params`:
```json
{"type": "session.start", "model": "robbyant/lingbot-world-v2-14b-causal-fast-diffusers",
"prompt": "The camera moves slowly forward through the scene.",
"image_reference": {"image_url": "data:image/png;base64,<first frame>"},
"width": 832, "height": 480, "num_frames": 33,
"extra_params": {"flow_shift": 5.0,
"camera_action_script": [[["w"], ["w"], ["w"]], [["a"], [], []], [[], [], []]]}}
```
`camera_action_script` must hold exactly one action list per generated
chunk, and a request generates `((num_frames - 1) // 4 + 1) // 3` chunks:
three for `num_frames: 33`, seven for `num_frames: 81`.
The bundled client speaks this protocol:
```bash
python examples/online_serving/streaming_video_generation/streaming_video_client.py \
--model robbyant/lingbot-world-v2-14b-causal-fast-diffusers \
--prompt "The camera moves slowly forward through the scene." \
--image-reference /path/to/first_frame.png \
--width 832 --height 480 --num-frames 33 --fps 16 --seed 42 \
--extra-params '{"camera_action_script":[[["w"],["w"],["w"]],[["a"],[],[]],[[],[],[]]]}' \
--output lingbot_world_v2_stream.mp4
```
A served request may instead point at a pose/intrinsics trajectory with
`extra_params.action_path`, which is resolved inside the trusted root set by
`model_config.lingbot_action_root` or `VLLM_OMNI_LINGBOT_ACTION_ROOT`; a server
started without that root configured accepts only `camera_action_script`.
Requested `width`/`height` must match `ar_diffusion_width`/`ar_diffusion_height`
in the deploy config, because the AR cache geometry is fixed at load time.
Blocks are decoded independently, so seams between chunks are possible; a
session-owned streaming decoder is tracked separately.
Mid-session `session.interaction` for camera control is not wired yet, so a
served rollout follows the script it started with.
The deprecated tick example remains available for the older
one-block-per-`generate()` control plane.
## Realtime identity and controls
@@ -104,10 +196,16 @@ tested commit.
## Current limitations
- Only the 14B causal-fast checkpoint is supported.
- The realtime control plane is internal; there is no public server transport yet.
- The tick control plane is internal; the public transport is the stepwise
`WS /v1/realtime/video` path, which cannot take mid-session camera
interaction yet.
- Stepwise serving requires an explicit
`--deploy-config vllm_omni/deploy/lingbot_world_v2_stepwise.yaml`; there is
no default deploy config for this model.
- AR-Diffusion stages currently require one replica because session-affine
routing across replicas is not implemented.
- One AR block is generated per request and `max_num_seqs` must be one.
- Tick mode generates one AR block per request. Stepwise mode generates many
AR blocks in one request. `max_num_seqs` must be one in both cases.
- Stateful streaming VAE decode is not implemented; the realtime example emits
latent chunks.
- SP/USP, pipeline/CFG parallelism, HSDP, VAE parallelism, quantization,
@@ -119,4 +217,7 @@ tested commit.
- Checkpoint: <https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers>
- Official implementation: <https://github.com/robbyant/lingbot-world-v2>
- Offline example: [`examples/offline_inference/diffusion/lingbot_world_v2.py`](../../examples/offline_inference/diffusion/lingbot_world_v2.py)
- Streaming serving client: [`examples/online_serving/streaming_video_generation/`](../../examples/online_serving/streaming_video_generation/README.md)
- Streaming serving end-to-end test: [`tests/e2e/online_serving/test_lingbot_world_v2_stepwise.py`](../../tests/e2e/online_serving/test_lingbot_world_v2_stepwise.py)
- Stepwise latent-level end-to-end test: [`tests/e2e/offline_inference/test_lingbot_world_v2_stepwise.py`](../../tests/e2e/offline_inference/test_lingbot_world_v2_stepwise.py)
- Realtime design: [`docs/design/feature/realtime_ar_diffusion.md`](../../docs/design/feature/realtime_ar_diffusion.md)
+184
View File
@@ -0,0 +1,184 @@
# Wan2.2 TI2V 5B — A100 80GB
> Single-GPU, native 720p text-to-video and image-to-video serving
## Summary
- Vendor: Wan-AI
- Model: `Wan-AI/Wan2.2-TI2V-5B-Diffusers`
- Task: Unified text-to-video (T2V) and image-to-video (I2V) generation
- Mode: Online serving with the OpenAI-compatible Videos API
- Hardware: 1x NVIDIA A100-SXM4-80GB
- Maintainer: Community
## When to use this recipe
Use this recipe to serve both T2V and I2V from one dense 5B checkpoint on one
A100 80GB. It qualifies the model's native 720p profile without CPU offload or
multi-GPU parallelism. The synchronous endpoint is used below because it
returns the MP4 and request timing in one response; the asynchronous Videos API
uses the same server and sampling parameters.
## Supported model contract
| Task | Required input | Qualified output | Endpoint |
| --- | --- | --- | --- |
| T2V | Text prompt; omit `input_reference` | 1280x704, 121 frames, 24 fps | `POST /v1/videos/sync` |
| I2V | Text prompt plus one image uploaded as `input_reference` | 1280x704, 121 frames, 24 fps | `POST /v1/videos/sync` |
The upstream checkpoint supports both tasks at 720p and 24 fps. For native
landscape 720p, use `1280x704`; portrait generation uses `704x1280`. Wan video
lengths follow the temporal VAE contract (`4k+1` frames). This recipe validates
one output per request with 50 denoising steps and seed 42.
## References
- Upstream model card:
<https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers>
- Online T2V example:
[`examples/online_serving/text_to_video`](../../examples/online_serving/text_to_video)
- Online I2V example:
[`examples/online_serving/image_to_video`](../../examples/online_serving/image_to_video)
- Offline I2V example:
[`examples/offline_inference/image_to_video`](../../examples/offline_inference/image_to_video)
- Community recipe tracker:
[vllm-project/vllm-omni#2645](https://github.com/vllm-project/vllm-omni/issues/2645)
## Hardware
- Accelerator model and per-device memory: NVIDIA A100-SXM4-80GB, 81,920 MiB
- Number of devices: 1
- Device interconnect: not applicable to this single-GPU profile
- Host memory: 1 TiB installed on the qualification host; CPU offload was not used
- Qualification scope: BF16, batch size 1, native FlashAttention, eager mode,
VAE tiling, T2V and I2V at 1280x704 / 121 frames / 24 fps
## Software environment
- OS: Ubuntu 24.04.3 LTS, Linux 6.8.0
- Python: 3.12.3
- NVIDIA driver / host toolkit: 550.127.05 / CUDA 12.8
- PyTorch: 2.13.0+cu129
- Transformers / Diffusers: 5.14.1 / 0.40.0
- vLLM: 0.28.0+cu129
- vLLM-Omni: commit `ff6e906a25de1ca3122c0d0d5250fde5a0ddbc7b`
The default vLLM 0.28.0 wheel targets CUDA 13.0. Driver 550 cannot run that
variant, so this qualification used the official CUDA 12.9 wheel:
```bash
uv venv --python 3.12 --seed
source .venv/bin/activate
uv pip install \
'https://github.com/vllm-project/vllm/releases/download/v0.28.0/vllm-0.28.0%2Bcu129-cp38-abi3-manylinux_2_28_x86_64.whl' \
--torch-backend=auto
uv pip install -e .
```
## Command
Start the server from the repository root:
```bash
vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--omni \
--port 8092 \
--vae-use-tiling \
--enforce-eager
```
In another terminal, download the public reference image:
```bash
wget -O cherry_blossom.jpg \
https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
```
Generate a text-only video:
```bash
curl --fail-with-body -D t2v-headers.txt \
-X POST http://127.0.0.1:8092/v1/videos/sync \
-F "prompt=Two anthropomorphic cats in comfortable boxing gear and bright gloves fight intensely on a spotlighted stage, cinematic lighting, smooth motion." \
-F "negative_prompt=blurry, low quality, static, distorted, artifacts, watermark, text" \
-F "width=1280" -F "height=704" \
-F "num_frames=121" -F "fps=24" \
-F "num_inference_steps=50" -F "guidance_scale=5.0" \
-F "flow_shift=5.0" -F "seed=42" \
--output wan22-ti2v-5b-t2v.mp4
```
Generate an image-conditioned video from the same server:
```bash
curl --fail-with-body -D i2v-headers.txt \
-X POST http://127.0.0.1:8092/v1/videos/sync \
-F "input_reference=@cherry_blossom.jpg" \
-F "prompt=Cherry blossoms sway gently in the breeze while petals drift across the path, cinematic, smooth natural motion." \
-F "negative_prompt=blurry, low quality, static, distorted, artifacts, watermark, text" \
-F "width=1280" -F "height=704" \
-F "num_frames=121" -F "fps=24" \
-F "num_inference_steps=50" -F "guidance_scale=5.0" \
-F "flow_shift=5.0" -F "seed=42" \
--output wan22-ti2v-5b-i2v.mp4
```
## Verification
Check both MP4 files with `ffprobe`:
```bash
for video in wan22-ti2v-5b-t2v.mp4 wan22-ti2v-5b-i2v.mp4; do
ffprobe -v error -select_streams v:0 -count_frames \
-show_entries stream=codec_name,width,height,r_frame_rate,nb_read_frames,duration \
-of default=noprint_wrappers=1 "$video"
done
```
Each file reported:
```text
codec_name=h264
width=1280
height=704
r_frame_rate=24/1
duration=5.041667
nb_read_frames=121
```
The synchronous response headers and an independent 200 ms `nvidia-smi`
sampler produced these steady-state results after server warmup:
| Task | HTTP | `X-Inference-Time-S` | Server peak | Observed device peak | MP4 size |
| --- | ---: | ---: | ---: | ---: | ---: |
| T2V | 200 | 288.267 s | 26,546 MB | 27,533 MiB | 7,321,179 bytes |
| I2V | 200 | 299.439 s | 29,766 MB | 30,333 MiB | 15,667,146 bytes |
Both outputs were decoded in full and visually inspected. The T2V output
contained the prompted boxing cats. The I2V first frame preserved the uploaded
cherry-blossom scene (5.81 mean absolute pixel error after resize and H.264
encoding), and later frames showed motion.
## Notes
- Memory usage: model loading used 21.2239 GiB. The measured I2V device peak
was 30,333 MiB, leaving more than 50 GiB free on this A100.
- `--vae-use-tiling` is retained in the known-good command even though this
profile has ample headroom; it bounds decode memory at 720p.
- `--enforce-eager` removes graph-compilation variability from the measured
profile. Remove it only after separately warming and measuring compiled mode.
- `flow_shift=5.0` is the 720p setting. The shared 480p examples use 12.0.
- The generated MP4 contains video only; this checkpoint does not generate audio.
- The first request includes neither model-load time nor the server's startup
warmup. Server initialization took approximately 26 seconds with a warm cache.
## Supported features
| Feature | Status for this profile | Shared guide |
| --- | --- | --- |
| OpenAI-compatible Videos API | Qualified for synchronous T2V and I2V | [Videos API](../../docs/serving/videos_api.md) |
| VAE tiling | Enabled and qualified on one A100 | [VAE parallelism and tiling](../../docs/user_guide/diffusion/parallelism/vae_parallelism.md) |
| Cache-DiT / TaylorSeer | Available but not enabled in the measurements above | [Cache-DiT](../../docs/user_guide/diffusion/cache_acceleration/cache_dit.md) |
| TeaCache | Available but not qualified in this recipe | [TeaCache](../../docs/user_guide/diffusion/cache_acceleration/teacache.md) |
| CPU or layerwise offload | Not needed for the qualified 80GB profile | [CPU offload](../../docs/user_guide/diffusion/cpu_offload.md) |
| TP / SP / CFG parallelism | Not applicable to this single-GPU profile | [Parallelism overview](../../docs/user_guide/diffusion/parallelism/overview.md) |
+12
View File
@@ -207,6 +207,18 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
-F 'extra_params={"depth":{"control_path":"/path/to/depth_control.mp4"},"max_frames":121,"resolution":"720","num_video_frames_per_chunk":121}' \
-o cosmos3_transfer_depth.mp4
# The same transfer control can be uploaded by a remote client instead of
# being placed on the server filesystem. Uploads are limited to 512 MiB.
curl -sS -X POST http://localhost:8000/v1/videos/sync \
-H "Accept: video/mp4" \
-F "model=nvidia/Cosmos3-Nano" \
-F "prompt=Generate a realistic scene following the provided world-state control." \
-F "input_reference=@/path/to/input.mp4;type=video/mp4" \
-F "control_reference=@/path/to/wsm.mp4;type=video/mp4" \
-F "control_type=wsm" \
-F 'extra_params={"wsm":{"control_weight":1.0},"max_frames":121,"resolution":"720","num_video_frames_per_chunk":121}' \
-o cosmos3_transfer_wsm.mp4
# Text-to-video-with-sound
curl -sS -X POST http://localhost:8000/v1/videos/sync \
-H "Accept: video/mp4" \
+3 -10
View File
@@ -36,19 +36,13 @@ Other hardware and configurations are welcome as community validation lands.
- OS: Linux
- Python: 3.10+
- Required package: `fish-speech` (for DAC codec)
- CUDA 12.8
- vLLM version: 0.19.0
- vLLM-Omni version or commit: c93359bb354a6aa5c14d062430cb85b2c4db251e
```bash
# Install PortAudio dependency if missing (required by pyaudio which is a dependency of fish-speech)
# For Ubuntu/Debian:
if ! dpkg -l | grep -q libportaudio2; then
sudo apt-get update && sudo apt-get install -y libportaudio2 portaudio19-dev
fi
pip install fish-speech
```
No extra packages are required: the DAC codec architecture is vendored in
vLLM-Omni (`vllm_omni/model_executor/models/fish_speech/dac_modules/`), so the
model runs out of the box, including in the official Docker image.
#### Kvcache attention fast path
@@ -129,7 +123,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \
- OS: Linux
- Python: 3.10+
- Required package: `fish-speech` (for DAC codec)
- CUDA 12.8
- vLLM version: 0.19.0+
+131
View File
@@ -457,5 +457,136 @@ def test_text_benchmark_still_reports_ttft(capsys):
assert "Time to First Token" in out, "text benchmarks must keep TTFT"
# ============================================================================
# Suppress text metrics for pure image / video endpoints
# ============================================================================
def _make_image_output(prompt_len: int) -> MixRequestFuncOutput:
output = MixRequestFuncOutput()
output.success = True
output.prompt_len = prompt_len
output.output_tokens = 0
output.generated_text = ""
output.ttft = 0.0
output.text_latency = 0.0
output.latency = 2.5
output.start_time = 0.0
output.itl = []
output.image_count = 1
output.image_generation_time_ms = 2000.0
output.image_pixels = 1024 * 1024
output.input_audio_duration = 0.0
output.error = ""
return output
def _make_video_output(prompt_len: int) -> MixRequestFuncOutput:
output = MixRequestFuncOutput()
output.success = True
output.prompt_len = prompt_len
output.output_tokens = 0
output.generated_text = ""
output.ttft = 0.0
output.text_latency = 0.0
output.latency = 5.0
output.start_time = 0.0
output.itl = []
output.video_duration = 2.0
output.video_frames = 48
output.video_generation_time_ms = 4000.0
output.video_rtf = 2.0
output.input_audio_duration = 0.0
output.error = ""
return output
_MEDIA_PERCENTILE_METRICS = ["e2el", "ttft"]
def test_pure_image_benchmark_omits_text_result(capsys):
"""Pure image generation must not print Text Result or invent peak tok/s."""
outputs = [_make_image_output(64), _make_image_output(64)]
metrics, actual_output_lens = calculate_metrics(
input_requests=[],
outputs=outputs,
dur_s=5.0,
tokenizer=None,
selected_percentiles=[99.0],
goodput_config_dict={},
task_type=TaskType.GENERATION,
selected_percentile_metrics=_MEDIA_PERCENTILE_METRICS,
max_concurrency=None,
request_rate=float("inf"),
benchmark_duration=5.0,
)
out = capsys.readouterr().out
assert actual_output_lens == [0, 0]
assert metrics.total_output == 0
assert metrics.max_output_tokens_per_s == 0.0
assert " Text Result " not in out
assert "Peak output token throughput" not in out
assert "Time to First Token" not in out
assert " Image Result " in out
def test_pure_video_benchmark_omits_text_result(capsys):
"""Pure video generation must not print Text Result or invent peak tok/s."""
outputs = [_make_video_output(63), _make_video_output(63)]
metrics, actual_output_lens = calculate_metrics(
input_requests=[],
outputs=outputs,
dur_s=10.0,
tokenizer=None,
selected_percentiles=[99.0],
goodput_config_dict={},
task_type=TaskType.GENERATION,
selected_percentile_metrics=_MEDIA_PERCENTILE_METRICS,
max_concurrency=None,
request_rate=float("inf"),
benchmark_duration=10.0,
)
out = capsys.readouterr().out
assert actual_output_lens == [0, 0]
assert metrics.total_output == 0
assert metrics.max_output_tokens_per_s == 0.0
assert " Text Result " not in out
assert "Peak output token throughput" not in out
assert "Time to First Token" not in out
assert " Video Result " in out
def test_image_with_generated_text_still_reports_text_result(capsys):
"""Image edits / AR text alongside images must keep Text Result."""
output = _make_image_output(64)
output.output_tokens = 8
output.generated_text = "revised caption"
output.ttft = 0.05
output.itl = [0.01] * 7
calculate_metrics(
input_requests=[],
outputs=[output],
dur_s=3.0,
tokenizer=_EmptyAwareTokenizer(),
selected_percentiles=[99.0],
goodput_config_dict={},
task_type=TaskType.GENERATION,
selected_percentile_metrics=_MEDIA_PERCENTILE_METRICS,
max_concurrency=None,
request_rate=float("inf"),
benchmark_duration=3.0,
)
out = capsys.readouterr().out
assert " Text Result " in out
assert "Time to First Token" in out
assert " Image Result " in out
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])
+339
View File
@@ -22,9 +22,16 @@ from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
)
from vllm_omni.benchmarks.patch.patch import (
MixRequestFuncOutput,
_add_video_extra_body_to_form,
_add_video_reference_to_form,
_apply_stage0_token_timings,
_apply_video_metrics_from_payload,
_attach_seed_tts_to_request_func_input,
_build_benchmark_session,
_omni_request_timeout_s,
async_request_openai_chat_omni_completions,
async_request_openai_image_edits_omni,
async_request_openai_image_generations_omni,
async_request_openai_realtime_duplex,
should_request_stage_metrics,
)
@@ -1196,5 +1203,337 @@ async def test_prompt_len_assigned_from_usage(mocker: MockerFixture):
)
class TestOmniRequestTimeout:
"""``--omni-request-timeout-s`` precedence: explicit value > 900 s default."""
_OVERRIDE = "vllm_omni.benchmarks.patch.patch._REQUEST_TIMEOUT_OVERRIDE_S"
def test_default_timeout_is_900s_when_not_set(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(self._OVERRIDE, None)
assert _omni_request_timeout_s() == 900.0
def test_explicit_value_wins_over_default(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(self._OVERRIDE, 123.5)
assert _omni_request_timeout_s() == 123.5
def test_non_positive_restores_legacy_6h_cap(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(self._OVERRIDE, 0)
assert _omni_request_timeout_s() == 6 * 60 * 60.0
async def test_benchmark_session_uses_configured_timeout(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(self._OVERRIDE, 42.0)
session = _build_benchmark_session(max_concurrency=8, ssl_setting=False)
try:
assert session.timeout.total == 42.0
assert session.connector.limit == 8
assert session.connector.limit_per_host == 8
finally:
await session.close()
async def test_hung_server_request_times_out_as_failed(self, monkeypatch: pytest.MonkeyPatch) -> None:
"""A server that accepts the request but never responds must surface as ``failed``."""
async def handler(reader, writer):
# Drain the request but never respond — the original hang failure mode.
try:
while await reader.read(4096):
pass
except ConnectionResetError:
pass
finally:
# Close our side of the socket: the client aborts after the
# timeout and only sends FIN, so without writer.close() the
# half-closed connection keeps Server.wait_closed() (which on
# Python 3.12+ waits for every accepted connection to drop)
# blocked forever.
writer.close()
try:
await writer.wait_closed()
except (ConnectionResetError, BrokenPipeError):
pass
monkeypatch.setattr(self._OVERRIDE, 1.0)
server = await asyncio.start_server(handler, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
session = _build_benchmark_session(max_concurrency=1, ssl_setting=False)
try:
request_input = _seed_tts_request_func_input()
request_input.api_url = f"http://127.0.0.1:{port}/v1/chat/completions"
output = await async_request_openai_chat_omni_completions(request_input, session)
assert output.success is False
assert output.error
finally:
await session.close()
server.close()
await server.wait_closed()
def test_video_rtf_prefers_generation_time_over_poll_latency():
"""RTF must not grow with client poll_interval overshoot in output.latency."""
payload = {
"duration_s": 2.0,
"num_frames": 48,
"fps": 24.0,
"stage_durations": {"stage_0_gen_ms": 4000.0},
}
request_body: dict[str, object] = {}
short_poll = MixRequestFuncOutput()
short_poll.latency = 4.2 # ~generation + small poll overshoot
_apply_video_metrics_from_payload(short_poll, payload, request_body)
long_poll = MixRequestFuncOutput()
long_poll.latency = 8.0 # same job, larger poll_interval_s overshoot
_apply_video_metrics_from_payload(long_poll, payload, request_body)
assert short_poll.video_generation_time_ms == pytest.approx(4000.0)
assert long_poll.video_generation_time_ms == pytest.approx(4000.0)
assert short_poll.video_rtf == pytest.approx(2.0)
assert long_poll.video_rtf == pytest.approx(2.0)
assert short_poll.video_rtf == long_poll.video_rtf
def test_video_rtf_falls_back_to_e2e_latency_without_generation_time():
output = MixRequestFuncOutput()
output.latency = 6.0
_apply_video_metrics_from_payload(
output,
{"duration_s": 2.0, "num_frames": 48, "fps": 24.0},
{},
)
assert output.video_generation_time_ms == 0.0
assert output.video_rtf == pytest.approx(3.0)
# 1x1 PNG used as a valid b64_json image payload in edit-client tests.
_MIN_PNG_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
@pytest.mark.asyncio
async def test_image_edits_defaults_to_non_streaming_json(mocker: MockerFixture) -> None:
"""Single-stage servers reject stream=true; default path must send stream=false + JSON."""
class MockJsonResponse:
status = 200
def __init__(self):
self._payload = {
"created": 1,
"data": [
{
"b64_json": _MIN_PNG_B64,
"stage_durations": {"stage_0_gen_ms": 12.5},
}
],
}
async def json(self):
return self._payload
async def text(self):
return json.dumps(self._payload)
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
return None
captured_stream: list[str] = []
real_add_field = None
def tracking_add_field(self, name, value=None, **kwargs):
if name == "stream":
captured_stream.append(str(value))
return real_add_field(self, name, value, **kwargs)
import aiohttp
real_add_field = aiohttp.FormData.add_field
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
mock_session = mocker.AsyncMock()
mock_session.post = mocker.MagicMock(return_value=MockJsonResponse())
request = RequestFuncInput(
model="single-stage-edit",
model_name="single-stage-edit",
prompt="make it sunny",
api_url="http://test.com/v1/images/edits",
prompt_len=4,
output_len=1,
multi_modal_content=[{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{_MIN_PNG_B64}"}}],
extra_body={"num_inference_steps": 4},
)
output = await async_request_openai_image_edits_omni(request, mock_session, pbar=None)
assert captured_stream == ["false"]
assert output.success is True
assert not output.error
assert output.image_count == 1
assert output.image_generation_time_ms == pytest.approx(12.5)
@pytest.mark.asyncio
async def test_image_edits_stream_true_uses_sse_path(mocker: MockerFixture) -> None:
"""Explicit stream=true keeps the multi-stage SSE client path."""
sse_chunk = (
b'data: {"type":"image","data":[{"b64_json":"' + _MIN_PNG_B64.encode() + b'"}]}\n\n' + b"data: [DONE]\n\n"
)
mock_response = MockResponse(200, [sse_chunk])
captured_stream: list[str] = []
import aiohttp
real_add_field = aiohttp.FormData.add_field
def tracking_add_field(self, name, value=None, **kwargs):
if name == "stream":
captured_stream.append(str(value))
return real_add_field(self, name, value, **kwargs)
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
mock_session = mocker.AsyncMock()
mock_session.post = mocker.MagicMock(return_value=mock_response)
request = RequestFuncInput(
model="multi-stage-edit",
model_name="multi-stage-edit",
prompt="edit",
api_url="http://test.com/v1/images/edits",
prompt_len=2,
output_len=1,
multi_modal_content=[{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{_MIN_PNG_B64}"}}],
extra_body={"stream": True},
)
output = await async_request_openai_image_edits_omni(request, mock_session, pbar=None)
assert captured_stream == ["true"]
assert output.success is True
assert output.image_count == 1
@pytest.mark.asyncio
async def test_image_generations_e2el_includes_json_body_consume(mocker: MockerFixture) -> None:
"""E2EL must include body transfer/decode, not stop at HTTP headers."""
class SlowJsonResponse:
status = 200
def __init__(self):
self._payload = {
"created": 1,
"data": [{"b64_json": _MIN_PNG_B64, "stage_durations": {"stage_0_gen_ms": 1.0}}],
}
async def json(self):
await asyncio.sleep(0.05)
return self._payload
async def text(self):
return json.dumps(self._payload)
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
return None
mock_session = mocker.AsyncMock()
mock_session.post = mocker.MagicMock(return_value=SlowJsonResponse())
request = RequestFuncInput(
model="img-gen",
model_name="img-gen",
prompt="a cat",
api_url="http://test.com/v1/images/generations",
prompt_len=2,
output_len=1,
extra_body={},
)
output = await async_request_openai_image_generations_omni(request, mock_session, pbar=None)
assert output.success is True
assert output.latency >= 0.05
def test_video_local_image_reference_not_forwarded_as_raw_extra_field(tmp_path, mocker: MockerFixture) -> None:
"""Local image_reference must only be uploaded via dedicated serializer."""
import aiohttp
ref_path = tmp_path / "ref.png"
ref_path.write_bytes(base64.b64decode(_MIN_PNG_B64))
captured: list[tuple[str, object]] = []
real_add_field = aiohttp.FormData.add_field
def tracking_add_field(self, name, value=None, **kwargs):
captured.append((str(name), value))
return real_add_field(self, name, value, **kwargs)
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
form = aiohttp.FormData()
extra_body = {
"num_inference_steps": 2,
"image_reference": str(ref_path),
}
request_body = {"model": "vid", "prompt": "p", **extra_body}
_add_video_extra_body_to_form(form, extra_body, request_body)
assert _add_video_reference_to_form(form, extra_body["image_reference"]) is True
field_names = [name for name, _ in captured]
assert "image_reference" not in field_names
assert field_names.count("input_reference") == 1
# Dedicated uploader sends file bytes, not the raw local path string.
uploaded = next(value for name, value in captured if name == "input_reference")
assert uploaded == base64.b64decode(_MIN_PNG_B64)
@pytest.mark.parametrize(
"reference",
[
{"image_url": "https://example.com/ref.png"},
{"file_id": "file-abc"},
[{"image_url": "https://example.com/a.png"}, {"file_id": "file-xyz"}],
],
)
def test_video_structured_image_reference_serialized_to_form(reference: object, mocker: MockerFixture) -> None:
"""Object-form image_reference must be JSON-serialized, not silently dropped."""
import aiohttp
captured: list[tuple[str, object]] = []
real_add_field = aiohttp.FormData.add_field
def tracking_add_field(self, name, value=None, **kwargs):
captured.append((str(name), value))
return real_add_field(self, name, value, **kwargs)
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
form = aiohttp.FormData()
extra_body = {"image_reference": reference}
request_body = {"model": "vid", "prompt": "p", **extra_body}
_add_video_extra_body_to_form(form, extra_body, request_body)
assert _add_video_reference_to_form(form, reference) is True
field_names = [name for name, _ in captured]
# Reserved key must not be double-forwarded as a generic extra field;
# exactly one dedicated image_reference field.
assert field_names.count("image_reference") == 1
assert "input_reference" not in field_names
payload = next(value for name, value in captured if name == "image_reference")
assert json.loads(payload) == reference
def test_video_unsupported_image_reference_raises() -> None:
import aiohttp
form = aiohttp.FormData()
with pytest.raises(ValueError, match="Unsupported image_reference"):
_add_video_reference_to_form(form, {"not_a_supported_key": "x"})
with pytest.raises(ValueError, match="Unsupported image_reference"):
_add_video_reference_to_form(form, "/tmp/does-not-exist-ref.png")
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])
@@ -1,9 +1,11 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# ruff: noqa: E402, I001
import argparse
import math
import os
import sys
import types
from pathlib import Path
import pytest
@@ -106,11 +108,9 @@ def test_resolve_seed_tts_root_downloads_only_requested_locale(monkeypatch, tmp_
captured["allow_patterns"] = allow_patterns
return str(downloaded_root)
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
types.SimpleNamespace(snapshot_download=fake_snapshot_download),
)
from vllm_omni.transformers_utils import repo_utils
monkeypatch.setattr(repo_utils.hf_api(), "snapshot_download", fake_snapshot_download)
resolved = resolve_seed_tts_root(
"zhaochenyang20/seed-tts-eval",
@@ -26,6 +26,7 @@ _MODULE_NAME = "vllm_omni.benchmarks.data_modules.seed_tts_dataset"
if _MODULE_NAME not in sys.modules:
_spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
assert _spec is not None and _spec.loader is not None
_mod = importlib.util.module_from_spec(_spec)
sys.modules[_MODULE_NAME] = _mod
_spec.loader.exec_module(_mod)
@@ -237,7 +238,8 @@ def test_seed_tts_design_dataset_rejects_missing_description(seed_tts_design_roo
assert req.seed_tts_utterance_id == "ok"
def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
@pytest.mark.parametrize("audio_seconds", [0.1, 61])
def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch, audio_seconds):
from vllm_omni.benchmarks.data_modules import seed_tts_eval
calls = {}
@@ -252,8 +254,9 @@ def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
return self
class FakeProcessor:
def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=False):
def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=False, **kwargs):
calls["return_attention_mask"] = return_attention_mask
calls["processor_kwargs"] = kwargs
assert sampling_rate == 16000
assert return_tensors == "pt"
assert len(wav) > 0
@@ -262,9 +265,10 @@ def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
attention_mask=FakeTensor("mask") if return_attention_mask else None,
)
def get_decoder_prompt_ids(self, *, language, task):
def get_decoder_prompt_ids(self, *, language, task, no_timestamps=True):
assert language == "english"
assert task == "transcribe"
calls["no_timestamps"] = no_timestamps
return [(1, 2)]
def batch_decode(self, predicted_ids, *, skip_special_tokens):
@@ -283,10 +287,18 @@ def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
monkeypatch.setattr(seed_tts_eval, "_en_model", FakeModel())
monkeypatch.setattr(seed_tts_eval, "_device", "cuda:1")
text = seed_tts_eval._transcribe_en_f32_16k(np.ones(1600, dtype=np.float32))
text = seed_tts_eval._transcribe_en_f32_16k(np.ones(int(audio_seconds * 16000), dtype=np.float32))
assert text == "hello"
assert calls["return_attention_mask"] is True
assert calls["input_features"].device == "cuda:1"
assert calls["generate_kwargs"]["attention_mask"].device == "cuda:1"
assert calls["generate_kwargs"]["forced_decoder_ids"] == [(1, 2)]
if audio_seconds > 30:
assert calls["processor_kwargs"] == {"truncation": False, "padding": "longest"}
assert calls["generate_kwargs"]["return_timestamps"] is True
assert calls["no_timestamps"] is False
else:
assert calls["processor_kwargs"] == {}
assert "return_timestamps" not in calls["generate_kwargs"]
assert calls["no_timestamps"] is True
+35
View File
@@ -167,6 +167,26 @@ def test_preprocess_serve_args_applies_safe_omniinteract_prompt_default(
{"bot_task": "recaption"},
{"backend", "print_stage", "bot_task"},
),
(
[
"--endpoint",
"/v1/images/edits",
"--image-edits-bot-task",
"think",
],
{"bot_task": "think"},
{"endpoint", "bot_task"},
),
(
[
"--backend",
"/v1/images/edits",
"--image-edits-bot-task",
"recaption",
],
{"bot_task": "recaption"},
{"backend", "bot_task"},
),
(
["--extra-body", '{"bot_task":"vanilla"}'],
{"bot_task": "vanilla"},
@@ -182,6 +202,7 @@ def test_omni_args_parse_and_preprocess(
parser = TrackingArgumentParser()
parser.add_argument("--extra-body", type=json.loads, default=None)
parser.add_argument("--backend", default="openai-chat-omni")
parser.add_argument("--endpoint", default="/v1/chat/completions")
add_omni_args(parser)
args = parser.parse_args(argv)
@@ -358,3 +379,17 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
)
assert bench_requests
assert all(url == expected_url for url in bench_requests), f"Unexpected target URLs: {bench_requests}"
def test_omni_request_timeout_s_flag_defaults_and_parses() -> None:
parser = TrackingArgumentParser()
add_omni_args(parser)
args = parser.parse_args([])
assert args.omni_request_timeout_s is None
args = parser.parse_args(["--omni-request-timeout-s", "0"])
assert args.omni_request_timeout_s == 0.0
args = parser.parse_args(["--omni-request-timeout-s", "3600"])
assert args.omni_request_timeout_s == 3600.0
+46 -8
View File
@@ -6,7 +6,6 @@ Unit tests for StageConfigFactory and related classes.
import importlib
import warnings
from dataclasses import dataclass
from pathlib import Path
from unittest.mock import patch
@@ -37,6 +36,7 @@ from vllm_omni.config.stage_config import (
normalize_pipeline_cli_overrides,
pipeline_cfg_resolver,
)
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -451,13 +451,8 @@ class TestStageConfigFactory:
def test_default_diffusion_with_parallel_config(self):
"""Test diffusion config calculates devices from parallel_config."""
@dataclass
class MockParallelConfig:
world_size: int = 4
kwargs = {
"parallel_config": MockParallelConfig(),
"parallel_config": DiffusionParallelConfig(tensor_parallel_size=4),
"cache_backend": "tea_cache",
}
configs = StageConfigFactory.create_default_diffusion(kwargs)
@@ -937,6 +932,44 @@ class TestPipelineRegistration:
assert isinstance(omni_config, VllmOmniConfig)
assert omni_config.stage_by_id(0).diffusion_config.model == "fake/model"
def test_cosyvoice3_deploy_resolves_hash_snapshot_without_hf_metadata(self, tmp_path):
snapshot = tmp_path / "models--FunAudioLLM--Fun-CosyVoice3-0.5B-2512" / "snapshots" / ("a" * 40)
snapshot.mkdir(parents=True)
(snapshot / "config.json").write_text("{}\n", encoding="utf-8")
model = str(snapshot)
deploy_path = get_deploy_config_path("cosyvoice3.yaml")
assert StageConfigFactory.try_infer_model_type(model, trust_remote_code=False) is None
assert StageConfigFactory.get_pipeline_config(model, trust_remote_code=False) is None
omni_config = StageConfigFactory.create_from_model(
model,
trust_remote_code=False,
cli_overrides={},
deploy_config_path=deploy_path,
)
legacy_configs, _ = StageConfigFactory.create_legacy_stage_configs_from_model(
model,
trust_remote_code=False,
cli_overrides={},
deploy_config_path=deploy_path,
)
assert omni_config is not None
assert omni_config.pipeline_config.model_type == "cosyvoice3"
assert [
(stage.stage_id, stage.stage_type, stage.worker_type, stage.model_stage)
for stage in omni_config.stage_configs
] == [
(0, StageType.LLM, "ar", "cosyvoice3_talker"),
(1, StageType.LLM, "generation", "cosyvoice3_code2wav"),
]
assert all(stage.model_config.model == model for stage in omni_config.stage_configs)
assert all(stage.stage_type is not StageType.DIFFUSION for stage in omni_config.stage_configs)
assert legacy_configs is not None
assert [stage.stage_type for stage in legacy_configs] == [StageType.LLM, StageType.LLM]
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"
@@ -1475,6 +1508,11 @@ stages:
deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml")))
stages = merge_pipeline_deploy(pipeline, deploy)
assert stages[0].yaml_engine_args["model_arch"] == "MiniMaxH3TextEncoder"
assert stages[1].yaml_engine_args["model_arch"] == "MiniMaxH3Pipeline"
assert stages[0].yaml_runtime["num_replicas"] == 1
assert stages[1].yaml_runtime["num_replicas"] == 1
assert stages[1].yaml_engine_args["model_loaded"] == {"text_encoder": False}
assert stages[0].yaml_engine_args["max_num_seqs"] == 1
assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
assert stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
@@ -1597,7 +1635,7 @@ stages:
assert stage.max_model_len == 1024
assert stage.max_num_batched_tokens == 1024
assert stage.max_num_seqs == 1
assert stage.gpu_memory_utilization == 0.7
assert stage.gpu_memory_utilization == 0.8
assert stage.skip_mm_profiling is True
assert stage.enforce_eager is True
assert stage.async_scheduling is False
+2 -2
View File
@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
EnvironmentVariableCategory.PUBLIC_OMNI: 24,
EnvironmentVariableCategory.INHERITED_VLLM: 20,
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
EnvironmentVariableCategory.MODEL_SPECIFIC: 62,
EnvironmentVariableCategory.MODEL_SPECIFIC: 61,
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
EnvironmentVariableCategory.INTERNAL: 2,
}
@@ -179,7 +179,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
)
assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
ModelEnvironmentVariableDisposition.PROMOTE: 37,
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 7,
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
ModelEnvironmentVariableDisposition.INTERNALIZE: 13,
ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5,
@@ -0,0 +1,52 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""The LingBot stepwise deploy config must keep serving on the AR-Diffusion path.
`WS /v1/realtime/video` only reaches the stepwise code when the stage selects the
AR-Diffusion engine and asks for streamed step execution, so those fields are
pinned here rather than left to drift.
"""
from __future__ import annotations
import pytest
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.config.stage_config import load_deploy_config
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
_DEPLOY = "lingbot_world_v2_stepwise.yaml"
def _stage():
deploy = load_deploy_config(get_deploy_config_path(_DEPLOY))
assert len(deploy.stages) == 1, "stepwise serving is a single diffusion stage"
return deploy.stages[0]
def test_stage_selects_the_ar_diffusion_engine_and_pipeline() -> None:
stage = _stage()
assert stage.model_class_name == "LingBotWorldCausalDMDPipeline"
# engine_backend/model_config reach the engine through the passthrough, the
# same way the DreamZero AR-Diffusion deploy config does.
assert stage.engine_extras["engine_backend"].endswith("ar_diffusion.engine.ARDiffusionEngine")
def test_stage_requests_streamed_step_execution_on_one_sequence() -> None:
stage = _stage()
assert stage.step_execution is True
assert stage.engine_extras["streaming_output"] is True
# AR-Diffusion binds one session per runner invocation; batching requests
# would break that contract, so the runner rejects max_num_seqs > 1.
assert stage.max_num_seqs == 1
def test_stage_declares_the_fixed_ar_cache_geometry() -> None:
stage = _stage()
model_config = stage.engine_extras["model_config"]
# The KV cache geometry is fixed at load time, so a served request must ask
# for exactly this resolution.
assert model_config["ar_diffusion_height"] == 480
assert model_config["ar_diffusion_width"] == 832
assert model_config["ar_diffusion_kv_config"]["gpu_memory_fraction"] == 0.6
+127
View File
@@ -261,6 +261,57 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras
VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
@pytest.mark.parametrize("stage_id", [0, 2], ids=["ar", "generation"])
def test_llm_additional_config_roundtrip_and_isolation(stage_id, monkeypatch):
from vllm_omni.engine import stage_init_utils
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
# Worker discovery requires hardware support; this test covers config transport.
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None)
additional_config = {"backend_options": {"enabled": True}}
pipeline = _resolve_pipeline_or_skip("minicpmo_4_5")
deploy = DeployConfig(
stages=[StageDeployConfig(stage_id=i, engine_extras={"additional_config": additional_config}) for i in (0, 2)]
)
config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
stage = config.stage_by_id(stage_id)
assert stage.runtime_config.additional_config == additional_config
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
assert engine_args["additional_config"] == additional_config
engine_args["additional_config"]["backend_options"]["enabled"] = False
assert stage.runtime_config.additional_config["backend_options"]["enabled"] is True
stage.runtime_config.additional_config["backend_options"]["enabled"] = False
assert config.stage_by_id(2 if stage_id == 0 else 0).runtime_config.additional_config == additional_config
assert additional_config["backend_options"]["enabled"] is True
@pytest.mark.parametrize("deploy_name", ["minicpmo_4_5", "minicpmo_4_5_2gpu", "minicpmo_4_5_3gpu"])
def test_minicpmo_npu_additional_config_reaches_engine_args(monkeypatch, deploy_name):
from vllm_omni.engine import stage_init_utils
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
from vllm_omni.platforms import current_omni_platform
monkeypatch.setattr(current_omni_platform, "device_name", "npu")
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None)
stage = _from_pipeline_key("minicpmo_4_5", deploy_config_path=deploy_name).stage_by_id(2)
expected = {"code2wav_enable_npu_graph": True, "code2wav_max_npu_graphs": 32}
assert stage.runtime_config.additional_config == expected
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
assert engine_args["additional_config"] == expected
def test_diffusion_additional_config_keeps_diffusion_owner():
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
additional_config = {"torchair_graph_config": {"enabled": True}}
stage = _from_pipeline_key("dreamzero", cli_overrides={"additional_config": additional_config}).stage_by_id(0)
assert stage.runtime_config.additional_config is None
assert stage.diffusion_config.additional_config == additional_config
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
assert engine_args["additional_config"] == additional_config
@pytest.mark.parametrize(
("engine_extras", "cli_overrides"),
[
@@ -532,6 +583,17 @@ def test_from_pipeline_config_dispatches_async_chunk_processors_without_mutating
assert pipeline.get_stage(1).custom_process_input_func is None
def test_joyai_code2wav_waits_for_full_payload():
config = _from_pipeline_key("joyai_vl_interaction")
talker = config.stage_by_id(1)
code2wav = config.stage_by_id(2)
assert talker.custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
assert code2wav.custom_process_input_func.endswith("talker2code2wav_token_only")
assert code2wav.connector_config.async_chunk is False
assert code2wav.model_config.requires_full_payload_input is True
def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
assert not hasattr(BaseVllmOmniStageConfig, "from_stage_config")
assert not hasattr(BaseVllmOmniStageConfig, "to_legacy_stage_config")
@@ -559,6 +621,7 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
def test_runtime_config_fields_match_structured_runtime_scope():
assert {f.name for f in fields(OmniStageRuntimeConfig)} == {
"additional_config",
"distributed_executor_backend",
"worker_cls",
"devices",
@@ -1253,6 +1316,70 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
assert cfg.diffusers_call_kwargs == {}
def test_diffusion_config_none_values_preserve_dataclass_defaults():
from vllm_omni.diffusion.data import OmniDiffusionConfig
normalized = OmniDiffusionConfig.normalize_init_kwargs(
{
"lora_scale": None,
"enable_sleep_mode": None,
"diffusers_load_kwargs": None,
}
)
assert "lora_scale" not in normalized
assert "enable_sleep_mode" not in normalized
assert normalized["diffusers_load_kwargs"] == {}
config = OmniDiffusionConfig.from_kwargs(
lora_scale=None,
enable_sleep_mode=None,
)
assert config.lora_scale == 1.0
assert config.enable_sleep_mode is False
@pytest.mark.parametrize(
("canonical_key", "alias_key", "canonical_value", "alias_value"),
[
("lora_scale", "static_lora_scale", 0.75, 0.25),
(
"quantization_config",
"diffusion_quantization_config",
{"method": "canonical"},
{"method": "diffusion-alias"},
),
(
"quantization_config",
"quantization",
{"method": "canonical"},
"legacy-alias",
),
("diffusion_kv_cache_dtype", "kv_cache_dtype", "fp8", "fp16"),
("diffusion_kv_cache_skip_steps", "kv_cache_skip_steps", "0-1", "2-3"),
("diffusion_kv_cache_skip_layers", "kv_cache_skip_layers", "1-2", "3-4"),
("streaming_output", "diffusion_streaming_output", False, True),
],
)
def test_diffusion_alias_conflicts_prefer_canonical_key(
canonical_key,
alias_key,
canonical_value,
alias_value,
):
from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs
normalized = normalize_omni_diffusion_kwargs(
{
canonical_key: canonical_value,
alias_key: alias_value,
}
)
assert normalized[canonical_key] == canonical_value
assert alias_key not in normalized
def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path, monkeypatch):
from vllm_omni.platforms import current_omni_platform
@@ -19,6 +19,9 @@ from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.outputs import ModelRunnerOutput
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
OmniChunkTransferAdapter,
)
# isort: on
@@ -62,6 +65,43 @@ def _make_update(prompt_token_ids: list[int] | None = None) -> StreamingUpdate:
)
def _make_talker_adapter(
*,
max_model_len: int = 100,
recompute_on_capacity: bool = True,
) -> OmniChunkTransferAdapter:
adapter = OmniChunkTransferAdapter.__new__(OmniChunkTransferAdapter)
adapter.receives_chunks = False
adapter._max_model_len = max_model_len
adapter._streaming_prompt_previous_chunks = 1
adapter._streaming_prompt_recompute_on_capacity = recompute_on_capacity
adapter._streaming_condition_lengths = {}
adapter._streaming_condition_seqs = {}
adapter.segment_finished_requests = set()
adapter.requests_num_chunks_sent = {}
return adapter
def _make_talker_update(
condition_len: int,
*,
reserve: int = 10,
condition_seq: int | None,
) -> StreamingUpdate:
update = _make_update([0] * condition_len)
update.model_intermediate_buffer = {
"native_duplex": True,
"ids": {"prompt": [1]},
"meta": {
"next_stage_prompt_len": condition_len,
"next_stage_generation_tokens": reserve,
},
}
if condition_seq is not None:
update.model_intermediate_buffer["meta"]["streaming_condition_seq"] = condition_seq
return update
def _run_resumable_segment_stop(
session: Request,
*,
@@ -538,6 +578,288 @@ def test_explicit_model_intermediate_prompt_replacement_releases_cache_and_water
sched.encoder_cache_manager.free.assert_called_once_with(session)
def test_talker_capacity_exact_fit_extends_from_declared_length_without_ids_prompt() -> None:
sched = _make_scheduler(stage_id=1)
adapter = _make_talker_adapter(max_model_len=100)
sched.chunk_transfer_adapter = adapter
session = _make_request()
session.external_req_id = "external-capacity-exact-fit"
session.prompt_token_ids = [0] * 70
session._all_token_ids.clear()
session._all_token_ids.extend(session.prompt_token_ids)
session._output_token_ids.clear()
session.append_output_token_ids([7, 8, 9, 10])
session.num_prompt_tokens = 70
session.num_computed_tokens = 74
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
sched.num_waiting_for_streaming_input = 1
adapter._streaming_condition_lengths[session.request_id] = 20
adapter._streaming_condition_seqs[session.request_id] = 0
update = _make_talker_update(16, condition_seq=1)
update.model_intermediate_buffer["ids"] = {"tts": [1, 2, 3]}
sched._update_request_as_session(session, update)
assert session.num_computed_tokens == 74
assert session.num_prompt_tokens == 90
assert list(session._all_token_ids[-20:]) == [7, 8, 9, 10] + [0] * 16
assert update.model_intermediate_buffer["meta"] == {
"next_stage_prompt_len": 16,
"next_stage_generation_tokens": 10,
"streaming_condition_seq": 1,
"streaming_prompt_recompute": False,
}
assert adapter._streaming_condition_lengths[session.request_id] == 16
assert adapter._streaming_condition_seqs[session.request_id] == 1
assert session._omni_segment_generation == 1
assert session.status == RequestStatus.WAITING
assert sched.num_waiting_for_streaming_input == 0
sched._free_request_blocks.assert_not_called()
sched.encoder_cache_manager.free.assert_not_called()
def test_talker_first_update_seeds_sender_condition_tracking() -> None:
sched = _make_scheduler(stage_id=1)
adapter = _make_talker_adapter(max_model_len=100)
sched.chunk_transfer_adapter = adapter
session = _make_request()
session.external_req_id = "external-capacity-first-update"
session.prompt_token_ids = [0] * 20
session._all_token_ids.clear()
session._all_token_ids.extend(session.prompt_token_ids)
session._output_token_ids.clear()
session.append_output_token_ids([7, 8, 9])
session.num_prompt_tokens = 20
session.num_computed_tokens = 23
session.model_intermediate_buffer = {
"native_duplex": True,
"meta": {
"next_stage_prompt_len": 20,
"next_stage_generation_tokens": 10,
"streaming_condition_seq": 0,
},
}
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
sched.num_waiting_for_streaming_input = 1
update = _make_talker_update(16, condition_seq=1)
sched._update_request_as_session(session, update)
assert update.model_intermediate_buffer["meta"]["streaming_condition_seq"] == 1
assert update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is False
assert adapter._streaming_condition_lengths[session.request_id] == 16
assert adapter._streaming_condition_seqs[session.request_id] == 1
assert session.num_computed_tokens == 23
assert session.num_prompt_tokens == 39
sched._free_request_blocks.assert_not_called()
sched.encoder_cache_manager.free.assert_not_called()
def test_talker_first_sliding_update_recomputes_from_seeded_condition() -> None:
sched = _make_scheduler(stage_id=1)
adapter = _make_talker_adapter(recompute_on_capacity=False)
sched.chunk_transfer_adapter = adapter
session = _make_request()
session.external_req_id = "external-sliding-first-update"
session.prompt_token_ids = [0] * 20
session._all_token_ids.clear()
session._all_token_ids.extend(session.prompt_token_ids)
session._output_token_ids.clear()
previous_codes = [7, 8, 9]
session.append_output_token_ids(previous_codes)
session.num_prompt_tokens = 20
session.num_computed_tokens = 23
session.model_intermediate_buffer = {
"native_duplex": True,
"meta": {
"next_stage_prompt_len": 20,
"next_stage_generation_tokens": 10,
"streaming_condition_seq": 0,
},
}
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
sched.num_waiting_for_streaming_input = 1
update = _make_talker_update(16, condition_seq=1)
sched._update_request_as_session(session, update)
assert session.num_computed_tokens == 0
assert session.num_prompt_tokens == 39
assert update.model_intermediate_buffer["ids"]["streaming_prompt_previous_codes"] == previous_codes
assert update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is True
assert adapter._streaming_condition_lengths[session.request_id] == 16
assert adapter._streaming_condition_seqs[session.request_id] == 1
sched._free_request_blocks.assert_called_once_with(session)
sched.encoder_cache_manager.free.assert_called_once_with(session)
@pytest.mark.parametrize(
("condition_seq", "error"),
[
(None, "missing streaming_condition_seq"),
(0, "expected=1, received=0"),
(2, "expected=1, received=2"),
],
)
def test_talker_invalid_condition_sequence_does_not_advance_tracking(mocker, condition_seq, error) -> None:
sched = _make_scheduler(stage_id=1)
adapter = _make_talker_adapter()
adapter.record_receive_failure = mocker.MagicMock()
adapter._streaming_condition_lengths["req-ar-streaming-test"] = 20
adapter._streaming_condition_seqs["req-ar-streaming-test"] = 0
sched.chunk_transfer_adapter = adapter
session = _make_request()
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
sched.num_waiting_for_streaming_input = 1
update = _make_talker_update(16, condition_seq=condition_seq)
sched._update_request_as_session(session, update)
failure = adapter.record_receive_failure.call_args.args
assert failure[0] == session.request_id
assert error in failure[1]
assert adapter._streaming_condition_lengths[session.request_id] == 20
assert adapter._streaming_condition_seqs[session.request_id] == 0
assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
def test_talker_first_update_that_cannot_fit_window_is_request_local(mocker) -> None:
sched = _make_scheduler(stage_id=1)
adapter = _make_talker_adapter()
adapter.record_receive_failure = mocker.MagicMock()
sched.chunk_transfer_adapter = adapter
session = _make_request()
session.prompt_token_ids = [0] * 20
session._all_token_ids.clear()
session._all_token_ids.extend(session.prompt_token_ids)
session._output_token_ids.clear()
session.append_output_token_ids(list(range(9)))
session.num_prompt_tokens = 20
session.num_computed_tokens = 29
session.model_intermediate_buffer = {
"native_duplex": True,
"meta": {
"next_stage_prompt_len": 20,
"next_stage_generation_tokens": 10,
"streaming_condition_seq": 0,
},
}
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
sched.num_waiting_for_streaming_input = 1
update = _make_talker_update(62, condition_seq=1)
sched._update_request_as_session(session, update)
failure = adapter.record_receive_failure.call_args.args
assert failure[0] == session.request_id
assert "sliding streaming prompt plus generation reserve exceeds max_model_len" in failure[1]
assert adapter._streaming_condition_lengths == {}
assert adapter._streaming_condition_seqs == {}
assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
def test_talker_capacity_overflow_recomputes_then_resumes_accumulation() -> None:
sched = _make_scheduler(stage_id=1)
adapter = _make_talker_adapter(max_model_len=100)
sched.chunk_transfer_adapter = adapter
session = _make_request()
session.external_req_id = "external-capacity-rollover"
session.prompt_token_ids = [0] * 70
session._all_token_ids.clear()
session._all_token_ids.extend(session.prompt_token_ids)
session._output_token_ids.clear()
previous_codes = [7, 8, 9, 10]
session.append_output_token_ids(previous_codes)
session.num_prompt_tokens = 70
session.num_computed_tokens = 74
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
sched.num_waiting_for_streaming_input = 1
adapter._streaming_condition_lengths[session.request_id] = 20
adapter._streaming_condition_seqs[session.request_id] = 0
adapter.segment_finished_requests.add(session.request_id)
adapter.requests_num_chunks_sent[session.external_req_id] = 74
overflow_update = _make_talker_update(17, condition_seq=1)
sched._update_request_as_session(session, overflow_update)
assert session.num_computed_tokens == 0
assert session.num_prompt_tokens == 41
assert session.prompt_token_ids == [0] * 41
assert overflow_update.model_intermediate_buffer["ids"]["streaming_prompt_previous_codes"] == previous_codes
assert overflow_update.model_intermediate_buffer["meta"]["streaming_condition_seq"] == 1
assert overflow_update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is True
assert adapter.segment_finished_requests == set()
assert adapter.requests_num_chunks_sent == {}
assert session._omni_segment_generation == 1
sched._free_request_blocks.assert_called_once_with(session)
sched.encoder_cache_manager.free.assert_called_once_with(session)
# Once the replacement has been recomputed, later conditions append to
# the new full-attention prefix until capacity is approached again.
next_codes = [21, 22, 23, 24, 25, 26]
session.append_output_token_ids(next_codes)
session.num_computed_tokens = 47
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
sched.num_waiting_for_streaming_input = 1
append_update = _make_talker_update(20, condition_seq=2)
sched._update_request_as_session(session, append_update)
assert session.num_computed_tokens == 47
assert session.num_prompt_tokens == 67
assert list(session._all_token_ids[-26:]) == next_codes + [0] * 20
assert append_update.model_intermediate_buffer["meta"]["streaming_condition_seq"] == 2
assert append_update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is False
assert session._omni_segment_generation == 2
assert sched._free_request_blocks.call_count == 1
assert sched.encoder_cache_manager.free.call_count == 1
# A later overflow replaces the accumulated prefix again rather than
# reverting to recompute-on-every-condition behavior.
second_rollover_codes = [31, 32, 33, 34, 35, 36]
session.num_computed_tokens = session.num_prompt_tokens
session.append_output_token_ids(second_rollover_codes)
session.num_computed_tokens += len(second_rollover_codes)
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
sched.num_waiting_for_streaming_input = 1
second_overflow_update = _make_talker_update(18, condition_seq=3)
sched._update_request_as_session(session, second_overflow_update)
assert session.num_computed_tokens == 0
assert session.num_prompt_tokens == 44
assert second_overflow_update.model_intermediate_buffer["ids"]["streaming_prompt_previous_codes"] == (
second_rollover_codes
)
assert second_overflow_update.model_intermediate_buffer["meta"]["streaming_condition_seq"] == 3
assert second_overflow_update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is True
assert session._omni_segment_generation == 3
assert sched._free_request_blocks.call_count == 2
assert sched.encoder_cache_manager.free.call_count == 2
def test_talker_invalid_capacity_update_is_reported_per_request(mocker) -> None:
sched = _make_scheduler(stage_id=1)
adapter = _make_talker_adapter(max_model_len=100)
adapter.record_receive_failure = mocker.MagicMock()
sched.chunk_transfer_adapter = adapter
session = _make_request()
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
sched.num_waiting_for_streaming_input = 1
update = _make_talker_update(91, condition_seq=0)
sched._update_request_as_session(session, update)
adapter.record_receive_failure.assert_called_once_with(
session.request_id,
"fresh streaming prompt plus generation reserve exceeds max_model_len: prompt=91, reserve=10, limit=100",
)
assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
assert sched.num_waiting_for_streaming_input == 1
sched._free_request_blocks.assert_not_called()
sched.encoder_cache_manager.free.assert_not_called()
def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None:
sched = _make_scheduler(stage_id=1)
session = _make_request()
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Performance benchmark CI runner for diffusion models.
@@ -356,6 +359,8 @@ def _resolve_offline_model(model: str) -> str:
"""
import huggingface_hub
from vllm_omni.transformers_utils.repo_utils import hf_api
if not model or os.path.isdir(model):
return model
@@ -380,9 +385,7 @@ def _resolve_offline_model(model: str) -> str:
if len(parts) >= 3:
repo_id = "/".join(parts[:2])
subfolder = "/".join(parts[2:])
from huggingface_hub import snapshot_download
snapshot_root = snapshot_download(
snapshot_root = hf_api().snapshot_download(
repo_id,
allow_patterns=[f"{subfolder}/**"],
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
@@ -391,9 +394,7 @@ def _resolve_offline_model(model: str) -> str:
if not huggingface_hub.constants.HF_HUB_OFFLINE:
return model
from huggingface_hub import snapshot_download
return snapshot_download(model, local_files_only=True)
return hf_api().snapshot_download(model, local_files_only=True)
class DiffusionServer:
@@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU contracts for capability-driven AR-Diffusion sessions."""
from __future__ import annotations
@@ -10,6 +11,7 @@ from types import SimpleNamespace
import pytest
import torch
from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.experimental.ar_diffusion.capability import (
ARDiffusionCrossAttentionKVSpec,
@@ -27,6 +29,22 @@ NEG = "negative"
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.fixture(autouse=True)
def platform_synchronize_calls(monkeypatch) -> list[None]:
"""Stub the platform barrier the runner takes at every chunk boundary.
These tests run on a CPU host, where the resolved platform is
``UnspecifiedOmniPlatform`` and ``synchronize()`` raises. Tests that need
a failing barrier override this with their own fake.
"""
calls: list[None] = []
monkeypatch.setattr(
"vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
SimpleNamespace(synchronize=lambda: calls.append(None)),
)
return calls
def lingbot_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
"""Single-kv_branch causal DMD: 3 latent frames/block + sink/window + text KV."""
return ARDiffusionKVCacheSpec(
@@ -114,6 +132,43 @@ class BatchCapablePipeline(CapablePipeline):
supports_request_batch = True
class StepCapablePipeline(CapablePipeline):
supports_step_execution = True
def prepare_encode(self, state, **kwargs):
del kwargs
return state
def denoise_step(self, input_batch, *, states=None, **kwargs):
del input_batch, states, kwargs
return None
def step_scheduler(self, state, noise_pred, **kwargs):
del state, noise_pred, kwargs
def post_decode(self, state, **kwargs):
del state, kwargs
from vllm_omni.diffusion.data import DiffusionOutput
return DiffusionOutput()
def scheduler_output(
request_ids: tuple[str, ...] = ("req-1",),
*,
finished: tuple[str, ...] = (),
) -> DiffusionSchedulerOutput:
"""Build the real scheduler payload; stepwise requests arrive as cached rows."""
return DiffusionSchedulerOutput(
step_id=0,
scheduled_new_reqs=[],
scheduled_cached_reqs=CachedRequestData(request_ids=list(request_ids)),
finished_req_ids=set(finished),
num_running_reqs=len(request_ids),
num_waiting_reqs=0,
)
def make_runner(
pipeline: object,
*,
@@ -140,6 +195,7 @@ def make_runner(
runner._sessions = OrderedDict()
runner._session_capacity = 0
runner._perf_e2e_times = []
runner._stepwise_chunk_started = {}
runner._preallocate_kv_cache(available_bytes=available_bytes)
return runner
@@ -395,12 +451,16 @@ def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
ctx.ensure_video_slots(torch.device("cpu"))
return object()
def synchronize_boom(device):
def synchronize_boom():
raise RuntimeError("asynchronous kernel failed")
monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom)
runner.device = torch.device("cuda")
# The runner defers the device barrier to the platform, so the failure is
# injected there rather than on a device-specific torch entry point.
monkeypatch.setattr(
"vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
SimpleNamespace(synchronize=synchronize_boom),
)
request = SimpleNamespace(
request_id="broken-request",
sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
@@ -424,6 +484,11 @@ def test_ar_runner_rejects_step_and_request_batch_modes():
make_runner(BatchCapablePipeline(lingbot_like_spec()))
def test_ar_runner_allows_step_execution_when_pipeline_implements_the_contract():
runner = make_runner(StepCapablePipeline(lingbot_like_spec()), step_execution=True)
assert runner.kv_cache is not None
def test_ar_runner_defensively_rejects_inherited_batch_and_step_entrypoints():
runner = object.__new__(ARDiffusionModelRunner)
with pytest.raises(RuntimeError, match="request-batch execution"):
@@ -432,6 +497,144 @@ def test_ar_runner_defensively_rejects_inherited_batch_and_step_entrypoints():
runner.execute_stepwise(None)
def test_execute_stepwise_binds_request_id_session_and_unbinds_after_call(monkeypatch):
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
pipeline = StepCapablePipeline(lingbot_like_spec())
runner = make_runner(pipeline, step_execution=True)
bound_during = []
def fake_execute(self, scheduler_output):
del self, scheduler_output
bound_during.append(pipeline.bound_state is not None)
return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False)])
monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
output = runner.execute_stepwise(scheduler_output())
assert bound_during == [True]
assert pipeline.bound_state is None
assert pipeline.binds == ["req-1"]
assert "req-1" in runner._sessions
assert output.get_request_output("req-1") is not None
assert not pipeline.closes
def test_execute_stepwise_closes_session_when_request_finishes(monkeypatch):
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
pipeline = StepCapablePipeline(lingbot_like_spec())
runner = make_runner(pipeline, step_execution=True)
def fake_execute(self, scheduler_output):
del self, scheduler_output
return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=True)])
monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
runner.execute_stepwise(scheduler_output())
assert pipeline.bound_state is None
assert pipeline.closes == ["req-1"]
assert "req-1" not in runner._sessions
def test_execute_stepwise_exception_releases_kv_and_does_not_resume_pages(monkeypatch):
pipeline = StepCapablePipeline(lingbot_like_spec())
runner = make_runner(pipeline, step_execution=True)
first = runner._get_or_create_session("req-1")
def boom(self, scheduler_output):
del self, scheduler_output
raise RuntimeError("step failed")
monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", boom)
with pytest.raises(RuntimeError, match="step failed"):
runner.execute_stepwise(scheduler_output())
assert pipeline.bound_state is None
assert pipeline.closes == ["req-1"]
assert "req-1" not in runner._sessions
second = runner._get_or_create_session("req-1")
assert second is not first
def test_execute_stepwise_closes_session_for_scheduler_aborted_request(monkeypatch):
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput
pipeline = StepCapablePipeline(lingbot_like_spec())
runner = make_runner(pipeline, step_execution=True)
runner._get_or_create_session("req-1")
def fake_execute(self, scheduler_output):
del self, scheduler_output
return BatchRunnerOutput.from_list([])
monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
# An aborted request is never scheduled again and never reports finished,
# so the retire path is the only thing that can free its KV.
runner.execute_stepwise(scheduler_output(request_ids=(), finished=("req-1",)))
assert "req-1" not in runner._sessions
assert pipeline.closes == ["req-1"]
def test_execute_stepwise_times_once_per_chunk_not_once_per_step(monkeypatch, platform_synchronize_calls):
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
pipeline = StepCapablePipeline(lingbot_like_spec())
runner = make_runner(pipeline, step_execution=True)
results = [None, None, None, DiffusionOutput()]
def fake_execute(self, scheduler_output):
del self, scheduler_output
return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False, result=results.pop(0))])
monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
for _ in range(4):
runner.execute_stepwise(scheduler_output())
# Four denoise steps, one emitted AR block: one timing sample, matching the
# one-sample-per-block meaning request mode already has.
assert len(runner._perf_e2e_times) == 1
assert len(platform_synchronize_calls) == 1
assert "req-1" not in runner._stepwise_chunk_started
def test_execute_stepwise_synchronize_exception_fails_closed(monkeypatch):
"""The chunk barrier is where an asynchronous device error surfaces. It must
clean up like an in-step failure instead of leaving the session, the cached
step state and the timing entry to the scheduler's next cycle."""
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
pipeline = StepCapablePipeline(lingbot_like_spec())
runner = make_runner(pipeline, step_execution=True)
runner.state_cache = {"req-1": object()}
def fake_execute(self, scheduler_output):
del self, scheduler_output
return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False, result=object())])
def synchronize_boom():
raise RuntimeError("asynchronous kernel failed")
monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
monkeypatch.setattr(
"vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
SimpleNamespace(synchronize=synchronize_boom),
)
with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
runner.execute_stepwise(scheduler_output())
assert pipeline.bound_state is None
assert pipeline.closes == ["req-1"]
assert "req-1" not in runner._sessions
assert "req-1" not in runner.state_cache
assert "req-1" not in runner._stepwise_chunk_started
assert not runner._perf_e2e_times
def test_model_specific_warmup_provider_is_consumed(monkeypatch):
requests = [object(), object()]
pipeline = WarmupPipeline(lingbot_like_spec(), requests)
+145 -7
View File
@@ -25,12 +25,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
)
from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
BLOCK = 16
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
@@ -155,14 +152,155 @@ def test_paged_pool_layout_exposes_flat_slot_views():
dtype=torch.float32,
device=torch.device("cpu"),
)
assert kv_pools[0].shape == (2, 4, BLOCK, 4, 64)
key_cache, value_cache = kv_pools[0]
assert key_cache.shape == value_cache.shape == (4, BLOCK, 4, 64)
assert k_pools[0].shape == (4 * BLOCK, 4, 64)
assert v_pools[0].shape == (4 * BLOCK, 4, 64)
# The flat views must alias their block-shaped cache, so a slot write is
# visible through the layout the attention kernel reads.
k_pools[0][BLOCK + 3].fill_(7)
v_pools[0][2 * BLOCK + 5].fill_(11)
assert torch.equal(kv_pools[0][0, 1, 3], k_pools[0][BLOCK + 3])
assert torch.equal(kv_pools[0][1, 2, 5], v_pools[0][2 * BLOCK + 5])
assert torch.equal(key_cache[1, 3], k_pools[0][BLOCK + 3])
assert torch.equal(value_cache[2, 5], v_pools[0][2 * BLOCK + 5])
def test_key_and_value_caches_do_not_share_storage():
"""K and V must be separate allocations, not halves of one tensor.
The paged-write custom op declares both as mutated. When they alias one
storage, inductor's reinplace pass can re-inplace only the first of the
two, and auto_functionalized_v2's clone of the entire second pool survives
into the compiled graph -- one full pool copied per compiled region per
denoising step.
"""
kv_pools, k_pools, v_pools = allocate_kv_pool_with_views(
num_blocks=4,
block_size=BLOCK,
num_layers=2,
num_kv_heads=4,
head_dim=64,
dtype=torch.float32,
device=torch.device("cpu"),
)
seen = set()
for layer, (key_cache, value_cache) in enumerate(kv_pools):
k_storage = key_cache.untyped_storage().data_ptr()
v_storage = value_cache.untyped_storage().data_ptr()
assert k_storage != v_storage, f"layer {layer}: K and V share one allocation"
assert k_pools[layer].untyped_storage().data_ptr() == k_storage
assert v_pools[layer].untyped_storage().data_ptr() == v_storage
seen.update((k_storage, v_storage))
# Every layer's K and V are distinct allocations as well.
assert len(seen) == 2 * len(kv_pools)
# Negative control: writing V must not disturb K.
k_pools[0].fill_(0)
v_pools[0].fill_(5)
assert torch.equal(k_pools[0], torch.zeros_like(k_pools[0]))
def _pool_sized_allocations(code: str, pool_numel: int) -> list[int]:
"""Sizes of generated allocations at least as large as one pool."""
import re
sizes = []
for shape in re.findall(r"empty_strided_\w+\(\(([\d, ]*?)\)", code):
dims = [int(part) for part in shape.replace(" ", "").strip(",").split(",") if part]
numel = 1
for dim_size in dims:
numel *= dim_size
if dims and numel >= pool_numel:
sizes.append(numel)
return sizes
def _generated_code_for_pools(key_cache, value_cache):
"""Compile a paged-write-shaped op over these pools, return the generated code."""
from torch._inductor import config as inductor_config
from torch._inductor.utils import run_and_get_code
class _Holder(torch.nn.Module):
def __init__(self):
super().__init__()
self.register_buffer("key_cache", key_cache)
self.register_buffer("value_cache", value_cache)
def forward(self, value):
return torch.ops.test_kv_pool_clone.write(self.key_cache, self.value_cache, value)
torch._dynamo.reset()
# Inductor caches compiled code across ``torch._dynamo.reset()``, so the
# second call in a process would otherwise return the first one's code and
# both branches of this test would describe the same graph.
with inductor_config.patch(force_disable_caches=True):
_, codes = run_and_get_code(torch.compile(_Holder(), backend="inductor", fullgraph=True), torch.ones(1))
return "\n".join(codes)
def test_compiled_paged_write_does_not_clone_the_pool():
"""The compiled graph must write the pools in place, not through a clone.
``test_key_and_value_caches_do_not_share_storage`` pins the mechanism this
change introduces -- two allocations rather than one. This pins the effect
that mechanism buys, which is the reason the change exists: with K and V as
halves of one tensor, inductor's reinplace pass re-inplaces only the first,
and ``auto_functionalized_v2``'s clone of the pool survives into the
generated code as a pool-sized allocation plus a copy, once per compiled
region per denoising step.
The shared layout is compiled here as a positive control. Without it, a
codegen change that stops emitting ``empty_strided_*`` would leave the
detector matching nothing and this test passing for the wrong reason.
"""
if "write" not in dir(getattr(torch.ops, "test_kv_pool_clone", object())):
@torch.library.custom_op("test_kv_pool_clone::write", mutates_args=("key_pool", "value_pool"))
def _write(key_pool: torch.Tensor, value_pool: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
key_pool[0, 0, 0, 0] = value[0]
value_pool[0, 0, 0, 0] = value[0]
return value * 2
@_write.register_fake
def _(key_pool, value_pool, value):
return torch.empty_like(value)
# The stand-in models the real op's contract; that declaration is what
# drives the reinplace pass, so a drift here would silently test nothing.
real_schema = torch.ops.vllm_omni.ar_diffusion_paged_write_attn.default._schema
real_mutated = [arg.name for arg in real_schema.arguments if arg.alias_info and arg.alias_info.is_write]
assert real_mutated == ["key_pool", "value_pool"], (
f"the real paged-write op now mutates {real_mutated}; this test still models two pools"
)
num_blocks, heads, dim = 8, 4, 16
cache_shape = (num_blocks, BLOCK, heads, dim)
pool_numel = num_blocks * BLOCK * heads * dim
# Positive control: one allocation, K and V aliasing it.
shared = torch.empty(2, *cache_shape)
shared_allocations = _pool_sized_allocations(_generated_code_for_pools(shared[0], shared[1]), pool_numel)
assert shared_allocations, (
"the shared layout produced no pool-sized allocation, so this test can no "
"longer tell the two layouts apart -- inductor's codegen or the detector changed"
)
# The allocator under test.
kv_pools, _, _ = allocate_kv_pool_with_views(
num_blocks=num_blocks,
block_size=BLOCK,
num_layers=1,
num_kv_heads=heads,
head_dim=dim,
dtype=torch.float32,
device=torch.device("cpu"),
)
key_cache, value_cache = kv_pools[0]
separate_allocations = _pool_sized_allocations(_generated_code_for_pools(key_cache, value_cache), pool_numel)
assert not separate_allocations, (
"the compiled graph allocates a pool-sized buffer, so the pool is being cloned "
f"rather than written in place: {separate_allocations} (control saw {shared_allocations})"
)
def test_build_manager_allocate_free_roundtrip():
@@ -0,0 +1,193 @@
# SPDX-License-Identifier: Apache-2.0
"""CUDA-only regression test for the AR-Diffusion paged KV write op.
Split out of ``test_kv_cache.py``: that module is swept into CI's
``core_model and cpu`` job via a module-level ``pytest.mark.cpu``, and
``hardware_test(..., num_cards=1)`` adds no ``skipif`` guard for a missing
GPU, so this test was being collected (and would fail to import a CUDA
device) on CPU-only runners. Keeping it in its own, non-``cpu`` module means
it is only ever collected by a job that explicitly selects ``cuda`` for this
path -- none does yet, so for now this is intentionally not wired into any
CI job.
"""
import pytest
import torch
from tests.helpers.mark import hardware_test
from vllm_omni.experimental.ar_diffusion.kv_cache import allocate_kv_pool_with_views
BLOCK = 16
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_real_paged_write_op_does_not_clone_the_pool_on_cuda():
"""Same assertion as the CPU test, against the real op's schema and CUDA codegen.
The CPU test compiles a stand-in with the same ``mutates_args`` contract,
because the real op needs a device and an attention kernel. That leaves
one thing unverified: whether Triton codegen reinplaces the same way,
which is the path the reported regression was observed on.
This does not call the real op directly: ``ar_diffusion_paged_write_attn``
runs FlashAttention's paged ``flash_attn_varlen_func`` internally, and
compiling that under ``fullgraph=True`` can fail for reasons that have
nothing to do with cloning (a missing kernel for this toy shape, a graph
break, or an autotune illegal-memory-access on the control pool). Instead
this registers a second stand-in that copies the real op's full argument
list and ``mutates_args`` declaration -- so reinplace sees the identical
contract the real op presents -- but skips straight to the pool writes,
the same way the real op's own implementation does before it calls
attention. That isolates "does Inductor clone this schema on CUDA" from
"does FlashAttention compile", which is a separate, already-covered
concern.
Note the pools passed here are the flat views, since that is what the op
takes -- it indexes them by slot and unflattens internally.
"""
import re
from torch._inductor import config as inductor_config
from torch._inductor.utils import run_and_get_code
device = torch.device("cuda")
num_blocks, heads, dim, num_tokens = 8, 4, 64, 8
pool_numel = num_blocks * BLOCK * heads * dim
# The size a genuine clone would produce: K and V pool-sized allocations,
# coalesced into one buffer the way `auto_functionalized_v2` was cloning
# them before this PR's fix. Matching this exact size (rather than "any
# allocation at least as large as one pool") avoids false positives from
# unrelated workspace/scratch buffers Inductor may emit in that range.
clone_numel = 2 * pool_numel
if "write" not in dir(getattr(torch.ops, "test_kv_pool_clone_schema", object())):
@torch.library.custom_op("test_kv_pool_clone_schema::write", mutates_args=("key_pool", "value_pool"))
def _write(
query: torch.Tensor,
k_curr: torch.Tensor,
v_curr: torch.Tensor,
k_act: torch.Tensor | None,
v_act: torch.Tensor | None,
key_pool: torch.Tensor,
value_pool: torch.Tensor,
block_size: int,
video_slots: torch.Tensor,
action_slots: torch.Tensor,
block_table: torch.Tensor,
query_start_loc: torch.Tensor,
seq_lens: torch.Tensor,
max_query_len: int,
max_seq_len: int,
softmax_scale: float,
) -> torch.Tensor:
# Same pool-write side effect as `_paged_write_attn_impl`, minus
# the `ar_diffusion_paged_attention` call -- that is the part
# this test is not trying to exercise.
key_pool[video_slots] = k_curr.to(key_pool.dtype)
value_pool[video_slots] = v_curr.to(value_pool.dtype)
if k_act is not None and v_act is not None and k_act.shape[0] > 0:
key_pool[action_slots] = k_act.to(key_pool.dtype)
value_pool[action_slots] = v_act.to(value_pool.dtype)
return query * 2
@_write.register_fake
def _(
query,
k_curr,
v_curr,
k_act,
v_act,
key_pool,
value_pool,
block_size,
video_slots,
action_slots,
block_table,
query_start_loc,
seq_lens,
max_query_len,
max_seq_len,
softmax_scale,
):
return torch.empty_like(query)
class _Holder(torch.nn.Module):
def __init__(self, key_flat, value_flat):
super().__init__()
self.register_buffer("key_flat", key_flat)
self.register_buffer("value_flat", value_flat)
def forward(self, query, k_curr, v_curr, video_slots, action_slots, block_table, query_start_loc, seq_lens):
return torch.ops.test_kv_pool_clone_schema.write(
query,
k_curr,
v_curr,
None,
None,
self.key_flat,
self.value_flat,
BLOCK,
video_slots,
action_slots,
block_table,
query_start_loc,
seq_lens,
num_tokens,
num_tokens,
dim**-0.5,
)
def compiled_pool_allocations(key_flat, value_flat) -> list[int]:
args = (
torch.randn(num_tokens, heads, dim, dtype=torch.bfloat16, device=device),
torch.randn(num_tokens, heads, dim, dtype=torch.bfloat16, device=device),
torch.randn(num_tokens, heads, dim, dtype=torch.bfloat16, device=device),
torch.arange(num_tokens, device=device),
torch.empty(0, dtype=torch.long, device=device),
torch.arange(num_blocks, device=device, dtype=torch.int32).unsqueeze(0),
torch.tensor([0, num_tokens], device=device, dtype=torch.int32),
torch.tensor([num_tokens], device=device, dtype=torch.int32),
)
torch._dynamo.reset()
# See the CPU test: inductor's cache outlives torch._dynamo.reset().
with inductor_config.patch(force_disable_caches=True):
_, codes = run_and_get_code(
torch.compile(_Holder(key_flat, value_flat), backend="inductor", fullgraph=True), *args
)
sizes = []
for shape in re.findall(r"empty_strided_cuda\(\(([\d, ]*?)\)", "\n".join(codes)):
dims = [int(part) for part in shape.replace(" ", "").strip(",").split(",") if part]
numel = 1
for dim_size in dims:
numel *= dim_size
if dims and numel >= pool_numel:
sizes.append(numel)
return sizes
# Positive control: one allocation, K and V as its two halves. This must
# produce an allocation of exactly `clone_numel` -- that is the clone
# signature the assertion below checks for, not "no allocation this big".
shared = torch.empty(2, num_blocks * BLOCK, heads, dim, dtype=torch.bfloat16, device=device)
control = compiled_pool_allocations(shared[0], shared[1])
assert clone_numel in control, (
f"the shared layout produced no {clone_numel}-element allocation ({control}), so this test can no "
"longer tell the two layouts apart -- inductor's codegen or the detector changed"
)
_, k_pools, v_pools = allocate_kv_pool_with_views(
num_blocks=num_blocks,
block_size=BLOCK,
num_layers=1,
num_kv_heads=heads,
head_dim=dim,
dtype=torch.bfloat16,
device=device,
)
separate = compiled_pool_allocations(k_pools[0], v_pools[0])
assert clone_numel not in separate, (
f"the compiled graph allocates a {clone_numel}-element buffer, so the pool is being cloned "
f"rather than written in place: {separate} (control saw {control})"
)
@@ -0,0 +1,174 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import weakref
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
from diffusers.models.autoencoders import AutoencoderKLWan
from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
DistributedAutoencoderKLWan,
OmniAutoencoderKLWan,
)
from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@pytest.fixture
def vae():
# Real causal convolutions and caches, with small random weights for CPU CI.
return OmniAutoencoderKLWan(
base_dim=4,
z_dim=2,
dim_mult=[1, 2],
num_res_blocks=1,
temperal_downsample=[True],
scale_factor_spatial=2,
scale_factor_temporal=2,
).eval()
@pytest.mark.parametrize("patch_size", [None, 2])
@pytest.mark.parametrize("return_dict", [False, True])
@torch.no_grad()
def test_chunked_decode_matches_diffusers(vae, patch_size, return_dict):
if patch_size is not None:
vae = OmniAutoencoderKLWan.from_config(vae.config, patch_size=patch_size, out_channels=12).eval()
z = torch.randn(1, 2, 3, 2, 2)
expected = AutoencoderKLWan._decode(vae, z).sample
actual = vae.decode(z, return_dict=return_dict)
torch.testing.assert_close(actual.sample if return_dict else actual[0], expected, rtol=0, atol=0)
chunks: list[torch.Tensor] = []
assert supports_chunked_vae_decode(vae)
assert not supports_chunked_vae_decode(object())
assert vae.decode_with_chunks(z, on_chunk=chunks.append) is None
assert [chunk.shape[2] for chunk in chunks] == [1, 2, 2]
torch.testing.assert_close(torch.cat(chunks, dim=2), expected, rtol=0, atol=0)
assert all(entry is None for entry in vae._feat_map)
@torch.no_grad()
def test_slicing_and_forward_hooks(vae):
z = torch.randn(2, 2, 2, 2, 2)
vae.enable_slicing()
batches = []
handle = vae.post_quant_conv.register_forward_pre_hook(lambda module, args: batches.append(args[0].shape[0]))
vae._hf_hook = SimpleNamespace(pre_forward=Mock())
try:
assert vae.decode(z).sample.shape[0] == 2
assert batches == [1, 1]
vae._hf_hook.pre_forward.assert_called_once_with(vae)
batches.clear()
with pytest.raises(ValueError, match="use_slicing"):
vae.decode_with_chunks(z, on_chunk=Mock())
assert batches == []
vae._hf_hook.pre_forward.reset_mock()
vae.decode_with_chunks(z[:1], on_chunk=Mock())
assert batches == [1]
vae._hf_hook.pre_forward.assert_called_once_with(vae)
finally:
handle.remove()
@torch.no_grad()
def test_callback_failure_drains_and_clears_cache(vae):
z = torch.randn(1, 2, 4, 2, 2)
decoded = Mock(return_value=None)
handle = vae.decoder.register_forward_hook(lambda *args: decoded())
error = RuntimeError("consumer failed")
callback = Mock(side_effect=error)
try:
with pytest.raises(RuntimeError, match="consumer failed") as exc:
vae.decode_with_chunks(z, on_chunk=callback)
assert exc.value is error
assert callback.call_count == 1
assert decoded.call_count == z.shape[2]
assert all(entry is None for entry in vae._feat_map)
finally:
handle.remove()
# The same instance can decode again after the failure.
chunks: list[torch.Tensor] = []
vae.decode_with_chunks(z, on_chunk=chunks.append)
torch.testing.assert_close(torch.cat(chunks, 2), vae.decode(z).sample, rtol=0, atol=0)
def test_chunked_decode_rejects_spatial_tiling(vae):
vae.enable_tiling(tile_sample_min_height=2, tile_sample_min_width=2)
with pytest.raises(ValueError, match="spatial tiling"):
vae.decode_with_chunks(torch.zeros(1, 2, 2, 2, 2), on_chunk=Mock())
@pytest.mark.parametrize("rank", [0, 1])
@pytest.mark.parametrize("fail_callback", [False, True])
def test_spatial_callback_ownership_and_bounded_retention(monkeypatch, rank, fail_callback):
# Isolate the temporal loop from halo kernels; weakrefs detect retention of
# decoded GPU-sized tensors even after the consumer has failed.
monkeypatch.setattr(wan_spatial_shard, "install_wan_spatial_shard_decode", lambda *args, **kwargs: None)
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (rank, 2))
references: list[weakref.ReferenceType[torch.Tensor]] = []
def decode(x, **kwargs):
assert sum(ref() is not None for ref in references) <= 1
chunk = x.clone()
references.append(weakref.ref(chunk))
return chunk
vae = SimpleNamespace(
config=SimpleNamespace(patch_size=None),
clear_cache=Mock(),
post_quant_conv=lambda z: z,
decoder=decode,
_feat_map=[],
)
callback = Mock(side_effect=RuntimeError("consumer failed") if fail_callback else None)
z = torch.zeros(1, 3, 6, 2, 2)
if rank == 0 and fail_callback:
with pytest.raises(RuntimeError, match="consumer failed"):
wan_spatial_shard.spatial_shard_decode(vae, z, group=None, on_chunk=callback)
else:
assert wan_spatial_shard.spatial_shard_decode(vae, z, group=None, on_chunk=callback) is None
assert len(references) == z.shape[2]
assert callback.call_count == (0 if rank else 1 if fail_callback else z.shape[2])
assert vae.clear_cache.call_count == 2
@pytest.mark.parametrize("rank", [0, 1])
@torch.no_grad()
def test_distributed_small_latent_callback_ownership(vae, monkeypatch, rank):
vae.is_distributed_enabled = lambda: True
vae.distributed_executor = SimpleNamespace(group=None)
monkeypatch.setattr(torch.distributed, "get_rank", lambda **kwargs: rank)
callback = Mock()
vae.decode_with_chunks(torch.zeros(1, 2, 2, 2, 2), on_chunk=callback)
assert callback.call_count == (2 if rank == 0 else 0)
@pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width", "tile"])
def test_distributed_chunked_routing(monkeypatch, mode):
vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
torch.nn.Module.__init__(vae)
vae.use_tiling = True
vae.use_slicing = False
vae.tile_sample_min_height = vae.tile_sample_min_width = 2
vae.spatial_compression_ratio = 2
vae.distributed_executor = SimpleNamespace(group=None, parallel_mode=mode, parallel_size=2)
monkeypatch.setattr(vae, "is_distributed_enabled", lambda: True)
monkeypatch.setattr(torch.distributed, "get_world_size", lambda **kwargs: 2)
decode = Mock()
monkeypatch.setattr(wan_spatial_shard, "spatial_shard_decode", decode)
callback = Mock()
z = torch.zeros(1, 2, 2, 2, 2)
if mode == "tile":
with pytest.raises(ValueError, match="require spatial-shard"):
vae.decode_with_chunks(z, on_chunk=callback)
decode.assert_not_called()
else:
assert vae.decode_with_chunks(z, on_chunk=callback) is None
decode.assert_called_once_with(
vae, z, group=None, return_dict=True, split_dim=mode.removeprefix("spatial_shard_"), on_chunk=callback
)
+25 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Shared test helpers for diffusion LoRA tests."""
from __future__ import annotations
@@ -26,15 +26,39 @@ class DummyBaseLayerWithLoRA(torch.nn.Module):
tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
] = []
self.reset_calls: int = 0
self.suspend_calls: int = 0
self.resume_calls: int = 0
self.active_slices: tuple[bool, ...] = ()
self.suspended_slices: tuple[bool, ...] | None = None
self.create_calls: int = 0
def set_lora(self, index: int, lora_a, lora_b):
assert index == 0
self.set_calls.append((lora_a, lora_b))
if isinstance(lora_b, list):
self.active_slices = tuple(b is not None for b in lora_b)
else:
self.active_slices = (True,)
self.suspended_slices = None
def reset_lora(self, index: int):
assert index == 0
self.reset_calls += 1
self.active_slices = ()
self.suspended_slices = None
def suspend_lora(self) -> None:
if self.suspended_slices is not None:
return
self.suspended_slices = self.active_slices
self.active_slices = (False,) * len(self.active_slices)
self.suspend_calls += 1
def resume_lora(self) -> None:
if self.suspended_slices is not None:
self.active_slices = self.suspended_slices
self.suspended_slices = None
self.resume_calls += 1
def create_lora_weights(self, max_loras, lora_config, model_config):
self.create_calls += 1
+125
View File
@@ -30,14 +30,38 @@ class _DummyLoRALayer:
tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
] = []
self.reset_calls: int = 0
self.suspend_calls: int = 0
self.resume_calls: int = 0
self.active_slices: tuple[bool, ...] = ()
self.suspended_slices: tuple[bool, ...] | None = None
def set_lora(self, index: int, lora_a, lora_b):
assert index == 0
self.set_calls.append((lora_a, lora_b))
if isinstance(lora_b, list):
self.active_slices = tuple(b is not None for b in lora_b)
else:
self.active_slices = (True,)
self.suspended_slices = None
def reset_lora(self, index: int):
assert index == 0
self.reset_calls += 1
self.active_slices = ()
self.suspended_slices = None
def suspend_lora(self) -> None:
if self.suspended_slices is not None:
return
self.suspended_slices = self.active_slices
self.active_slices = (False,) * len(self.active_slices)
self.suspend_calls += 1
def resume_lora(self) -> None:
if self.suspended_slices is not None:
self.active_slices = self.suspended_slices
self.suspended_slices = None
self.resume_calls += 1
# Aliases for backward compatibility within this file
@@ -723,3 +747,104 @@ def test_lora_manager_discovers_unet_component(monkeypatch):
assert "unet.down_block.proj" in manager._lora_modules
# Verify the module was actually replaced in the tree (not just recorded)
assert isinstance(pipeline.unet.down_block.proj, _DummyBaseLayerWithLoRA)
def _suspend_harness(monkeypatch, adapter_id: int = 7, rank: int = 2):
"""Manager wired to a single dummy layer, ready to activate `adapter_id`."""
import vllm_omni.diffusion.lora.manager as manager_mod
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
lora_model = _DummyLM(rank=rank)
manager = DiffusionLoRAManager(
pipeline=_DummyPipeline(),
device=torch.device("cpu"),
dtype=torch.bfloat16,
)
monkeypatch.setattr(manager, "_load_adapter", lambda _req: (lora_model, type("PH", (), {"r": rank})()))
manager._registered_adapters = {adapter_id: lora_model}
manager._lora_modules = {"transformer.foo": lora_model.transformer.foo}
return manager, lora_model.transformer.foo
def test_reactivating_same_adapter_resumes_without_rebinding(monkeypatch):
"""activate -> deactivate -> activate at the same scale must re-arm the
saved mask instead of re-uploading every layer."""
adapter_id = 7
req = _dummy_lora_request(adapter_id)
manager, layer = _suspend_harness(monkeypatch, adapter_id)
manager.set_active_adapter(req, lora_scale=0.5)
assert len(layer.set_calls) == 1
assert layer.active_slices == (True,)
manager.set_active_adapter(None)
assert manager._suspended_adapter_id == adapter_id
assert layer.suspend_calls == 1
assert layer.active_slices == (False,)
# Suspending must not tear the upload down.
assert layer.reset_calls == 0
manager.set_active_adapter(req, lora_scale=0.5)
assert layer.resume_calls == 1
assert len(layer.set_calls) == 1, "resume must not rebind"
assert layer.active_slices == (True,)
assert manager._active_adapter_id == adapter_id
assert manager._suspended_adapter_id is None
@pytest.mark.parametrize("second_scale", [0.5, 0.25])
def test_scale_or_adapter_change_forces_rebind(monkeypatch, second_scale):
"""A different scale, or a different adapter, must take the full bind
path rather than re-arming a stale mask."""
adapter_id = 7
manager, layer = _suspend_harness(monkeypatch, adapter_id)
manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5)
manager.set_active_adapter(None)
assert manager._suspended_adapter_id == adapter_id
# Same id at a new scale, or a different id entirely.
other_id = adapter_id if second_scale != 0.5 else adapter_id + 1
if other_id != adapter_id:
manager._registered_adapters[other_id] = manager._registered_adapters[adapter_id]
manager.set_active_adapter(_dummy_lora_request(other_id), lora_scale=second_scale)
assert len(layer.set_calls) == 2, "must rebind"
assert manager._suspended_adapter_id is None
assert manager._active_adapter_id == other_id
def test_partial_packed_mask_survives_suspend_resume(monkeypatch):
"""A packed layer where only some slices carry LoRA must come back with
the same mask, not an all-True one."""
adapter_id = 7
manager, layer = _suspend_harness(monkeypatch, adapter_id)
# Emulate a packed bind that left slice 1 empty.
layer.set_lora(0, [torch.ones(2, 2), None], [torch.ones(2, 2), None])
manager._active_adapter_id = adapter_id
manager._update_adapter_scale(adapter_id, 0.5)
assert layer.active_slices == (True, False)
manager.set_active_adapter(None)
assert layer.active_slices == (False, False)
manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5)
assert layer.active_slices == (True, False), "mask must survive unchanged"
def test_removing_suspended_adapter_drops_the_upload(monkeypatch):
"""A removed adapter can never be resumed, so its weights must not stay
in the stacked buffers."""
adapter_id = 7
manager, layer = _suspend_harness(monkeypatch, adapter_id)
manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5)
manager.set_active_adapter(None)
assert manager._suspended_adapter_id == adapter_id
manager.remove_adapter(adapter_id)
assert manager._suspended_adapter_id is None
assert layer.reset_calls == 1, "the upload must be torn down"
assert layer.suspended_slices is None
@@ -13,7 +13,6 @@ from unittest.mock import MagicMock
import pytest
import torch
import torch.nn as nn
from huggingface_hub import snapshot_download
from safetensors.torch import save_file
from vllm.config.load import LoadConfig
@@ -32,6 +31,7 @@ from vllm_omni.diffusion.models.helios import HeliosPipeline
from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
from vllm_omni.diffusion.registry import initialize_model
from vllm_omni.quantization.component_config import ComponentQuantizationConfig
from vllm_omni.transformers_utils.repo_utils import hf_api
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -41,7 +41,7 @@ model_path = "hf-internal-testing/tiny-helios-modular-pipe"
@pytest.fixture(scope="module")
def prefetch_helios_model():
"""Downloads the tiny helios model prior to running a test."""
snapshot_download(model_path)
hf_api().snapshot_download(model_path)
@pytest.fixture(scope="function")
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import contextlib
@@ -14,10 +14,10 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def test_prefetch_subfolders_propagates_revision(monkeypatch):
calls = []
def fake_snapshot_download(**kwargs):
def fake_snapshot_download(self, **kwargs):
calls.append(kwargs)
monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", fake_snapshot_download)
monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext())
hub_prefetch.prefetch_subfolders(
@@ -0,0 +1,30 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
import torch
from vllm_omni.diffusion.models.utils import make_attention_mask
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@pytest.mark.parametrize("seq_lengths", [[16], [16, 16]])
def test_attention_mask_is_skipped_for_dense_batch(seq_lengths):
hidden_states = torch.empty(len(seq_lengths), 16, 64)
assert make_attention_mask(hidden_states, seq_lengths) is None
def test_attention_mask_preserves_variable_length_padding():
hidden_states = torch.empty(2, 16, 64)
attention_mask = make_attention_mask(hidden_states, [16, 12])
expected = torch.tensor(
[
[True] * 16,
[True] * 12 + [False] * 4,
]
)
torch.testing.assert_close(attention_mask, expected)
@@ -0,0 +1,96 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""L1 tests for checkpoint-to-pipeline routing of Boogu-Image Turbo."""
import json
import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def test_turbo_registry_entry_loads_the_turbo_class():
from vllm_omni.diffusion.models.boogu_image import BooguImageTurboPipeline
from vllm_omni.diffusion.registry import DiffusionModelRegistry
assert DiffusionModelRegistry._try_load_model_cls("BooguImageTurboPipeline") is BooguImageTurboPipeline
def test_turbo_entry_disables_request_batching():
from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline, BooguImageTurboPipeline
assert BooguImagePipeline.supports_request_batch
assert not BooguImageTurboPipeline.supports_request_batch
def test_turbo_process_funcs_resolve_to_the_shared_boogu_helpers():
from vllm_omni.diffusion.models import boogu_image
from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS, _DIFFUSION_PRE_PROCESS_FUNCS
module = boogu_image.pipeline_boogu_image
for funcs in (_DIFFUSION_PRE_PROCESS_FUNCS, _DIFFUSION_POST_PROCESS_FUNCS):
name = funcs["BooguImageTurboPipeline"]
assert name == funcs["BooguImagePipeline"]
assert callable(getattr(module, name))
def test_turbo_inherits_the_base_multimodal_metadata():
from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
metadata = get_diffusion_model_metadata("BooguImageTurboPipeline")
assert metadata == get_diffusion_model_metadata("BooguImagePipeline")
assert metadata.supports_multimodal_inputs
assert metadata.max_multimodal_image_inputs == 1
def test_turbo_model_index_resolves_and_enriches(tmp_path):
from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
(tmp_path / "model_index.json").write_text(
json.dumps({"_class_name": "BooguImageTurboPipeline"}),
encoding="utf-8",
)
assert resolve_model_class_name(str(tmp_path)) == "BooguImageTurboPipeline"
config = OmniDiffusionConfig(model=str(tmp_path))
config.enrich_config()
assert config.model_class_name == "BooguImageTurboPipeline"
assert config.supports_multimodal_inputs
assert config.max_multimodal_image_inputs == 1
def test_edit_turbo_override_selects_the_turbo_class(tmp_path):
"""Edit-Turbo publishes ``BooguImagePipeline``; ``--model-class-name`` must win."""
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.registry import DiffusionModelRegistry
(tmp_path / "model_index.json").write_text(
json.dumps({"_class_name": "BooguImagePipeline"}),
encoding="utf-8",
)
config = OmniDiffusionConfig(model=str(tmp_path), model_class_name="BooguImageTurboPipeline")
config.enrich_config()
assert config.model_class_name == "BooguImageTurboPipeline"
assert DiffusionModelRegistry._try_load_model_cls(config.model_class_name)._is_turbo
def test_edit_turbo_without_override_keeps_the_base_class(tmp_path):
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.registry import DiffusionModelRegistry
(tmp_path / "model_index.json").write_text(
json.dumps({"_class_name": "BooguImagePipeline"}),
encoding="utf-8",
)
config = OmniDiffusionConfig(model=str(tmp_path))
config.enrich_config()
assert config.model_class_name == "BooguImagePipeline"
assert not DiffusionModelRegistry._try_load_model_cls(config.model_class_name)._is_turbo
@@ -1180,6 +1180,21 @@ def test_transfer_fps_matches_resolved_frame_rate_precedence() -> None:
assert cfg.fps == sp.resolved_frame_rate == 12.0
@pytest.mark.parametrize("use_path", [False, True], ids=["image", "path"])
def test_transfer_pil_conversion_returns_writable_array(tmp_path, use_path: bool) -> None:
from vllm_omni.diffusion.models.cosmos3 import transfer
image = Image.new("RGB", (5, 4), "red")
value = image
if use_path:
value = tmp_path / "control.png"
image.save(value)
array = transfer._pil_to_uint8_rgb(value)
assert array.flags.writeable
def test_transfer_vae_executor_requires_distributed_vae() -> None:
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
@@ -2386,6 +2401,42 @@ def test_forward_transfer_runs_multichunk_overlap_path(
torch.testing.assert_close(captured["targets"][1][:, :, 0], torch.full((1, 3, 16, 16), -0.2))
def test_forward_transfer_non_output_rank_uses_canonical_envelope(
make_cosmos3_pipeline,
sequential_cfg_parallel,
) -> None:
pipeline = make_cosmos3_pipeline()
pipeline.vae.distributed_executor = SimpleNamespace(rank=1)
pipeline.vae.is_distributed_enabled = lambda: True
pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16, "1,1")
pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (_ids(2), _mask(), _ids(1), _mask())
pipeline._set_flow_shift = lambda *_args, **_kwargs: None
decoded = torch.zeros(1, 3, 1, 16, 16)
pipeline._decode_latents = lambda latents: decoded
request = SimpleNamespace(
prompts=[{"prompt": "transfer", "modalities": ["video"]}],
sampling_params=make_sampling_params(
height=16,
width=16,
num_inference_steps=1,
guidance_scale=1.0,
extra_args={
"edge": {"control": torch.zeros(3, 1, 16, 16, dtype=torch.uint8)},
"max_frames": 1,
"num_video_frames_per_chunk": 1,
},
),
)
output = pipeline.forward(request)
assert set(output.output) == {"payload", "metadata"}
assert set(output.output["payload"]) == {"video"}
torch.testing.assert_close(output.output["payload"]["video"], decoded)
assert output.output["metadata"] == {"video": {"fps": 24.0}}
def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -> None:
"""With cache-dit active the uncond pass is kept even outside the guidance
interval (so cache-dit's has_separate_cfg parity stays in phase), and the
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -81,9 +81,8 @@ def test_from_config_loads_local_diffusers_component(tmp_path, monkeypatch: pyte
def test_from_config_downloads_component_from_hf_repo(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None:
import huggingface_hub
from vllm_omni.diffusion.models.cosmos3 import sound_tokenizer
from vllm_omni.transformers_utils import repo_utils
cache_dir = tmp_path / "hf"
_write_component(cache_dir, checkpoint_name=DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME)
@@ -95,7 +94,7 @@ def test_from_config_downloads_component_from_hf_repo(tmp_path, monkeypatch: pyt
calls.append((repo_id, revision, allow_patterns))
return str(cache_dir)
monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
monkeypatch.setattr(repo_utils.hf_api(), "snapshot_download", fake_snapshot_download)
sound_tokenizer.Cosmos3SoundTokenizer.from_config(
SimpleNamespace(
@@ -383,6 +383,117 @@ def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None:
)
def test_gen_sp_plan_auto_pads_hidden_states_and_rope_together() -> None:
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
prepare_plan = Cosmos3VFMTransformer._sp_plan["gen_sp_prepare"]
assert set(prepare_plan) == {0, 1, 2}
assert all(spec.auto_pad for spec in prepare_plan.values())
assert [prepare_plan[index].split_dim for index in range(3)] == [1, 1, 1]
def test_gen_sp_prepare_resets_branch_local_padding_metadata() -> None:
from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GenSPPrepare
ctx = ForwardContext(sp_padding_size=1, sp_original_seq_len=12121)
hidden = torch.zeros(1, 4, 8)
rope = torch.zeros(1, 4, 1, 2)
with override_forward_context(ctx):
output = Cosmos3GenSPPrepare()(hidden, rope, rope)
assert output[0] is hidden
assert output[1] is rope
assert output[2] is rope
assert ctx.sp_padding_size == 0
assert ctx.sp_original_seq_len is None
def test_sp_attention_masks_padded_gen_and_joint_und_keys() -> None:
from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3CrossAttention
class RecordingAttention(nn.Module):
def __init__(self) -> None:
super().__init__()
self.metadata: Any | None = None
def forward(self, query, key, value, metadata):
del key, value
self.metadata = metadata
return query
module = object.__new__(Cosmos3CrossAttention)
nn.Module.__init__(module)
module.num_heads_local = 2
module.head_dim = 4
recording_attention = RecordingAttention()
module.attn = recording_attention
q = torch.zeros(1, 2, 2, 4)
k = torch.zeros_like(q)
v = torch.zeros_like(q)
k_und = torch.zeros(1, 3, 2, 4)
v_und = torch.zeros_like(k_und)
ctx = ForwardContext(sp_padding_size=1, sp_original_seq_len=3)
with override_forward_context(ctx):
output = module._forward_sp(q, k, v, k_und, v_und)
assert output.shape == (1, 2, 8)
assert recording_attention.metadata is not None
assert recording_attention.metadata.attn_mask.tolist() == [[True, True, True, False]]
assert recording_attention.metadata.joint_attn_mask.tolist() == [[True, True, True]]
def test_ulysses_2d_mask_uses_key_length_with_kv_only_joint_context(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.parallel import ulysses
def fake_all_to_all(
_group: object,
tensor: torch.Tensor,
_scatter_idx: int,
_gather_idx: int,
_use_sync: bool,
) -> torch.Tensor:
return torch.cat([tensor, tensor], dim=1)[:, :, :1]
monkeypatch.setattr(ulysses.SeqAllToAll4D, "apply", staticmethod(fake_all_to_all))
strategy = ulysses.UlyssesParallelAttention(
sp_group=SimpleNamespace(
ulysses_group=object(),
ulysses_world_size=2,
ulysses_rank=0,
ring_world_size=1,
),
scatter_idx=2,
gather_idx=1,
use_sync=False,
)
query = torch.zeros(1, 2, 2, 4)
key = torch.zeros_like(query)
value = torch.zeros_like(query)
joint_query = torch.zeros(1, 0, 2, 4)
joint_key = torch.zeros(1, 3, 2, 4)
metadata = AttentionMetadata(
attn_mask=torch.tensor([[True, True, True, False]]),
joint_attn_mask=torch.ones(1, 3, dtype=torch.bool),
joint_query=joint_query,
joint_key=joint_key,
joint_value=torch.zeros_like(joint_key),
joint_strategy="front",
)
query_out, key_out, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
assert query_out.shape[1] == 4
assert key_out.shape[1] == 7
assert metadata_out is not None
assert metadata_out.attn_mask.tolist() == [[True, True, True, True, True, True, False]]
def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> None:
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
@@ -721,7 +832,7 @@ def test_forward_returns_video_plus_optional_modality_predictions(
assert [tuple(tensor.shape) for tensor in output] == expected_shapes
def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch: pytest.MonkeyPatch) -> None:
def test_forward_with_sound_odd_sequence_uses_auto_padding_contract(monkeypatch: pytest.MonkeyPatch) -> None:
import vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 as cosmos3_module
model = cosmos3_module.Cosmos3VFMTransformer(
@@ -732,16 +843,17 @@ def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch
)
monkeypatch.setattr(cosmos3_module, "_get_ulysses_state", lambda: (2, 0, None))
with pytest.raises(ValueError, match=r"GEN sequence length \(3 = video tokens 2 \+ sound tokens 1\)"):
model(
hidden_states=torch.zeros(1, 2, 1, 1, 2),
timestep=torch.tensor([1.0]),
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
text_mask=torch.ones(1, 2, dtype=torch.long),
video_shape=(1, 1, 2),
fps=24.0,
sound_latents=torch.zeros(1, 3, 1),
)
output = model(
hidden_states=torch.zeros(1, 2, 1, 1, 2),
timestep=torch.tensor([1.0]),
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
text_mask=torch.ones(1, 2, dtype=torch.long),
video_shape=(1, 1, 2),
fps=24.0,
sound_latents=torch.zeros(1, 3, 1),
)
assert [tuple(tensor.shape) for tensor in output] == [(1, 2, 1, 1, 2), (1, 3, 1)]
def test_sound_latent_frames_padded_for_sequence_parallel(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -0,0 +1,177 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""The request's seed must reach GR00T's flow-matching noise.
The runner builds ``sampling_params.generator`` from ``sampling_params.seed``, the pipeline hands
that generator to the policy, and the action head draws its initial noise from it, so the same
seed reproduces the same action chunk and different seeds give independent ones. These tests pin
both ends of that plumbing.
"""
from __future__ import annotations
from types import SimpleNamespace
import numpy as np
import pytest
import torch
from torch import nn
from transformers.feature_extraction_utils import BatchFeature
from vllm_omni.diffusion.models.gr00t import pipeline_gr00t
from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config
from vllm_omni.diffusion.models.gr00t.modeling.gr00t_n1d7 import Gr00tN1d7ActionHead
from vllm_omni.diffusion.models.gr00t.pipeline_gr00t import Gr00tN1d7Pipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
ROBOT_OBS = {
"images": {"cam": np.zeros((1, 1, 8, 8, 3), dtype=np.uint8)},
"state": {"joint": np.zeros((1, 1, 2), dtype=np.float32)},
"prompt": "pick the cube",
}
class _OptionsRecordingPolicy:
"""Stand-in for ``Gr00tPolicy``; the pipeline only needs ``language_key`` and ``get_action``."""
def __init__(self, *, model_path: str, embodiment_tag: str, device: str, strict: bool) -> None:
del model_path, embodiment_tag, device, strict # the pipeline's constructor contract; unused here
self.language_key = "annotation.language.language_instruction"
self.seen_options: dict[str, object] | None = None
def get_action(self, obs: dict[str, object], options: dict[str, object] | None = None):
self.seen_options = options
return {"arm": np.zeros((1, 2, 2), dtype=np.float32)}, {}
@pytest.fixture
def pipeline(monkeypatch) -> Gr00tN1d7Pipeline:
monkeypatch.setattr(pipeline_gr00t, "Gr00tPolicy", _OptionsRecordingPolicy)
od_config = SimpleNamespace(
model="nvidia/GR00T-N1.7-3B",
model_config={"embodiment_tag": "LIBERO_PANDA", "strict": False},
custom_pipeline_args={},
)
return Gr00tN1d7Pipeline(od_config=od_config)
def _request(**sampling_kwargs) -> OmniDiffusionRequest:
return OmniDiffusionRequest(
prompt="pick",
request_id="req",
sampling_params=OmniDiffusionSamplingParams(extra_args={"robot_obs": ROBOT_OBS}, **sampling_kwargs),
)
def test_runner_generator_is_passed_through_untouched(pipeline):
"""In serving the runner has already built ``sampling_params.generator``; the head must get that object."""
generator = torch.Generator().manual_seed(99)
assert pipeline.forward(_request(generator=generator)).error is None
assert pipeline.policy.seen_options["generator"] is generator
def test_seed_without_generator_builds_one_on_the_pipeline_device(pipeline):
"""Callers that bypass the runner (direct use, tests) still get the seed honoured."""
pipeline.forward(_request(seed=1234))
generator = pipeline.policy.seen_options["generator"]
assert isinstance(generator, torch.Generator)
assert generator.initial_seed() == 1234
assert generator.device.type == pipeline.device
def test_auto_assigned_seed_reaches_the_head(pipeline):
"""``OmniDiffusionRequest`` seeds every request, so the head is never left on the global RNG."""
req = _request()
assert req.sampling_params.seed is not None
pipeline.forward(req)
assert pipeline.policy.seen_options["generator"].initial_seed() == req.sampling_params.seed
def _captured_warnings(monkeypatch) -> list[tuple]:
"""vLLM loggers do not propagate to the root logger, so ``caplog`` cannot see them; record calls instead."""
captured: list[tuple] = []
monkeypatch.setattr(pipeline_gr00t.logger, "warning", lambda *args, **kwargs: captured.append(args))
return captured
def test_unexpected_generator_form_warns_and_falls_back_to_the_seed(pipeline, monkeypatch):
"""GR00T draws one noise tensor per request, so a list of several generators cannot be honoured."""
warnings = _captured_warnings(monkeypatch)
generators = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)]
pipeline.forward(_request(generator=generators, seed=7))
assert len(warnings) == 1
assert "torch.Generator" in warnings[0][0]
assert pipeline.policy.seen_options["generator"].initial_seed() == 7
def test_unexpected_generator_form_without_a_seed_warns_and_leaves_the_global_rng(pipeline, monkeypatch):
warnings = _captured_warnings(monkeypatch)
generators = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)]
pipeline.forward(_request(generator=generators))
assert len(warnings) == 1
assert pipeline.policy.seen_options is None
def _noise_only_head(action_horizon: int, action_dim: int) -> Gr00tN1d7ActionHead:
"""An action head whose single denoising step is the identity, so ``action_pred`` is the initial noise."""
config = Gr00tN1d7Config(
max_action_dim=action_dim,
action_horizon=action_horizon,
num_inference_timesteps=1,
add_pos_embed=False,
use_alternate_vl_dit=False,
)
head = Gr00tN1d7ActionHead.__new__(Gr00tN1d7ActionHead)
nn.Module.__init__(head)
head.config = config
head.action_dim = config.max_action_dim
head.action_horizon = config.action_horizon
head.num_inference_timesteps = config.num_inference_timesteps
head.num_timestep_buckets = config.num_timestep_buckets
head.action_encoder = lambda actions, timesteps, embodiment_id: actions
head.model = lambda hidden_states, encoder_hidden_states, timestep: hidden_states
head.action_decoder = lambda hidden_states, embodiment_id: torch.zeros_like(hidden_states)
return head
def _predict(head: Gr00tN1d7ActionHead, options) -> torch.Tensor:
batch, seq, dim = 2, 3, head.action_dim
return head.get_action_with_features(
backbone_features=torch.zeros(batch, seq, dim),
state_features=torch.zeros(batch, 1, dim),
embodiment_id=torch.zeros(batch, dtype=torch.long),
backbone_output=BatchFeature(data={}),
action_input=BatchFeature(data={}),
options=options,
)["action_pred"]
def test_head_draws_its_noise_from_the_request_generator():
head = _noise_only_head(action_horizon=4, action_dim=3)
expected = torch.randn(2, 4, 3, generator=torch.Generator().manual_seed(5))
assert torch.equal(_predict(head, {"generator": torch.Generator().manual_seed(5)}), expected)
assert not torch.equal(_predict(head, {"generator": torch.Generator().manual_seed(6)}), expected)
def test_head_without_generator_stays_on_the_global_rng():
"""Upstream Isaac-GR00T parity: no generator means a bare ``torch.randn`` on the global RNG."""
head = _noise_only_head(action_horizon=4, action_dim=3)
torch.manual_seed(5)
expected = torch.randn(2, 4, 3)
torch.manual_seed(5)
assert torch.equal(_predict(head, None), expected)
@@ -1,4 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
from typing import Any
import numpy as np
import pytest
@@ -12,12 +16,12 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class FakeGr00tPolicy:
instances = []
instances: list["FakeGr00tPolicy"] = []
def __init__(self, **kwargs):
self.kwargs = kwargs
self.reset_calls = 0
self.seen_obs = None
self.seen_obs: dict[str, Any] | None = None
self.embodiment_tag = SimpleNamespace(value="fake_embodiment")
self.language_key = "annotation.language.language_instruction"
self.modality_configs = {
@@ -40,7 +44,7 @@ class FakeGr00tPolicy:
)
FakeGr00tPolicy.instances.append(self)
def get_action(self, obs):
def get_action(self, obs, options=None):
self.seen_obs = obs
return {
"arm": np.array([[[1.0, 2.0]]], dtype=np.float64),
@@ -107,6 +111,7 @@ def test_forward_returns_dict_actions_in_output():
assert actions["arm"].dtype == np.float32
np.testing.assert_allclose(actions["arm"], np.array([[[1.0, 2.0]]], dtype=np.float32))
policy = FakeGr00tPolicy.instances[0]
assert policy.seen_obs is not None
assert "video" in policy.seen_obs
assert policy.seen_obs["language"] == {"annotation.language.language_instruction": [["pick the cube"]]}
assert "images" not in policy.seen_obs
@@ -0,0 +1,292 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""GR00T-N1.7 eval image transform: chosen from ``processor_config.json`` and bit-exact with Isaac-GR00T.
``GOLDEN`` holds sha256 digests of the ``(C, H, W)`` uint8 output of Isaac-GR00T's own eval transform
(``build_image_transformations_albumentations(...)[1]`` applied through ``apply_with_replay``) for the
synthetic frames built below, at Isaac-GR00T 51d4c89 / albumentations 1.4.18 / opencv 5.0.0.93. The frame
generators are integer-only so they are identical on every platform.
Regenerating a digest needs Isaac-GR00T installed, whose transformers pin conflicts with this repo's, so it
cannot run here: in an Isaac-GR00T environment push each ``synthetic_frames()`` frame through that transform
with ``letter_box_transform`` False and True and hash the output with ``digest()``.
"""
from __future__ import annotations
import hashlib
import json
import numpy as np
import pytest
import torch
import torchvision.transforms.v2 as transforms
from vllm_omni.diffusion.models.gr00t.modeling import processing_gr00t_n1d7 as processing
from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import (
AlbumentationsEvalImageTransform,
Gr00tN1d7Processor,
LetterBoxTransform,
_build_eval_image_transform,
_fractional_center_crop,
_smallest_edge_resize,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
# nvidia/GR00T-N1.7-3B and nvidia/GR00T-N1.7-DROID both ship exactly these processor_kwargs.
TARGET, CROP, EDGE, FRACTION = [256, 256], [230, 230], 256, 0.95
MODALITY_CONFIGS = {
"new_embodiment": {
"video": {"delta_indices": [0], "modality_keys": ["cam"]},
"state": {"delta_indices": [0], "modality_keys": ["joint"]},
"action": {"delta_indices": [0, 1], "modality_keys": ["joint"]},
}
}
# Isaac-GR00T eval transform output for synthetic_frames(); key = "<frame>|letterbox=<flag>".
GOLDEN: dict[str, tuple[tuple[int, int, int], str]] = {
"gradient_256x256|letterbox=False": (
(3, 256, 256),
"4e90a1a3c49539b3648aad752bba59af452bd7231e5712e871ab941cbea1d715",
),
"checker_256x256|letterbox=False": (
(3, 256, 256),
"d2225b94efea5ab3f1c3e6dbae52b7afa3e0f7172d5de24439a2e729ab56080c",
),
"noise_256x256|letterbox=False": (
(3, 256, 256),
"a427ae3b84ebcc6e87470c79c0e45cf684770d6e31158a49f8736685e0da72c6",
),
"noise_512x768|letterbox=False": (
(3, 256, 383),
"45bc3a309f3d06c3f31051911bc8a544ff3f067b023b999dd9c2b7e6029ecc93",
),
"gradient_256x256|letterbox=True": (
(3, 256, 256),
"4e90a1a3c49539b3648aad752bba59af452bd7231e5712e871ab941cbea1d715",
),
"checker_256x256|letterbox=True": (
(3, 256, 256),
"d2225b94efea5ab3f1c3e6dbae52b7afa3e0f7172d5de24439a2e729ab56080c",
),
"noise_256x256|letterbox=True": ((3, 256, 256), "a427ae3b84ebcc6e87470c79c0e45cf684770d6e31158a49f8736685e0da72c6"),
"noise_512x768|letterbox=True": ((3, 256, 256), "52a3db185e6ee4f00617b87711dd198b8f2e2ad53865713d34e56225a4ec61a2"),
}
def _hash_noise(height: int, width: int, seed: int) -> np.ndarray:
"""Integer-only pseudo-noise (murmur-style mixing on uint32); no RNG or float, so identical everywhere."""
ys, xs = np.mgrid[0:height, 0:width].astype(np.uint32)
channels = []
for channel in range(3):
v = xs * np.uint32(73856093) ^ ys * np.uint32(19349663) ^ np.uint32((seed * 3 + channel) * 83492791)
v ^= v >> np.uint32(13)
v *= np.uint32(0x5BD1E995)
v ^= v >> np.uint32(15)
channels.append((v & np.uint32(0xFF)).astype(np.uint8))
return np.stack(channels, axis=-1)
def synthetic_frames() -> dict[str, np.ndarray]:
"""Deterministic HWC uint8 RGB frames: smooth, hard-edged, noisy, and a non-square noisy one.
The non-square frame is 512x768 so both INTER_AREA downscales (2x without letterbox, 3x with it)
are integer factors, which OpenCV computes with exact integer sums: the goldens do not depend on
the CPU architecture. The 243->256 upscale is a fixed-point bilinear on uint8, also exact.
"""
ys, xs = np.mgrid[0:256, 0:256].astype(np.int64)
frames = {
"gradient_256x256": np.stack([xs % 256, ys % 256, (xs + ys) // 2 % 256], axis=-1).astype(np.uint8),
"checker_256x256": np.stack(
[((ys // 16 + xs // 16) % 2) * 255, ((ys // 8 + xs // 8) % 2) * 255, ((ys // 32) % 2) * 255],
axis=-1,
).astype(np.uint8),
"noise_256x256": _hash_noise(256, 256, seed=1),
"noise_512x768": _hash_noise(512, 768, seed=2),
}
return frames
def digest(tensor: torch.Tensor) -> str:
return hashlib.sha256(np.ascontiguousarray(tensor.numpy()).tobytes()).hexdigest()
class _FakeTokenizer:
padding_side = "right"
class _FakeVlmProcessor:
"""The two attributes ``Gr00tN1d7Processor`` and its collator touch on the Qwen3-VL processor."""
def __init__(self) -> None:
self.tokenizer = _FakeTokenizer()
def apply_chat_template(self, *_args: object, **_kwargs: object) -> str:
return "<prompt>"
@pytest.fixture(autouse=True)
def fake_vlm_processor(monkeypatch: pytest.MonkeyPatch) -> None:
"""The real Qwen3-VL processor is a gated download and irrelevant to the image pipeline."""
monkeypatch.setattr(processing, "build_processor", lambda _model_name, _kwargs: _FakeVlmProcessor())
def _processor(**kwargs) -> Gr00tN1d7Processor:
return Gr00tN1d7Processor(
modality_configs=MODALITY_CONFIGS, image_crop_size=CROP, image_target_size=TARGET, **kwargs
)
def _has_letterbox(compose: transforms.Compose) -> bool:
return any(isinstance(step, LetterBoxTransform) for step in compose.transforms)
def test_processor_builds_the_pipeline_named_by_processor_config():
albumentations = _processor(use_albumentations=True, shortest_image_edge=EDGE, crop_fraction=FRACTION)
transform = albumentations.eval_image_transform
assert isinstance(transform, AlbumentationsEvalImageTransform)
assert (transform.max_size, transform.fraction, transform.letter_box_transform) == (EDGE, FRACTION, False)
assert albumentations.use_albumentations is True
torchvision = _processor(use_albumentations=False)
assert isinstance(torchvision.eval_image_transform, transforms.Compose)
assert not _has_letterbox(torchvision.eval_image_transform)
letterboxed = _processor(use_albumentations=False, letter_box_transform=True)
assert _has_letterbox(letterboxed.eval_image_transform)
assert _processor(use_albumentations=True, letter_box_transform=True).eval_image_transform.letter_box_transform
def test_legacy_config_without_the_flags_keeps_the_torchvision_recipe():
"""Checkpoints predating the flags behave like Isaac-GR00T's defaults; training-only keys stay ignored."""
processor = _processor(state_dropout_prob=0.2, extra_augmentation_config=None)
assert processor.use_albumentations is False
assert processor.letter_box_transform is False
assert isinstance(processor.eval_image_transform, transforms.Compose)
assert not _has_letterbox(processor.eval_image_transform)
def test_fraction_and_max_size_fall_back_like_isaac_gr00t():
transform = _build_eval_image_transform(
TARGET, CROP, shortest_image_edge=None, crop_fraction=None, use_albumentations=True
)
assert isinstance(transform, AlbumentationsEvalImageTransform)
assert transform.fraction == CROP[0] / TARGET[0]
assert transform.max_size == TARGET[0]
with pytest.raises(ValueError, match="crop_fraction is None"):
_build_eval_image_transform(TARGET, None, crop_fraction=None, use_albumentations=True)
def test_geometry_follows_isaac_rounding_rules():
frame = _hash_noise(540, 640, seed=3)
resized = _smallest_edge_resize(frame, EDGE)
assert resized.shape == (256, 303, 3) # round() per dimension, shortest edge -> 256
cropped = _fractional_center_crop(resized, FRACTION)
assert cropped.shape == (243, 287, 3) # int() truncation: 243.2 -> 243, 287.85 -> 287
assert np.array_equal(cropped, resized[6:249, 8:295]) # (256-243)//2, (303-287)//2
square = frame[:256, :256]
assert _smallest_edge_resize(square, EDGE) is square # scale == 1.0 is a no-op, not a resample
@pytest.mark.parametrize("case", sorted(GOLDEN))
def test_albumentations_recipe_is_bit_exact_with_isaac_gr00t(case: str):
name, flag = case.split("|letterbox=")
transform = _build_eval_image_transform(
TARGET,
CROP,
shortest_image_edge=EDGE,
crop_fraction=FRACTION,
use_albumentations=True,
letter_box_transform=(flag == "True"),
)
out = transform(synthetic_frames()[name])
expected_shape, expected_digest = GOLDEN[case]
assert out.dtype == torch.uint8
assert tuple(out.shape) == expected_shape
assert digest(out) == expected_digest
def test_letterbox_is_gated_on_both_recipes():
frame = _hash_noise(540, 640, seed=3)
def albumentations(letterbox: bool) -> torch.Tensor:
return _build_eval_image_transform(
TARGET,
CROP,
shortest_image_edge=EDGE,
crop_fraction=FRACTION,
use_albumentations=True,
letter_box_transform=letterbox,
)(frame)
padded, unpadded = albumentations(True), albumentations(False)
assert tuple(padded.shape) == (3, 256, 256) # 540x640 -> 640x640 -> 256x256 -> 243x243 -> 256x256
assert tuple(unpadded.shape) == (3, 256, 302) # aspect preserved: 256x303 -> 243x287 -> 256x302
tv_padded = _build_eval_image_transform(TARGET, CROP, use_albumentations=False, letter_box_transform=True)(frame)
tv_unpadded = _build_eval_image_transform(TARGET, CROP, use_albumentations=False, letter_box_transform=False)(frame)
assert tuple(tv_padded.shape) == tuple(tv_unpadded.shape) == (3, 256, 256)
assert not torch.equal(tv_padded, tv_unpadded)
assert (tv_padded[:, :2, :] == 0).all() and not (tv_unpadded[:, :2, :] == 0).all() # black bars only when gated
def test_albumentations_recipe_rejects_frames_outside_the_contract():
transform = _build_eval_image_transform(TARGET, CROP, use_albumentations=True)
with pytest.raises(ValueError, match="HWC uint8"):
transform(np.zeros((256, 256, 3), dtype=np.float32))
with pytest.raises(ValueError, match="HWC uint8"):
transform(np.zeros((3, 256, 256), dtype=np.uint8))
with pytest.raises(ValueError, match="HWC uint8"):
transform(np.zeros((256, 256), dtype=np.uint8))
def test_get_vlm_inputs_runs_the_selected_transform_per_frame():
processor = _processor(use_albumentations=True, shortest_image_edge=EDGE, crop_fraction=FRACTION)
frames = np.stack([_hash_noise(540, 640, seed=3)] * 2) # (T, H, W, C) like policy.py hands over
vlm_inputs = processor._get_vlm_inputs(
image_keys=["cam"],
images={"cam": frames},
masks=None,
image_transform=processor.eval_image_transform,
language="x",
)
images = vlm_inputs["vlm_content"]["images"]
assert len(images) == 2
assert all(img.size == (302, 256) for img in images) # PIL reports (W, H); no letterbox for this checkpoint
def test_flags_round_trip_through_save_pretrained_and_from_pretrained(tmp_path):
saved = _processor(
use_albumentations=True, letter_box_transform=True, shortest_image_edge=EDGE, crop_fraction=FRACTION
)
saved.save_pretrained(tmp_path)
kwargs = json.loads((tmp_path / "processor_config.json").read_text())["processor_kwargs"]
assert kwargs["use_albumentations"] is True
assert kwargs["letter_box_transform"] is True
loaded = Gr00tN1d7Processor.from_pretrained(tmp_path)
assert loaded.use_albumentations is True
assert loaded.letter_box_transform is True
assert isinstance(loaded.eval_image_transform, AlbumentationsEvalImageTransform)
assert loaded.eval_image_transform.letter_box_transform is True
def test_from_pretrained_defaults_missing_flags_to_isaac_gr00t_defaults(tmp_path):
(tmp_path / "processor_config.json").write_text(
json.dumps(
{
"processor_class": "Gr00tN1d7Processor",
"processor_kwargs": {
"modality_configs": MODALITY_CONFIGS,
"image_crop_size": CROP,
"image_target_size": TARGET,
},
}
)
)
(tmp_path / "statistics.json").write_text("{}")
loaded = Gr00tN1d7Processor.from_pretrained(tmp_path)
assert loaded.use_albumentations is False
assert loaded.letter_box_transform is False
assert isinstance(loaded.eval_image_transform, transforms.Compose)
@@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -13,6 +14,7 @@ from vllm_omni.diffusion.models.lingbot_world.actions import (
LingBotCameraControlReducer,
integrate_lingbot_camera_actions,
parse_lingbot_camera_action_frames,
parse_lingbot_camera_action_script,
)
from vllm_omni.experimental.ar_diffusion.session import (
ARDiffusionSession,
@@ -292,3 +294,8 @@ def test_session_reducer_keeps_prompt_and_actions_atomic_and_resets() -> None:
assert consumer.ticks[1].controls[0].data["frames"] == ((), (), ())
asyncio.run(exercise())
def test_parse_camera_action_script_rejects_wrong_chunk_width() -> None:
with pytest.raises(ValueError, match="exactly 3 per-latent-frame"):
parse_lingbot_camera_action_script([[["w"], ["w"]]], frames_per_chunk=3)
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -15,8 +15,10 @@ from diffusers.utils.torch_utils import randn_tensor as _diffusers_randn_tensor
from PIL import Image
from torch import nn
import vllm_omni.diffusion.models.lingbot_world.dmd_block as lingbot_dmd_block
import vllm_omni.diffusion.models.lingbot_world.pipeline as lingbot_pipeline
from tests.diffusion.models.wan2_2.conftest import noop_progress_bar
from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
from vllm_omni.diffusion.models.lingbot_world.actions import (
integrate_lingbot_camera_actions,
)
@@ -24,6 +26,8 @@ from vllm_omni.diffusion.models.lingbot_world.camera import CameraTrajectory as
from vllm_omni.diffusion.models.lingbot_world.camera import (
build_plucker_embedding as _real_build_plucker_embedding,
)
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.utils import StepRequestState
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
ARDiffusionControlInput,
ARDiffusionTickRequest,
@@ -253,6 +257,8 @@ class _SamplingParams:
if extra_args is not None:
self.extra_args.update(extra_args)
self.latents = None
self.guidance_scale = None
self.guidance_scale_2 = None
class _RequestBatch:
@@ -345,7 +351,6 @@ def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
"FlowUniPCMultistepScheduler": _FakeScheduler,
"CausalLingBotWorldTransformer3DModel": _FakeTransformerFactory,
"get_local_device": lambda: torch.device("cpu"),
"set_forward_context_denoise_step_idx": lambda index: None,
"prefetch_subfolders": prefetch_subfolders,
"from_pretrained_with_prefetch": from_pretrained_with_prefetch,
"load_transformer_config": load_transformer_config,
@@ -358,6 +363,9 @@ def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
}
for name, value in replacements.items():
monkeypatch.setattr(lingbot_pipeline, name, value)
# The DMD block math lives in its own module; stub the symbols it imports.
monkeypatch.setattr(lingbot_dmd_block, "set_forward_context_denoise_step_idx", lambda index: None)
monkeypatch.setattr(lingbot_dmd_block, "randn_tensor", _diffusers_randn_tensor)
monkeypatch.setattr(lingbot_pipeline, "_loader_state", loader_state, raising=False)
@@ -774,6 +782,7 @@ def test_denoise_state_stays_fp32_while_transformer_inputs_use_model_dtype() ->
return torch.zeros(shape, device=device, dtype=dtype)
module.randn_tensor = randn
lingbot_dmd_block.randn_tensor = randn
result = pipeline(_request())
assert requested_noise_dtypes == [torch.float32] * 4
@@ -1139,7 +1148,7 @@ def test_first_frame_condition_and_camera_fold_match_transformer_contract() -> N
module = _load_pipeline_module()
transformer = _RecordingTransformer()
pipeline = _pipeline(module, transformer=transformer)
module.randn_tensor = lambda shape, **kwargs: torch.full(
module.randn_tensor = lingbot_dmd_block.randn_tensor = lambda shape, **kwargs: torch.full(
shape,
-99.0,
device=kwargs["device"],
@@ -1745,6 +1754,288 @@ def test_request_cache_becomes_unreachable_after_transformer_error() -> None:
assert not hasattr(pipeline, "transformer_cache")
class _FakeARState:
def __init__(self, session_id: str = "req-1") -> None:
self.session_id = session_id
self.commits: list[str] = []
def get_kv_caches(self, branch, *, seq_len, commit_current):
del branch, seq_len, commit_current
return [SimpleNamespace()]
def commit_paged_context(self, branch):
self.commits.append(branch)
def clear_cross_attention(self):
return None
def is_cross_attention_populated(self, branch, name):
del branch, name
return True
def get_cross_attention_kv(self, branch, name):
del branch, name
zeros = torch.zeros(1, 1, 2, 4)
return [{"k": zeros, "v": zeros} for _ in range(2)]
def _empty_action_script(num_chunks: int):
return tuple(((), (), ()) for _ in range(num_chunks))
def _stepwise_state(
*,
request_id: str = "req-1",
num_frames: int = 21,
script=None,
seed: int = 17,
):
num_chunks = ((num_frames - 1) // 4 + 1) // 3
sampling = _SamplingParams(num_frames=num_frames, seed=seed)
sampling.extra_args["_lingbot_camera_trajectory"] = None
sampling.extra_args["_lingbot_camera_actions"] = None
sampling.extra_args["_lingbot_camera_action_script"] = (
script if script is not None else _empty_action_script(num_chunks)
)
return StepRequestState(
request_id=request_id,
sampling=sampling,
prompt=_prompt(),
)
def _run_stepwise(pipeline, state):
outputs = []
pipeline.prepare_encode(state)
while not state.request_denoise_completed:
noise = pipeline.denoise_step(None, states=[state])
pipeline.step_scheduler(state, noise)
if state.chunk_denoise_completed:
outputs.append(pipeline.post_decode(state))
return outputs
def test_pipeline_declares_step_execution_support() -> None:
module = _load_pipeline_module()
pipeline = _pipeline(module)
assert pipeline.supports_step_execution is True
assert isinstance(pipeline, SupportsStepExecution)
assert supports_step_execution(pipeline) is True
def test_preprocess_materializes_camera_action_script_without_action_path() -> None:
module = _load_pipeline_module()
sampling = _SamplingParams(include_action=False)
sampling.extra_args["camera_action_script"] = [[["w"], ["w"], ["w"]], [["a"], [], []]]
request = OmniDiffusionRequest(prompt=_prompt(), sampling_params=sampling, request_id="req-1")
result = module.get_lingbot_world_pre_process_func(_od_config())(request)
assert result.sampling_params.extra_args["_lingbot_camera_trajectory"] is None
assert result.sampling_params.extra_args["_lingbot_camera_actions"] is None
assert result.sampling_params.extra_args["_lingbot_camera_action_script"] == (
(("w",), ("w",), ("w",)),
(("a",), (), ()),
)
@pytest.mark.parametrize(
"extra_args",
[
pytest.param({}, id="request_mode"),
pytest.param(_tick_extra_args(chunk_index=0), id="tick_mode"),
],
)
def test_forward_rejects_a_stepwise_camera_action_script(extra_args) -> None:
"""A script only steers step execution, so request mode must not drop it silently."""
module = _load_pipeline_module()
pipeline = _pipeline(module)
sampling = _SamplingParams(
include_action=False,
extra_args={
**extra_args,
"_lingbot_camera_trajectory": None,
"_lingbot_camera_action_script": _empty_action_script(1),
},
)
with pytest.raises(ValueError, match="camera_action_script is read only by LingBot step execution"):
pipeline(_request(sampling=sampling))
def test_stepwise_progress_metadata_and_commit_trace() -> None:
module = _load_pipeline_module()
transformer = _RecordingTransformer()
pipeline = _pipeline(module, transformer=transformer)
pipeline._ar_height = 16
pipeline._ar_width = 16
state = _stepwise_state(num_frames=21)
fake = _FakeARState(state.request_id)
with pipeline.bind_ar_diffusion_state(state.request_id, fake):
outputs = _run_stepwise(pipeline, state)
assert pipeline._ar_diffusion_kv_state is None
assert len(outputs) == 2
assert [output.chunk_index for output in outputs] == [0, 1]
assert all(output.total_chunks == 2 for output in outputs)
assert outputs[-1].finished is True
assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True] * 2
assert [call["start_frame"] for call in transformer.calls] == [0] * 5 + [3] * 5
assert fake.commits == ["main", "main"]
for chunk_index, output in enumerate(outputs):
metadata = output.output["metadata"]["ar_diffusion"]
assert metadata == {
"session_id": "req-1",
"request_id": "req-1",
"chunk_index": chunk_index,
"applied_event_ids": [],
}
assert output.output["payload"]["latents"].shape == (1, 16, 3, 2, 2)
def test_stepwise_matches_tick_transformer_trace_and_latents() -> None:
module = _load_pipeline_module()
transformer = _RecordingTransformer()
pipeline = _pipeline(module, transformer=transformer)
pipeline._ar_height = 16
pipeline._ar_width = 16
script = ((("w",), ("w",), ("w",)), (("d",), ("d",), ("d",)))
fake = _FakeARState("world-1")
with pipeline.bind_ar_diffusion_state("world-1", fake):
tick_latents = []
for chunk_index, frames in enumerate(script):
tick = ARDiffusionTickRequest(
session_id="world-1",
request_id="legacy-lingbot-request",
chunk_index=chunk_index,
controls=(
ARDiffusionControlInput(
track="camera",
schema="lingbot.camera_actions.v1",
data={"mode": "frames", "frames": [list(frame) for frame in frames]},
),
),
)
sampling = _SamplingParams(extra_args=tick.to_extra_args(), seed=17)
sampling.extra_args["_lingbot_camera_trajectory"] = None
sampling.extra_args["_lingbot_camera_actions"] = frames
result = pipeline(_request(sampling=sampling))
tick_latents.append(result.output["payload"]["latents"].clone())
tick_calls = list(transformer.calls)
transformer.calls.clear()
pipeline.close_ar_diffusion_session("world-1")
stepwise_state = _stepwise_state(request_id="req-1", num_frames=21, script=script, seed=17)
stepwise_fake = _FakeARState(stepwise_state.request_id)
with pipeline.bind_ar_diffusion_state(stepwise_state.request_id, stepwise_fake):
stepwise_outputs = _run_stepwise(pipeline, stepwise_state)
assert [call["update_cache"] for call in transformer.calls] == [call["update_cache"] for call in tick_calls]
assert [call["start_frame"] for call in transformer.calls] == [call["start_frame"] for call in tick_calls]
for tick_latent, output in zip(tick_latents, stepwise_outputs, strict=True):
torch.testing.assert_close(output.output["payload"]["latents"], tick_latent)
def test_stepwise_trajectory_camera_matches_request_mode_under_non_uniform_speed(monkeypatch) -> None:
"""A trajectory that slows down between blocks must condition every stepwise
chunk exactly as request mode does. Request mode embeds the whole trajectory
once; embedding per chunk would re-normalize each block's framewise
translations by its own largest step and present a slow block as full-speed
motion."""
from vllm_omni.diffusion.models.lingbot_world import camera as camera_module
module = _load_pipeline_module()
# The fixture stubs the ray embedding with a pose-blind ramp; parity needs
# the real geometry so a scale mismatch actually shows up.
monkeypatch.setattr(module, "build_plucker_embedding", camera_module.build_plucker_embedding)
transformer = _RecordingTransformer()
pipeline = _pipeline(module, transformer=transformer)
pipeline._ar_height = 16
pipeline._ar_width = 16
# Six latent frames -> two 3-frame blocks. The first block moves at full
# speed and everything after it at a tenth of that, so a per-chunk
# normalization would rescale the second block by 10x.
steps = torch.tensor([1.0] * 3 + [0.1] * 18)
poses = torch.eye(4).repeat(21, 1, 1)
poses[:, 2, 3] = torch.cumsum(steps, dim=0) - steps[0]
trajectory = _CameraTrajectory(
poses=poses,
intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(21, 1),
)
sampling = _SamplingParams(num_frames=21)
sampling.extra_args["_lingbot_camera_trajectory"] = trajectory
pipeline(_request(sampling=sampling))
request_cameras = [call["camera_hidden_states"] for call in transformer.calls]
transformer.calls.clear()
state = _stepwise_state(num_frames=21)
state.sampling.extra_args["_lingbot_camera_trajectory"] = trajectory
state.sampling.extra_args.pop("_lingbot_camera_action_script")
with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
_run_stepwise(pipeline, state)
stepwise_cameras = [call["camera_hidden_states"] for call in transformer.calls]
# Two blocks x (four probes + one commit), in the same order on both paths.
assert len(request_cameras) == len(stepwise_cameras) == 10
for stepwise, request in zip(stepwise_cameras, request_cameras, strict=True):
torch.testing.assert_close(stepwise, request)
def test_stepwise_emits_latents_without_touching_the_vae() -> None:
"""Latent mode must stay latent: streaming decode is opt-in per request."""
module = _load_pipeline_module()
pipeline = _pipeline(module, transformer=_RecordingTransformer())
pipeline._ar_height = 16
pipeline._ar_width = 16
state = _stepwise_state(num_frames=21)
fake = _FakeARState(state.request_id)
with pipeline.bind_ar_diffusion_state(state.request_id, fake):
outputs = _run_stepwise(pipeline, state)
assert [set(output.output["payload"]) for output in outputs] == [{"latents"}, {"latents"}]
assert pipeline.vae.decode_inputs == []
def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> None:
"""A non-latent request must stream pixels under the primary "video" key."""
import diffusers.video_processor as video_processor_module
module = _load_pipeline_module()
processed = []
class VideoProcessor:
def __init__(self, *, vae_scale_factor):
assert vae_scale_factor == 8
def postprocess_video(self, video, *, output_type):
processed.append((tuple(video.shape), output_type))
return f"frames-{len(processed)}"
monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor)
pipeline = _pipeline(module, transformer=_RecordingTransformer())
pipeline._ar_height = 16
pipeline._ar_width = 16
state = _stepwise_state(num_frames=21)
state.sampling.output_type = "np"
fake = _FakeARState(state.request_id)
with pipeline.bind_ar_diffusion_state(state.request_id, fake):
outputs = _run_stepwise(pipeline, state)
# One decode per AR block, each seeing only that block's three latent frames.
assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2
assert [output_type for _, output_type in processed] == ["np", "np"]
assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}]
# Identity metadata survives the switch to pixel output.
assert [output.output["metadata"]["ar_diffusion"]["chunk_index"] for output in outputs] == [0, 1]
def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None:
module = _load_pipeline_module()
resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry(
@@ -1755,7 +2046,7 @@ def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> No
assert resolved is module.LingBotWorldCausalDMDPipeline
assert cache_acceleration_disabled
assert preprocess_name == "get_lingbot_world_pre_process_func"
request = SimpleNamespace(prompt=_prompt(), sampling_params=_SamplingParams())
request = OmniDiffusionRequest(prompt=_prompt(), sampling_params=_SamplingParams(), request_id="req-1")
assert preprocess(request) is request
assert isinstance(request.sampling_params.extra_args["_lingbot_camera_trajectory"], _CameraTrajectory)
assert module.LingBotWorldCausalDMDPipeline.__name__ == "LingBotWorldCausalDMDPipeline"
@@ -5,6 +5,7 @@
import json
from dataclasses import replace
from threading import Lock
from types import SimpleNamespace
from typing import Any
@@ -500,6 +501,32 @@ def test_ltx25_checkpoint_selects_distilled_two_stage_profile(tmp_path, monkeypa
assert pipe.pipeline_recipe is LTX25_DISTILLED_TWO_STAGE_RECIPE
def test_ltx_runtime_profiles_phase_methods_separately(monkeypatch):
pipe = object.__new__(LTX2DistilledTwoStagePipeline)
pipe.enable_diffusion_pipeline_profiler = True
pipe._profiler_lock = Lock()
pipe._stage_durations = {}
monkeypatch.setattr(ltx2_runtime.current_omni_platform, "is_available", lambda: False)
with pipe._profile_phase_method("run_phase", "generate_lowres"):
pass
with pipe._profile_phase_method("run_phase", "refine"):
pass
with pipe._profile_phase_method("decode_phase"):
pass
output = DiffusionOutput()
assert pipe._attach_profiled_stage_durations(output) is output
expected_keys = {
"LTX2DistilledTwoStagePipeline.generate_lowres.run_phase",
"LTX2DistilledTwoStagePipeline.refine.run_phase",
"LTX2DistilledTwoStagePipeline.decode_phase",
}
assert set(pipe.stage_durations) == expected_keys
assert set(output.stage_durations) == expected_keys
assert all(duration >= 0.0 for duration in pipe.stage_durations.values())
def test_ltx25_full_config_enrichment_uses_selected_transformer_subfolder(tmp_path):
from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -103,7 +103,7 @@ def test_prepare_audio_latents_samples_directly_in_packed_token_space():
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
def test_prepare_audio_latents_pads_short_audio_for_sp():
pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=2)
latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for LTX guidance and forward-parallel behavior."""
@@ -85,14 +85,14 @@ class TestCFGParallelHelpers:
LTX_GUIDANCE_EXECUTOR.validate_guidance_world_size(plan, guidance_world_size)
def test_guidance_parallel_dummy_warms_full_guidance_plan(self, monkeypatch):
def test_guidance_parallel_preserves_full_guidance_plan(self, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 2)
pipe = object.__new__(LTX2Pipeline)
req = SimpleNamespace(is_dummy_run=lambda: True)
req = SimpleNamespace()
request_inputs = SimpleNamespace(guidance=LTX23_ONE_STAGE_RECIPE.request_guidance)
guidance_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None)
@@ -10,6 +10,7 @@ from dataclasses import replace
from types import SimpleNamespace
from typing import Any
import huggingface_hub
import numpy as np
import pytest
import torch
@@ -175,7 +176,19 @@ def test_ltx_device_uint8_output_rejects_invalid_contract(video, error):
_prepare_ltx2_video_output(video, do_normalize=True)
def test_ltx_profiler_covers_video_output_transport(tmp_path, monkeypatch):
@pytest.mark.parametrize(
"parallel_config",
[
None,
SimpleNamespace(ulysses_degree=4),
SimpleNamespace(tensor_parallel_size=2),
SimpleNamespace(use_hsdp=True, hsdp_shard_size=2),
SimpleNamespace(ulysses_degree=2, tensor_parallel_size=2),
SimpleNamespace(ulysses_degree=2, use_hsdp=True, hsdp_shard_size=2),
],
ids=["single", "ulysses", "tp", "hsdp", "tp-ulysses", "hsdp-ulysses"],
)
def test_ltx_supported_parallel_configs_keep_output_profiling(tmp_path, monkeypatch, parallel_config):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]}))
@@ -193,7 +206,13 @@ def test_ltx_profiler_covers_video_output_transport(tmp_path, monkeypatch):
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", capture_profiler_setup)
LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False))
LTX2Pipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
parallel_config=parallel_config,
enable_diffusion_pipeline_profiler=False,
)
)
assert "video_processor.postprocess_video" in profiler_setup["profiler_targets"]
assert "_prepare_video_output_for_transport" in profiler_setup["profiler_targets"]
@@ -284,11 +303,11 @@ def test_ltx_artifact_uses_source_revision_and_hub_fallback(
filename = "ltx-sidecar.safetensors"
calls = []
def fake_download(**kwargs):
def fake_download(self, **kwargs):
calls.append(kwargs)
return "/cache/ltx-sidecar.safetensors"
monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
monkeypatch.setattr(huggingface_hub.HfApi, "hf_hub_download", fake_download)
assert (
resolve_ltx_artifact(
@@ -313,7 +332,9 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
filename = "ltx-sidecar.safetensors"
expected = tmp_path / filename
expected.write_bytes(b"sidecar")
monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup"))
monkeypatch.setattr(
huggingface_hub.HfApi, "hf_hub_download", lambda *_args, **_kwargs: pytest.fail("unexpected Hub lookup")
)
assert resolve_ltx_artifact(
str(tmp_path),
@@ -327,11 +348,11 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
def test_ltx_artifact_local_model_missing_sidecar_falls_back_to_hub(tmp_path, monkeypatch):
calls = []
def fake_download(**kwargs):
def fake_download(self, **kwargs):
calls.append(kwargs)
return "/cache/ltx-sidecar.safetensors"
monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
monkeypatch.setattr(huggingface_hub.HfApi, "hf_hub_download", fake_download)
assert (
resolve_ltx_artifact(
@@ -383,14 +404,42 @@ def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expecte
)
def test_ltx_rejects_advanced_uaa_before_component_initialization():
@pytest.mark.parametrize(
"pipeline_cls",
[
LTX2Pipeline,
LTX2DistilledOneStagePipeline,
LTX2TwoStagePipeline,
LTX2DistilledTwoStagePipeline,
LTX2T2VDMD2Pipeline,
LTX2I2VDMD2Pipeline,
],
)
@pytest.mark.parametrize(
("parallel_config", "error"),
[
(SimpleNamespace(ulysses_mode="advanced_uaa"), "does not support ulysses_mode='advanced_uaa'"),
(SimpleNamespace(ring_degree=2), "pure Ulysses sequence parallelism only"),
(SimpleNamespace(ulysses_degree=2, ring_degree=2), "pure Ulysses sequence parallelism only"),
(SimpleNamespace(allgather_degree=2), "pure Ulysses sequence parallelism only"),
],
ids=["advanced-uaa", "ring", "hybrid", "allgather-kv"],
)
def test_ltx_rejects_unsupported_sp_before_component_initialization(pipeline_cls, parallel_config, error, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
def unexpected_model_access(*_args, **_kwargs):
pytest.fail("Unsupported SP must be rejected before accessing model files or initializing components")
monkeypatch.setattr(ltx2_runtime, "detect_ltx_model_version", unexpected_model_access)
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", unexpected_model_access)
od_config = SimpleNamespace(
model="unused",
parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"),
parallel_config=parallel_config,
)
with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"):
LTX2Pipeline(od_config=od_config)
with pytest.raises(ValueError, match=error):
pipeline_cls(od_config=od_config)
def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch):
@@ -1409,39 +1458,19 @@ def test_ltx_distilled_forward_rejects_fixed_recipe_overrides(direct_kwargs, sam
pipeline.forward(req, **direct_kwargs)
def test_ltx_distilled_dummy_run_uses_fixed_recipe_values():
from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
req = DiffusionRequestBatch(
[
OmniDiffusionRequest(
prompt="dummy run",
sampling_params=OmniDiffusionSamplingParams(
height=512,
width=512,
num_frames=LTX2DistilledPipeline.dummy_run_num_frames,
num_inference_steps=1,
guidance_scale=0.0,
),
request_id=DUMMY_DIFFUSION_REQUEST_ID,
)
]
)
pipeline = _make_ltx_request_pipe(LTX2DistilledPipeline)
seen = {}
def fake_run_recipe(_req, request_inputs, **_kwargs):
seen["request_inputs"] = request_inputs
return "dummy-output"
object.__setattr__(pipeline, "_run_recipe", fake_run_recipe)
assert pipeline.forward(req) == "dummy-output"
assert seen["request_inputs"].num_frames == 1
assert seen["request_inputs"].num_inference_steps == 8
assert seen["request_inputs"].guidance == LTXGuidanceSpec.positive_only()
@pytest.mark.parametrize(
"pipeline_cls",
[
LTX2Pipeline,
LTX2DistilledOneStagePipeline,
LTX2TwoStagePipeline,
LTX2DistilledTwoStagePipeline,
LTX2T2VDMD2Pipeline,
LTX2I2VDMD2Pipeline,
],
)
def test_ltx_pipelines_disable_startup_dummy_run(pipeline_cls):
assert pipeline_cls.dummy_run_num_frames == 0
def test_ltx_custom_sigmas_bypass_scheduler_shifting():
@@ -0,0 +1,100 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
from datetime import timedelta
from types import SimpleNamespace
import pytest
import torch
import torch.distributed as dist
import torch.nn.functional as F
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.models.ltx2.ltx2_sequence_parallel import (
LTX2VideoToAudioParallelAttention,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def _run_video_to_audio_parity(rank: int, world_size: int, master_port: int) -> None:
dist.init_process_group(
"gloo",
init_method=f"tcp://127.0.0.1:{master_port}",
rank=rank,
world_size=world_size,
timeout=timedelta(seconds=60),
)
try:
# Non-contiguous SP subgroups match TP-fastest rank ordering. Different
# seeds across groups catch collectives that accidentally use WORLD.
group_count = world_size // 2
ulysses_group = dist.group.WORLD
if group_count > 1:
for offset in range(group_count):
ranks = list(range(offset, world_size, group_count))
group = dist.new_group(ranks)
if rank in ranks:
ulysses_group = group
ulysses_rank = dist.get_rank(ulysses_group)
sp_group = SimpleNamespace(
ring_world_size=1,
ulysses_group=ulysses_group,
ulysses_world_size=2,
ulysses_rank=ulysses_rank,
)
strategy = LTX2VideoToAudioParallelAttention(sp_group)
for batch_size, masked in [(1, False), (2, True)]:
torch.manual_seed(17 + rank % group_count)
query = torch.randn(batch_size, 3, 4, 8)
global_key = torch.randn(batch_size, 6, 4, 8)
global_value = torch.randn(batch_size, 6, 4, 8)
key = global_key.chunk(2, dim=1)[ulysses_rank].contiguous()
value = global_value.chunk(2, dim=1)[ulysses_rank].contiguous()
mask = None
if masked:
mask = torch.ones(batch_size, 1, 1, 6, dtype=torch.bool)
mask[0, ..., -1] = False
mask[1, ..., -2:] = False
metadata = AttentionMetadata(attn_mask=mask) if masked else None
local_query, local_key, local_value, local_metadata, ctx = strategy.pre_attention(
query, key, value, metadata
)
assert local_metadata is metadata
local_output = F.scaled_dot_product_attention(
local_query.transpose(1, 2),
local_key.transpose(1, 2),
local_value.transpose(1, 2),
attn_mask=local_metadata.attn_mask if local_metadata is not None else None,
).transpose(1, 2)
actual = strategy.post_attention(local_output, ctx)
expected = F.scaled_dot_product_attention(
query.transpose(1, 2),
global_key.transpose(1, 2),
global_value.transpose(1, 2),
attn_mask=mask,
).transpose(1, 2)
torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
finally:
dist.destroy_process_group()
@pytest.mark.parametrize("world_size", [2, 4], ids=["world", "subgroups"])
def test_ltx_video_to_audio_sp_matches_replicated_attention(unused_tcp_port, world_size):
torch.multiprocessing.spawn(
_run_video_to_audio_parity,
args=(world_size, unused_tcp_port),
nprocs=world_size,
)
def test_ltx_video_to_audio_sp_rejects_ring_parallelism():
sp_group = SimpleNamespace(ring_world_size=2)
with pytest.raises(NotImplementedError, match="ring_degree must be 1"):
LTX2VideoToAudioParallelAttention(sp_group)
@@ -141,16 +141,54 @@ def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
@pytest.mark.parametrize("rope_type", ["interleaved", "split"])
def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
root_plan = LTX2VideoTransformer3DModel._build_sp_plan(rope_type)[""]
def test_ltx_sp_plan_only_shards_video_stream(rope_type):
plan = LTX2VideoTransformer3DModel._build_sp_plan(rope_type)
root_plan = plan[""]
video_timestep = root_plan["timestep"]
audio_timestep = root_plan["audio_timestep"]
assert set(root_plan) == {"hidden_states", "keyframes_mask", "timestep"}
assert root_plan["hidden_states"].split_dim == 1
assert root_plan["keyframes_mask"].split_dim == 1
assert root_plan["timestep"].split_dim == 1
assert root_plan["timestep"].expected_dims == 2
assert not root_plan["timestep"].split_output
assert set(plan) == {"", "rope", "cross_attn_rope", "proj_out"}
assert video_timestep.split_dim == audio_timestep.split_dim == 1
assert video_timestep.expected_dims == audio_timestep.expected_dims == 2
assert not video_timestep.split_output
assert not audio_timestep.split_output
def test_ltx_attention_assigns_video_only_sp_modes(monkeypatch):
modes = {}
class FakeAttention(nn.Module):
def __init__(self, *, prefix, sequence_parallel_mode="local", **_kwargs):
super().__init__()
modes[prefix.rsplit(".", 1)[-1]] = sequence_parallel_mode
class FakeFeedForward(nn.Module):
def __init__(self, *_args, **_kwargs):
super().__init__()
monkeypatch.setattr(ltx2_transformer, "LTX2Attention", FakeAttention)
monkeypatch.setattr(ltx2_transformer, "LTX2FeedForward", FakeFeedForward)
LTX2VideoTransformerBlock(
dim=8,
num_attention_heads=2,
attention_head_dim=4,
cross_attention_dim=8,
audio_dim=8,
audio_num_attention_heads=2,
audio_attention_head_dim=4,
audio_cross_attention_dim=8,
prefix="transformer_blocks.0",
)
assert modes == {
"attn1": "video_self",
"audio_attn1": "local",
"attn2": "local",
"audio_attn2": "local",
"audio_to_video_attn": "local",
"video_to_audio_attn": "video_to_audio",
}
def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):
@@ -0,0 +1,171 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Full-transformer numeric parity for LTX Ulysses sequence parallelism.
This FP32/SDPA check verifies that Ulysses preserves the complete transformer
mathematical path. It does not guarantee pixel-level E2E parity: production
BF16 and compiled attention kernels can change rounding order, and diffusion
can amplify that drift across denoising steps.
"""
from __future__ import annotations
import os
import pytest
import torch
from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.config import set_current_diffusion_config
from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.parallel_state import (
destroy_distributed_env,
init_distributed_environment,
initialize_model_parallel,
)
from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.parallel]
_FP32_REL_L2 = 1e-5
_TINY_CONFIG = {
"in_channels": 4,
"out_channels": 4,
"num_attention_heads": 2,
"attention_head_dim": 4,
"cross_attention_dim": 8,
"vae_scale_factors": (1, 1, 1),
"pos_embed_max_pos": 4,
"base_height": 2,
"base_width": 2,
"audio_in_channels": 4,
"audio_out_channels": 4,
"audio_num_attention_heads": 2,
"audio_attention_head_dim": 4,
"audio_cross_attention_dim": 8,
"audio_scale_factor": 1,
"audio_pos_embed_max_pos": 4,
"audio_sampling_rate": 16,
"audio_hop_length": 4,
"num_layers": 1,
"caption_channels": 8,
"rope_type": "split",
"use_prompt_embeddings": False,
"use_keyframes_abs_pos_embedding": True,
}
def _relative_l2(actual: torch.Tensor, expected: torch.Tensor) -> float:
return ((actual - expected).double().norm() / expected.double().norm()).item()
def _omni_config(sp_size: int) -> OmniDiffusionConfig:
return OmniDiffusionConfig.from_kwargs(
model="test",
dtype=torch.float32,
parallel_config=DiffusionParallelConfig(
pipeline_parallel_size=1,
data_parallel_size=1,
tensor_parallel_size=1,
sequence_parallel_size=sp_size,
ulysses_degree=sp_size,
ring_degree=1,
cfg_parallel_size=1,
),
diffusion_attention_backend="TORCH_SDPA",
)
def _inputs(device: torch.device) -> dict[str, torch.Tensor | int | bool]:
generator = torch.Generator().manual_seed(123)
return {
"hidden_states": torch.randn(1, 8, 4, generator=generator, device="cpu").to(device),
"audio_hidden_states": torch.randn(1, 4, 4, generator=generator, device="cpu").to(device),
"encoder_hidden_states": torch.randn(1, 3, 8, generator=generator, device="cpu").to(device),
"audio_encoder_hidden_states": torch.randn(1, 3, 8, generator=generator, device="cpu").to(device),
"timestep": torch.linspace(0.1, 0.8, 8, device=device).view(1, 8),
"audio_timestep": torch.linspace(0.2, 0.5, 4, device=device).view(1, 4),
"keyframes_mask": torch.tensor([[[1.0], [0.0], [1.0], [0.0], [0.0], [1.0], [0.0], [0.0]]], device=device),
"sigma": torch.tensor([0.1], device=device),
"audio_sigma": torch.tensor([0.2], device=device),
"num_frames": 2,
"height": 2,
"width": 2,
"audio_num_frames": 4,
"return_dict": False,
}
def _run_transformer_parity(rank: int, world_size: int, master_port: int) -> None:
device = torch.device(f"{current_omni_platform.device_type}:{rank}")
current_omni_platform.set_device(device)
os.environ.update(
{
"RANK": str(rank),
"LOCAL_RANK": str(rank),
"WORLD_SIZE": str(world_size),
"MASTER_ADDR": "localhost",
"MASTER_PORT": str(master_port),
}
)
try:
init_distributed_environment()
initialize_model_parallel(sequence_parallel_size=world_size, ulysses_degree=world_size)
sp_config = _omni_config(world_size)
torch.manual_seed(42)
with set_forward_context(omni_diffusion_config=sp_config), set_current_diffusion_config(sp_config):
model = LTX2VideoTransformer3DModel(**_TINY_CONFIG).to(device).eval()
with torch.no_grad():
for _, parameter in sorted(model.named_parameters()):
torch.nn.init.normal_(parameter, mean=0.0, std=0.02)
inputs = _inputs(device)
reference_config = _omni_config(1)
with torch.no_grad(), set_forward_context(omni_diffusion_config=reference_config):
expected_video, expected_audio = model(**inputs)
apply_sequence_parallel(
model,
SequenceParallelConfig(ulysses_degree=world_size, ring_degree=1),
model._sp_plan,
)
with torch.no_grad(), set_forward_context(omni_diffusion_config=sp_config):
get_forward_context().sp_plan_hooks_applied = True
actual_video, actual_audio = model(**inputs)
v2a_attention = model.transformer_blocks[0].video_to_audio_attn.attn
assert v2a_attention._video_to_audio_strategy is not None
assert torch.isfinite(expected_video).all()
assert torch.isfinite(expected_audio).all()
for modality, actual, expected in (
("video", actual_video, expected_video),
("audio", actual_audio, expected_audio),
):
drift = _relative_l2(actual, expected)
assert drift <= _FP32_REL_L2, (
f"LTX SP{world_size} {modality} output drifted rel_l2={drift:.3e} from SP1 (bound {_FP32_REL_L2:.1e})"
)
finally:
destroy_distributed_env()
@pytest.mark.full_model
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
def test_ltx_ulysses_transformer_matches_sp1(unused_tcp_port) -> None:
"""Exercise real plan hooks, RoPE/timestep sharding, and attention routing."""
world_size = 2
device_count = current_omni_platform.device_count()
assert device_count >= world_size, (
f"LTX Ulysses parity requires {world_size} accelerator devices; only {device_count} are visible"
)
torch.multiprocessing.spawn(
_run_transformer_parity,
args=(world_size, unused_tcp_port),
nprocs=world_size,
)
@@ -3,6 +3,7 @@
import json
import sys
from contextlib import contextmanager
from multiprocessing.reduction import ForkingPickler
from types import SimpleNamespace
from typing import Any, TypeVar
@@ -13,6 +14,8 @@ import pytest
import torch
import torch.nn as nn
from vllm_omni.diffusion.models.minimax_h3.lora import TurboSpec, parse_turbo_filename
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
_ItemT = TypeVar("_ItemT")
@@ -24,6 +27,13 @@ def _append_and_return(items: list[_ItemT], item: _ItemT, result: _ResultT) -> _
return result
def _turbo_spec(filename: str) -> TurboSpec:
"""Build the spec a loaded artifact of this name would carry."""
spec = parse_turbo_filename(filename)
assert spec is not None
return spec
def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -167,9 +177,9 @@ def _write_partition_index(path, *, partition, tasks):
def test_modular_diffusers_index_is_resolved_generically(tmp_path):
from vllm_omni.config.resolver import resolve_omni_config
from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
from vllm_omni.entrypoints.utils import resolve_model_config_path
(tmp_path / "modular_model_index.json").write_text(
json.dumps(
@@ -183,7 +193,16 @@ def test_modular_diffusers_index_is_resolved_generically(tmp_path):
assert is_diffusion_model(str(tmp_path))
assert resolve_model_class_name(str(tmp_path)) == "MiniMaxH3ModularPipeline"
assert resolve_model_config_path(str(tmp_path)) is None
resolved = resolve_omni_config(
str(tmp_path),
trust_remote_code=False,
deploy_config_path=None,
cli_overrides=None,
stage_overrides=None,
strategy_config_path=None,
)
assert resolved.config_path is None
assert resolved.stage_by_id(0).engine_args.model_class_name == "MiniMaxH3ModularPipeline"
config = OmniDiffusionConfig(model=str(tmp_path))
config.enrich_config()
@@ -244,9 +263,12 @@ def test_combined_task_inference_and_transformer_routing():
assert pipeline._resolve_task(None, {"image": object()}) == "fl2va"
assert pipeline._resolve_task(None, {"audio": object()}) == "ref2va"
assert pipeline._resolve_task(None, {"video": object()}) == "ref2va"
assert pipeline._resolve_task("fl2va", {"image": object()}, has_turbo_lora=True) == "fl2va"
with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"):
pipeline._resolve_task("ref2va", {}, has_turbo_lora=True)
fl2v_spec = _turbo_spec("minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors")
ref2v_spec = _turbo_spec("minimax_h3_ref2v_turbo_8step_v1.0_768p_bf16.safetensors")
assert pipeline._resolve_task("fl2va", {"image": object()}, turbo_spec=fl2v_spec) == "fl2va"
with pytest.raises(OmniClientError, match="serves"):
pipeline._resolve_task("ref2va", {}, turbo_spec=fl2v_spec)
assert pipeline._resolve_task("ref2va", {"image": object()}, turbo_spec=ref2v_spec) == "ref2va"
pipeline.partition = "ref2va"
pipeline.supported_tasks = frozenset({"ref2va"})
assert pipeline._resolve_task(None, {"image": object()}) == "ref2va"
@@ -265,14 +287,16 @@ def test_combined_task_inference_and_transformer_routing():
"component_partition",
"source_partitions",
"expected_tasks",
"load_text_encoder",
),
[
(None, "combined", 2, "FL2VA", ["FL2VA", "Ref2VA"], {"t2va", "fl2va", "ref2va"}),
("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}),
("ref2va", "ref2va", 1, "Ref2VA", ["Ref2VA"], {"ref2va"}),
(None, "combined", 2, "FL2VA", ["FL2VA", "Ref2VA"], {"t2va", "fl2va", "ref2va"}, True),
("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}, True),
("ref2va", "ref2va", 1, "Ref2VA", ["Ref2VA"], {"ref2va"}, True),
("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}, False),
],
)
def test_pipeline_loads_task_selected_dits_and_shared_components_once(
def test_pipeline_loads_task_selected_components_with_encoder_ownership(
monkeypatch,
tmp_path,
task_type,
@@ -281,6 +305,7 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
component_partition,
source_partitions,
expected_tasks,
load_text_encoder,
):
from vllm_omni.diffusion.data import (
DiffusionParallelConfig,
@@ -392,6 +417,11 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
revision=None,
task_type=task_type,
quantization_config=None,
model_loaded={
"transformer": True,
"vae": True,
"text_encoder": load_text_encoder,
},
parallel_config=DiffusionParallelConfig(
cfg_parallel_size=1,
text_encoder_tp_size=1,
@@ -403,34 +433,54 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
assert pipeline.supported_tasks == expected_tasks
assert len(created["dit"]) == expected_dits
component_path = tmp_path / component_partition
assert created["text_encoder"] == [str(component_path / "text_encoder")]
assert created["video_vae"] == [str(component_path / "video_vae")]
assert created["audio_vae"] == [str(component_path / "audio_vae")]
assert len(tokenizer_calls) == 1
assert len(processor_calls) == 1
expected_dit_modules = ["transformer", "transformers_ref"] if expected_dits == 2 else ["transformer"]
assert pipeline._dit_modules == expected_dit_modules
assert hasattr(pipeline, "transformers_ref") is (expected_dits == 2)
if load_text_encoder:
assert created["text_encoder"] == [str(component_path / "text_encoder")]
assert len(tokenizer_calls) == 1
assert len(processor_calls) == 1
assert pipeline.text_encoder is not None
assert pipeline.tokenizer is not None
assert pipeline.processor is not None
assert pipeline._encoder_modules == ["text_encoder"]
assert pipeline.text_encoder_tp_size == 1
else:
assert created["text_encoder"] == []
assert tokenizer_calls == []
assert processor_calls == []
assert pipeline.text_encoder is None
assert pipeline.tokenizer is None
assert pipeline.processor is None
assert pipeline._encoder_modules == []
assert pipeline.text_encoder_tp_size == 0
assert pipeline.text_encoder_group is None
if expected_partition == "ref2va":
assert pipeline._transformer_for_task("ref2va") is pipeline.transformer
assert [source.model_or_path for source in pipeline.weights_sources] == [
*(str(tmp_path / partition_name) for partition_name in source_partitions),
str(component_path),
*((str(component_path),) if load_text_encoder else ()),
]
assert [source.subfolder for source in pipeline.weights_sources] == [
*("transformer" for _ in source_partitions),
"text_encoder",
*(("text_encoder",) if load_text_encoder else ()),
]
expected_prefixes = ["transformer."]
if expected_dits == 2:
expected_prefixes.append("transformers_ref.")
expected_prefixes.append("text_encoder.")
if load_text_encoder:
expected_prefixes.append("text_encoder.")
assert [source.prefix for source in pipeline.weights_sources] == expected_prefixes
expected_patterns = (
pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS
if expected_partition == "combined"
else pipeline_module.MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[expected_partition]
)
if load_text_encoder:
expected_patterns = (
pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS
if expected_partition == "combined"
else pipeline_module.MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[expected_partition]
)
else:
expected_patterns = pipeline_module.MINIMAX_H3_DIFFUSION_DOWNLOAD_PATTERNS[expected_partition]
assert download_calls == [
{
"model_name_or_path": "MiniMaxAI/MiniMax-H3",
@@ -1510,6 +1560,97 @@ def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
assert next(video_vae.parameters()).dtype == torch.float32
def test_keyframe_encode_pins_and_restores_cudnn_settings():
from vllm_omni.diffusion.models.minimax_h3.vae import (
_minimax_h3_keyframe_encode_context,
)
cudnn = torch.backends.cudnn
original = (
cudnn.enabled,
cudnn.benchmark,
cudnn.deterministic,
cudnn.allow_tf32,
)
try:
cudnn.enabled = False
cudnn.benchmark = True
cudnn.deterministic = False
cudnn.allow_tf32 = False
with _minimax_h3_keyframe_encode_context(torch.device("cuda")):
assert cudnn.enabled
assert not cudnn.benchmark
assert cudnn.deterministic
assert cudnn.allow_tf32
assert not cudnn.enabled
assert cudnn.benchmark
assert not cudnn.deterministic
assert not cudnn.allow_tf32
finally:
(
cudnn.enabled,
cudnn.benchmark,
cudnn.deterministic,
cudnn.allow_tf32,
) = original
@pytest.mark.parametrize("fail_encode", [False, True])
def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode):
from PIL import Image
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
events = []
input_image = Image.new("RGB", (32, 32))
@contextmanager
def track_context(device):
assert device == torch.device("cpu")
events.append("enter")
try:
yield
finally:
events.append("exit")
class FakeModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(torch.ones(1))
self.parallel_tiling = True
def encode_images(self, image, *, use_fp16_latent):
assert events == ["enter"]
assert image is input_image
assert use_fp16_latent
events.append("encode")
if fail_encode:
raise RuntimeError("encode failed")
return [torch.ones(1, 1, 2, 2)]
video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE)
torch.nn.Module.__init__(video_vae)
video_vae.model = FakeModel()
video_vae.config_dict = {
"latent_channels": 1,
"latents_mean": [0.0],
"latents_std": [1.0],
}
monkeypatch.setattr(vae_module, "_minimax_h3_keyframe_encode_context", track_context)
if fail_encode:
with pytest.raises(RuntimeError, match="encode failed"):
video_vae.encode_image(input_image)
else:
rows = video_vae.encode_image(input_image)
torch.testing.assert_close(rows, torch.ones(1, 4))
assert events == ["enter", "encode", "exit"]
assert video_vae.model.parallel_tiling
def test_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
Loaded 100 of 439 files, more files were not shown because too many files have changed in this diff. Show more