Compare commits

..
Author SHA1 Message Date
Hongsheng Liu 65c86dec44 [Docs] Clarify TP1 DLO mmap support scope
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-16 10:54:24 +08:00
Hongsheng Liu 5ce549bc57 Merge branch 'main' into codex-dlo-host-weight-plan 2026-08-16 00:13:33 +08:00
yiminghub2024andHongsheng Liu 444485650b Add MiniMax-H3 recipe for NPU 950PR (#6120)
Signed-off-by: yiminghub2024 <482890@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-16 00:04:38 +08:00
hurukawa 2a315e1a2d [Perf] Support request-level batching for Wan2.2 pipelines (#5676)
Signed-off-by: nagisa-kun <1434936049@qq.com>
Signed-off-by: nagisa <1434936049@qq.com>
2026-08-15 23:58:36 +08:00
Hongsheng Liu fa2f5b4571 [Diffusion] Preserve non-DiT loading with DLO plans
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-15 23:57:03 +08:00
Hongsheng Liu 11687f73bd [Diffusion] Rebase DLO plan and generalize lifecycle example
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-15 23:40:24 +08:00
Hongsheng Liu 88614175c0 [Diffusion] Address DLO host-weight review feedback
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-15 23:32:27 +08:00
Hongsheng Liu e2796b6168 [Diffusion] Add loader-owned host-weight plans for DLO
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-15 23:32:27 +08:00
SYLARandHongsheng Liu d1e230c95b [Diffusion][Quantization] Enable MiniMax-H3 global FP8 with DLO (#5910)
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-15 23:01:26 +08:00
m0g3r 09b6a1604a [Bugfix] Restore per-stage runtime env during launch (#6214)
Signed-off-by: m0g3r <87276771+m0g3r@users.noreply.github.com>
2026-08-15 13:00:27 +00:00
Yueqian Lin 555336e887 [CI] Refresh Voxtral-4B-TTS perf baselines from a clean window (#6032)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-08-15 11:51:56 +00:00
Samit 136cb27b05 [Bugfix] Fix lost async diffusion outputs in request-level batching (#6023)
Signed-off-by: samithuang <285365963@qq.com>
2026-08-15 18:27:51 +08:00
SYLAR ce6ab0ca46 [Model] Add LTX-2.5 support (#6070)
Signed-off-by: lishunyang12 <lishunyang12@163.com>
2026-08-15 18:27:23 +08:00
NumberWan f39b8684eb [BugFix][Nightly CI] Adjust Qwen-Image accuracy thresholds for stable FA-deterministic mode (#5963)
Signed-off-by: NumberWan <wantszkin2003@gmail.com>
2026-08-15 09:38:30 +00:00
MrlixiangWE e57574ae29 [Bugfix] Carry ec_transfer_params and num_cache_creation_tokens on OmniRequestOutput (#6152)
Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>
2026-08-15 08:37:27 +00:00
7b76b6446d [Model][Performance] Optimize MiniMax-H3 strict Ulysses boundaries (#6173)
Signed-off-by: mokashliu <mokashliu@tencent.com>
Co-authored-by: mokashliu <mokashliu@tencent.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-08-15 16:27:52 +08:00
a70d4d4b3c [codex] Fix diffusion worker timeout for large broadcast payloads (#4845)
Signed-off-by: jiangweifeng.dev <jiangweifeng.dev@bytedance.com>
Co-authored-by: jiangweifeng.dev <jiangweifeng.dev@bytedance.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-15 16:03:04 +08:00
Hongsheng Liu 1260ceaae6 [Docs] Split diffusion attention and CPU offload guides (#6075)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-15 14:41:09 +08:00
Ma1oneZhang 6f36fd3b78 [Bugfix] Preserve object storage URIs during model resolution (#5036)
Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
2026-08-15 14:11:48 +08:00
35cacf020c [Model] Add Gepard-1.0 native-AR FSQ/NanoCodec TTS (offline inference) (#5666)
Signed-off-by: Mengjie Zhao <zmj0129@gmail.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-08-14 22:06:05 -07:00
wangyu d671e6b0db [CI] Expand E2E source_file_dependencies for shared model code (#5994)
Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com>
2026-08-15 12:01:20 +08:00
MrlixiangWE b95c73742d [Bugfix] Stop reading the removed OmniRequestOutput.request_output accessor (#6172)
Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>
2026-08-15 07:07:28 +08:00
Alex Brooks fe54ecc388 [Frontend] Loosen Validation for TTS Models with No Uploaded Speakers (#5878)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-08-14 17:05:11 -06:00
165 changed files with 14089 additions and 1759 deletions

No files matched your search

+26 -91
View File
@@ -76,6 +76,9 @@ steps:
- tests/e2e/online_serving/test_qwen3_tts_customvoice.py
- tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
- vllm_omni/model_executor/models/qwen3_tts/
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
- vllm_omni/model_executor/models/common/snake_activation.py
- vllm_omni/model_executor/models/common/alias_free_activation.py
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
- vllm_omni/deploy/qwen3_tts.yaml
commands:
@@ -92,6 +95,9 @@ steps:
- tests/e2e/online_serving/test_qwen3_tts_base.py
- tests/e2e/offline_inference/test_qwen3_tts_base.py
- vllm_omni/model_executor/models/qwen3_tts/
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
- vllm_omni/model_executor/models/common/snake_activation.py
- vllm_omni/model_executor/models/common/alias_free_activation.py
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
- vllm_omni/deploy/qwen3_tts.yaml
- .buildkite/common/scripts/run_cov_split.sh
@@ -109,27 +115,15 @@ steps:
-- -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: l4_1
- label: "TTS · Higgs-Audio-V3 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_higgs_audio_v3.py
- vllm_omni/model_executor/models/higgs_audio_v3/
- vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py
- vllm_omni/deploy/higgs_multimodal_qwen3.yaml
commands:
- |
timeout 20m bash -c "
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_USE_DEEP_GEMM=0
export VLLM_MOE_USE_DEEP_GEMM=0
pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'advanced_model' --run-level 'advanced_model'
"
mirror_hardwares: l4_1
- label: "Omni · Qwen3-Omni Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_qwen3_omni.py
- tests/e2e/online_serving/test_qwen3_omni.py
- vllm_omni/model_executor/models/qwen3_omni/
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
- vllm_omni/model_executor/models/common/snake_activation.py
- vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
- vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
- vllm_omni/deploy/qwen3_omni_moe.yaml
timeout_in_minutes: 50
@@ -141,6 +135,11 @@ steps:
source_file_dependencies:
- tests/e2e/offline_inference/test_minicpmo_4_5.py
- tests/e2e/online_serving/test_minicpmo_4_5.py
- vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/models/step_audio2/
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
- vllm_omni/model_executor/models/cosyvoice3/utils.py
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
- vllm_omni/deploy/minicpmo_4_5.yaml
timeout_in_minutes: 50
commands:
@@ -162,70 +161,35 @@ steps:
- tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
- vllm_omni/deploy/minicpmo_4_5.yaml
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
- vllm_omni/experimental/fullduplex/
- vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/models/step_audio2/
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
- vllm_omni/model_executor/models/cosyvoice3/utils.py
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
- vllm_omni/core/sched/
- vllm_omni/outputs/multimodal_accumulation.py
- vllm_omni/engine/stage_init_utils.py
- vllm_omni/engine/async_omni_engine.py
- vllm_omni/engine/orchestrator.py
- vllm_omni/entrypoints/cli/serve.py
- vllm_omni/entrypoints/openai/api_server.py
- vllm_omni/entrypoints/openai/duplex_capability.py
- vllm_omni/config/stage_config.py
timeout_in_minutes: 50
commands:
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: h100_1
- label: "TTS · VoxCPM2 Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_voxcpm2_tts.py
- tests/e2e/online_serving/test_voxcpm2_tts.py
- vllm_omni/model_executor/models/voxcpm2/
- vllm_omni/deploy/voxcpm2.yaml
timeout_in_minutes: 30
commands:
- |
timeout 30m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
"
mirror_hardwares: l4_1
- label: "Diffusion · Qwen Image Edit Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image_edit.py
- vllm_omni/diffusion/models/qwen_image/
timeout_in_minutes: 20
commands:
- pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: h100_1
- label: "Diffusion · Bagel Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_bagel.py
- tests/e2e/online_serving/test_bagel.py
- vllm_omni/model_executor/models/bagel/
- vllm_omni/diffusion/models/bagel/
- vllm_omni/model_executor/stage_input_processors/bagel.py
- vllm_omni/deploy/bagel.yaml
- .buildkite/common/scripts/run_cov_split.sh
- pyproject.toml
# 20 -> 40: the coverage split runs the model twice.
timeout_in_minutes: 40
commands:
- |
.buildkite/common/scripts/run_cov_split.sh \
--model-id bagel \
--offline tests/e2e/offline_inference/test_bagel.py \
--online tests/e2e/online_serving/test_bagel.py \
-- -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: h100_1
- label: "Diffusion · Qwen Image Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image.py
- vllm_omni/diffusion/models/qwen_image/
- vllm_omni/diffusion/models/dmd2/
- vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py
commands:
- |
timeout 40m bash -c '
@@ -234,23 +198,12 @@ steps:
'
mirror_hardwares: h100_1
- label: "Diffusion · Qwen Image Layered Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image_layered.py
- vllm_omni/diffusion/models/qwen_image/
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'advanced_model and cuda' --run-level 'advanced_model'
"
mirror_hardwares: h100_1
- label: "Diffusion · Wan22 Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_wan22_t2v.py
- tests/e2e/online_serving/test_wan22_t2v.py
- vllm_omni/diffusion/models/wan2_2/
- vllm_omni/diffusion/models/dmd2/
commands:
- |
timeout 40m bash -c "
@@ -258,21 +211,3 @@ steps:
pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
"
mirror_hardwares: h100_1
- label: "Diffusion · Z Image Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_z_image.py
- tests/e2e/online_serving/test_z_image.py
- vllm_omni/diffusion/models/z_image/
timeout_in_minutes: 30
commands:
- pytest -s -v tests/e2e/offline_inference/test_z_image.py tests/e2e/online_serving/test_z_image.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
mirror_hardwares: l4_1
- label: "Diffusion · FLUX.2 Klein Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_flux2_klein.py
- vllm_omni/diffusion/models/flux2_klein/
timeout_in_minutes: 30
commands:
- pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
mirror_hardwares: l4_1
+13
View File
@@ -459,6 +459,12 @@ steps:
- pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx.py -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · LTX-2.5 Online Test"
timeout_in_minutes: 60
commands:
- pytest -s -v tests/e2e/online_serving/test_ltx25.py -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test"
timeout_in_minutes: 180
commands:
@@ -487,6 +493,13 @@ steps:
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx23_two_stage_layer_fused_i2v]' -m full_model --run-level full_model
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2V · LTX-2.5 Official Parity Test"
timeout_in_minutes: 180
commands:
- pytest -s -v tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py::test_ltx25_distilled_two_stage_matches_official -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py::test_ltx25_full_matches_official -m full_model --run-level full_model
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2V · Perf Test"
key: nightly-diffusion-x2v-performance
timeout_in_minutes: 180
+34 -69
View File
@@ -169,6 +169,10 @@ steps:
source_file_dependencies:
- tests/e2e/online_serving/test_qwen3_omni.py
- vllm_omni/model_executor/models/qwen3_omni/
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
- vllm_omni/model_executor/models/common/snake_activation.py
- vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
- vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
- vllm_omni/deploy/qwen3_omni_moe.yaml
commands:
@@ -181,6 +185,11 @@ steps:
- label: "Omni · MiniCPM-o 4.5 Offline Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_minicpmo_4_5.py
- vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/models/step_audio2/
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
- vllm_omni/model_executor/models/cosyvoice3/utils.py
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
- vllm_omni/deploy/minicpmo_4_5.yaml
commands:
- |
@@ -192,6 +201,11 @@ steps:
- label: "Omni · MiniCPM-o 4.5 Online Test"
source_file_dependencies:
- tests/e2e/online_serving/test_minicpmo_4_5.py
- vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/models/step_audio2/
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
- vllm_omni/model_executor/models/cosyvoice3/utils.py
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
- vllm_omni/deploy/minicpmo_4_5.yaml
commands:
- |
@@ -210,16 +224,24 @@ steps:
- tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
- vllm_omni/deploy/minicpmo_4_5.yaml
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
- vllm_omni/experimental/fullduplex/
- vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/models/step_audio2/
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
- vllm_omni/model_executor/models/cosyvoice3/utils.py
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
- vllm_omni/core/sched/
- vllm_omni/outputs/multimodal_accumulation.py
- vllm_omni/engine/stage_init_utils.py
- vllm_omni/engine/async_omni_engine.py
- vllm_omni/engine/orchestrator.py
- vllm_omni/entrypoints/cli/serve.py
- vllm_omni/entrypoints/openai/api_server.py
- vllm_omni/entrypoints/openai/duplex_capability.py
- vllm_omni/config/stage_config.py
commands:
- |
timeout 20m bash -c "
@@ -239,25 +261,13 @@ steps:
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and cuda'
mirror_hardwares: h100_1
- label: "TTS · VoxCPM2 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_voxcpm2_tts.py
- vllm_omni/model_executor/models/voxcpm2/
- vllm_omni/deploy/voxcpm2.yaml
timeout_in_minutes: 30
commands:
- |
timeout 30m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: l4_1
- label: "TTS · Qwen3-TTS CustomVoice Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen3_tts_customvoice.py
- vllm_omni/model_executor/models/qwen3_tts/
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
- vllm_omni/model_executor/models/common/snake_activation.py
- vllm_omni/model_executor/models/common/alias_free_activation.py
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
- vllm_omni/deploy/qwen3_tts.yaml
commands:
@@ -272,6 +282,9 @@ steps:
source_file_dependencies:
- tests/e2e/online_serving/test_qwen3_tts_base.py
- vllm_omni/model_executor/models/qwen3_tts/
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
- vllm_omni/model_executor/models/common/snake_activation.py
- vllm_omni/model_executor/models/common/alias_free_activation.py
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
- vllm_omni/deploy/qwen3_tts.yaml
commands:
@@ -284,52 +297,11 @@ steps:
"
mirror_hardwares: l4_1
- label: "TTS · Higgs-Audio-V3 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_higgs_audio_v3.py
- vllm_omni/model_executor/models/higgs_audio_v3/
- vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py
- vllm_omni/deploy/higgs_multimodal_qwen3.yaml
commands:
- |
timeout 20m bash -c "
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_USE_DEEP_GEMM=0
export VLLM_MOE_USE_DEEP_GEMM=0
pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: l4_1
- label: "Diffusion · Qwen Image Edit Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image_edit.py
- vllm_omni/diffusion/models/qwen_image/
commands:
- |
timeout 20m bash -c "
pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: h100_1
- label: "Diffusion · Bagel Test"
source_file_dependencies:
- tests/e2e/online_serving/test_bagel.py
- vllm_omni/model_executor/models/bagel/
- vllm_omni/diffusion/models/bagel/
- vllm_omni/model_executor/stage_input_processors/bagel.py
- vllm_omni/deploy/bagel.yaml
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_bagel.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: h100_1
- label: "Diffusion · Wan22 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_wan22_t2v.py
- vllm_omni/diffusion/models/wan2_2/
- vllm_omni/diffusion/models/dmd2/
commands:
- |
timeout 40m bash -c "
@@ -342,6 +314,9 @@ steps:
source_file_dependencies:
- tests/e2e/online_serving/test_cosmos3.py
- vllm_omni/diffusion/models/cosmos3/
- vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py
- vllm_omni/experimental/world_models/
- vllm_omni/model_extras/cosmos3.py
commands:
- |
timeout 40m bash -c "
@@ -354,6 +329,8 @@ steps:
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image.py
- vllm_omni/diffusion/models/qwen_image/
- vllm_omni/diffusion/models/dmd2/
- vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py
commands:
- |
timeout 40m bash -c "
@@ -361,15 +338,3 @@ steps:
pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: h100_1
- label: "Diffusion · Qwen Image Layered Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image_layered.py
- vllm_omni/diffusion/models/qwen_image/
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: h100_1
+16
View File
@@ -108,6 +108,22 @@ steps:
--run-level "full_model"
mirror_hardwares: h100_1
# Own step: the image has no NanoCodec, so this file import-skips out of
# the mark-selected TTS steps above. The overrides keep nemo_toolkit[tts]
# from downgrading transformers; see
# examples/offline_inference/text_to_speech/README.md.
- label: "TTS · Gepard-1.0 · L4"
timeout_in_minutes: 60
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0" > /tmp/gepard-overrides.txt
- uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" "nemo_toolkit[tts]==2.7.3" --overrides /tmp/gepard-overrides.txt
- >-
pytest -s -v tests/e2e/offline_inference/test_gepard_tts.py
-m "tts and cuda"
--run-level "full_model"
mirror_hardwares: l4_1
- label: "Diffusion · L4"
timeout_in_minutes: 60
commands:
+1
View File
@@ -254,6 +254,7 @@ outputs/
results/
generated/
output_*/
/tests/e2e/accuracy/artifacts/
# Configuration overrides
configs/local.yaml
+9 -6
View File
@@ -329,11 +329,13 @@ On PR builds, `upload_pipeline.py` logs `skip '…' (no changes under …)` for
### Per-model coverage
Pilot (v1): two Merge-tier (L3) jobs in `.buildkite/cuda/test-merge.yml` upload
per-model, per-entry-mode coverage as Buildkite artifacts — "Diffusion · Bagel
Test" and "TTS · Qwen3-TTS Base Test". Both run on a single GPU so the pilot is
cheap to reproduce. This is a pilot, not a rollout: no dashboard/visualization
lives in this repo, and nothing is gated on coverage.
Pilot (v1): one Merge-tier (L3) job in `.buildkite/cuda/test-merge.yml` uploads
per-model, per-entry-mode coverage as Buildkite artifacts — "TTS · Qwen3-TTS Base
Test". It runs on a single GPU so the pilot is cheap to reproduce. (A second
diffusion pilot, "Diffusion · Bagel Test", was retired from merge; reintroduce
another diffusion coverage job if a second pilot is needed again.) This is a
pilot, not a rollout: no dashboard/visualization lives in this repo, and nothing
is gated on coverage.
**Naming convention**: `coverage-<model_id>-<mode>-<step_id>.xml.gz`, where
`<model_id>` is the model's directory name under
@@ -410,7 +412,8 @@ runs a coverage job needs the same.
List both `run_cov_split.sh` and `pyproject.toml` in every opted-in job's
`source_file_dependencies` — both change what the job measures, so without them a
change there is filtered out of normal PR builds and only surfaces in a later
nightly. `tests/buildkite/test_upload_pipeline.py` asserts this for the pilots.
nightly. `tests/buildkite/test_upload_pipeline.py` covers the filter behavior with
a synthetic job (it does not pin real merge labels).
Editing only the surrounding CI YAML still does not schedule them, so a PR that
touches just the wiring needs a full E2E run (or the commands run on a GPU host)
to produce artifacts. When checking a new model's
+17 -1
View File
@@ -117,7 +117,22 @@ After (async D2H):
### Batch Split
When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently.
When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently. If `OUTPUT_READY` arrives before `execute_batch` registers the split map, `execute_batch` directly adopts the early output.
### Reliability, Lifecycle & Timeout Behavior
1. **Drain Before Memory Release (`drain_async_outputs` / `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S`)**:
When workers transition into sleep states (`handle_sleep_task`) or invoke memory-releasing methods, `WorkerProc.drain_async_outputs()` blocks up to `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S` (10.0s) until all in-flight background D2H/SHM packing tasks finish. This acts as a synchronization barrier preventing worker sleep cycles from releasing device tensors while the side CUDA stream is still actively reading them. Draining typically completes within 10–50 ms under normal operation.
2. **Timeout Diagnostics (`_ASYNC_OUTPUT_TIMEOUT`)**:
The engine waits for `wait_output_ready()` with a timeout of `_ASYNC_OUTPUT_TIMEOUT` (30.0s). If a timeout occurs, the executor logs detailed diagnostic bookkeeping (pending futures, cached outputs, and batch split maps) to identify whether the stall occurred in the worker, pump, or waiter.
3. **Queue Serialization & Atomic Future Resolution**:
- `result_mq` writes are serialized with `_result_mq_lock` to protect the single-writer `MessageQueue` against concurrent writes from the worker main loop (`COMPUTE_DONE`) and background thread (`OUTPUT_READY`).
- `ResultPumpThread` unpacks shared-memory payloads first, and then atomically resolves waiting futures or populates `_completed_outputs` under `_futures_lock`, avoiding race conditions with concurrent `wait_output_ready()` invocations.
4. **Storage-Aware IPC Packing**:
`pack_diffusion_output_shm()` evaluates `max(view_bytes, storage_bytes)` against `_SHM_TENSOR_THRESHOLD` (1 MB) to ensure that per-request slices sharing a large batch storage are moved to POSIX shared memory rather than serialized through pickle over the MessageQueue inline buffer.
## Model Coverage
@@ -131,6 +146,7 @@ All models running in request-mode (`step_execution=False`, the default) automat
|---|---|---|
| **HunyuanImage-3.0** | Image | `False` |
| **Qwen-Image** | Image | `True` |
| **LTX-2.3** | Video / Audio | `False` |
Other models with `step_execution=False` are also supported but not yet verified.
@@ -4,7 +4,7 @@ This document describes distributed layerwise offload (DLO) for diffusion
models. DLO keeps only a small number of DiT blocks on the accelerator and
streams the remaining blocks from host memory. The distributed backend can
either shard those host-side weights across an existing parallel group or keep
the standard loader's rank-local weights and avoid an additional collective.
complete rank-local block sources and avoid an additional collective.
For user-facing commands, see the
[distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
@@ -14,10 +14,20 @@ and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/rec
DLO is implemented for multi-device diffusion execution. The default
AllGather path is the primary path for DP and SP deployments. The
`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
for standard-loader sharding, workstation bring-up, and systems where an
additional DLO collective is undesirable, but it does not reduce host weight
storage across ranks.
`--dlo-no-use-allgather` path streams complete blocks independently and adds no
DLO weight collective.
Host storage is selected separately from the transfer protocol. The loader can
produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on
the same node can share immutable checkpoint pages when direct mmap is
selected, while the ordinary-loader fallback still keeps a private runtime
copy per process.
The Phase A shared-mmap support boundary is TP1. TP greater than one is an
ordinary-loader compatibility path: DLO can consume the resulting TP-local
tensors, but those configurations do not use checkpoint mmap and must not be
used to claim shared-mmap host-memory savings.
The compatibility matrix below describes the current implementation. The
unit-level guards are covered, but not every parallelism combination has a
@@ -41,6 +51,36 @@ TP is deliberately not used as DLO's AllGather group. HSDP has its own
parameter-sharding lifecycle and is not allowed to be sharded a second time by
DLO's AllGather path.
### The loader owns host-weight planning
Before it decides whether ordinary weight materialization can be skipped, the
diffusion loader builds one `HostWeightPlan`. A direct-checkpoint mmap plan is
accepted only when preflight proves all of the following:
- every required DiT parameter and persistent buffer has exactly one source;
- runtime names, checkpoint keys, shapes, and dtypes match;
- the runtime topology is TP1 without HSDP or online quantization; and
- every custom loader operation is represented by a loader-owned checkpoint
adapter.
The exact plan object is handed to DLO. The backend does not rescan checkpoint
files, repeat the capability decision, or reconstruct names from its block
topology. If preflight fails, the loader materializes weights normally and DLO
consumes those runtime tensors.
The plan owns only dedicated DiT component sources. If a pipeline also exposes
ordinary sources for a text encoder or another non-DiT component, the loader
still consumes those sources and includes their loaded names in its strict
coverage check. Only the source prefixes covered by the plan skip ordinary
materialization. A source that mixes DiT and non-DiT weights fails closed to
the complete ordinary-loader path because it cannot be skipped safely as a
unit.
This boundary keeps checkpoint semantics out of DLO and avoids model-pipeline
flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
transforms. Model-specific direct-layout knowledge, when required, lives in a
checkpoint adapter beside the ordinary loader.
### AllGather path
With the default `dlo_use_allgather=True`, each rank stores approximately
@@ -61,6 +101,15 @@ Buffers: [current slot] [prefetch slot] [current slot]
The backend uses two shared device buffers, so accelerator weight residency is
bounded by the largest streamed blocks rather than the complete model.
When direct checkpoint mmap is selected, the checkpoint mappings are only the
source used to prepare each rank's persistent shard. They can be closed after
shard preparation. Across the AllGather group, those private shards total
approximately one runtime model copy.
An effective DLO group size of one performs no collective, even when
`dlo_use_allgather=True`; it follows the rank-local transfer path described
below.
When DP is greater than one, the engine can process one request per DP rank in
the same denoising wave. Because AllGather is a collective, all participating
requests must take the same execution path at every denoising step.
@@ -68,9 +117,20 @@ requests must take the same execution path at every denoising step.
### Rank-local path without DLO AllGather
With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
one. The regular model loader remains responsible for preparing each rank's
weights, including TP-local tensors or HSDP-managed parameters. DLO then
streams those rank-local tensors block by block using H2D copies only.
one and streams complete blocks using H2D copies only. The host backing may be
either a loader-approved checkpoint mapping or ordinary runtime tensors.
For direct mmap, each process retains immutable safetensors views and uses two
bounded pinned host staging slots. Processes on the same node that map the same
files share physical checkpoint pages through the OS page cache. This removes
the persistent private full-model copy per pure-DP process, but each process
still packs and transfers every complete block. Sharing is node-local; each
node has its own page cache.
When direct mmap preflight fails, the regular model loader remains responsible
for preparing each rank's weights, including TP-local tensors or HSDP-managed
parameters. In that fallback, each pure-DP process keeps a private full runtime
copy.
This mode means:
@@ -80,21 +140,18 @@ This mode means:
not shard weights across SP ranks.
- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
DLO's additional weight AllGather is disabled.
- Pure DP deployments keep a full host-side model copy per rank, subject to
shared OS page-cache behavior.
- Pure DP deployments share one checkpoint-backed copy per node when direct
mmap is selected; the ordinary-loader fallback keeps one private runtime
copy per rank.
- The scheduler does not require a synchronized DP request wave for DLO.
This path is intentionally not implemented by reusing the AllGather mmap
loader. It relies on the standard loader so model-specific TP/HSDP transforms
remain intact.
## Parallelism compatibility
| Parallelism | DLO + AllGather | DLO without AllGather |
|---|---|---|
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. |
| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
| **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. |
| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
### Combined dimensions
@@ -122,11 +179,15 @@ AllGather DP multi-concurrency requires:
The no-AllGather path does not impose these DLO-specific synchronized-wave
requirements.
The mmap loader is used only by the supported DLO+AllGather path when the
model has a compatible checkpoint layout. Online quantization is incompatible
with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
quantization. Models that do not meet the mmap requirements use the regular
loader path.
Direct checkpoint mmap can back either transfer path. It is currently limited
to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
ordinary loader. Online quantization remains incompatible with DLO AllGather;
use `--dlo-no-use-allgather` or disable online quantization.
A normalized runtime mmap cache, built through the ordinary loader, is the
proposed general mechanism for sharing transformed TP or quantized layouts.
That cache and its publication/lifecycle protocol are intentionally outside
this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
## Validation coverage
@@ -134,15 +195,75 @@ Current source-level validation includes:
- HSDP + DLO + AllGather rejection;
- HSDP + DLO without AllGather acceptance at configuration level;
- TP rejection in the DLO+AllGather mmap path;
- loader preflight fallback for TP, HSDP, online quantization, unknown custom
loaders, missing keys, and shape/dtype mismatches;
- exact loader-to-backend plan transfer and ordinary-loader fallback;
- rank-local mmap source retention, bounded two-slot staging, and adapter
transforms without parameter-side flags;
- resident-layer requests requiring no-AllGather;
- DP request-wave validation for denoising-step compatibility;
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
tests.
The highest-value missing coverage is end-to-end numerical comparison against
ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
on the target CUDA/NCCL or CANN/HCCL hardware.
### B300 parallel-topology smoke matrix
A four-GPU B300 smoke test covered MiniMax-H3 FL2VA with the same prompt, seed,
CUDNN attention backend, 256x256 output, two denoising steps, and
`dlo_resident_layers=0`. The TP2 rows used DiT DP2xTP2 with the text encoder and
VAEs at TP1. They validate the ordinary-loader fallback only, not direct mmap
or shared-mmap host-memory savings.
| Configuration | Result | Warm E2E | Peak device memory | Host PSS |
|---|---:|---:|---:|---:|
| DP4xTP1 AllGather | Passed, 4 concurrent requests | 2.87 s / 4 requests | 13.84 GiB | 211.99 GiB |
| DP4xTP1 no-AllGather | Passed, 1 request | 15.02 s | 13.23 GiB | 187.77 GiB |
| DP2xTP2 AllGather | Passed, 2 concurrent requests | 4.16 s / 2 requests | 12.50 GiB | 211.97 GiB |
| DP2xTP2 no-AllGather | Passed, 1 request | 3.51 s | 11.88 GiB | 314.01 GiB |
Within each topology, the AllGather and no-AllGather video and audio outputs
were byte-identical. All four runs completed without an `ERROR` or traceback
and released their device allocations. For DP4xTP1, no-AllGather direct mmap
reduced total PSS by 24.22 GiB (11.4%) and `Private_Dirty` from 211.33 to
125.32 GiB (40.7%) relative to AllGather. For DP2xTP2, preflight selected the
ordinary loader as designed; no-AllGather PSS was 314.01 GiB, about 48% above
AllGather, because DP replicas did not share checkpoint-backed runtime
weights. This is a functional and memory smoke test, not a production-quality
performance or output-quality benchmark.
### Host-memory measurement
A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the
ordinary-loader fallback with direct mmap. Both runs used
DP=2, TP=1, no DLO AllGather, BF16 weights, two denoising steps, and a
256x256 four-second request. The ordinary-loader workers were sampled after
initialization. The mmap workers were sampled after one completed request, so
the checkpoint working set had been faulted into the page cache; this is the
more conservative point for mmap.
The values below come from `/proc/<worker>/smaps_rollup` and include the whole
worker, not only the DiT. The stable rank-to-rank difference comes from other
pipeline components, so each worker should be compared with the same worker in
the other storage mode.
| Worker | Ordinary RSS | mmap RSS | Ordinary PSS | mmap PSS | PSS reduction |
|---|---:|---:|---:|---:|---:|
| DP worker 0 | 168.27 GiB | 132.76 GiB | 167.84 GiB | 101.43 GiB | 66.40 GiB |
| DP worker 1 | 116.19 GiB | 79.97 GiB | 115.73 GiB | 48.64 GiB | 67.09 GiB |
| **Two-worker total** | — | — | **283.56 GiB** | **150.08 GiB** | **133.48 GiB (47.1%)** |
The direct-mmap workers each reported 62.45 GiB `Shared_Clean` but only
31.20 GiB `Pss_File`, which is the proportional charge expected when the same
resident checkpoint pages are mapped by two workers. `Private_Dirty` also fell
from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker
1, a reduction of about 97–98 GiB per worker. RSS understates this benefit
because it counts a shared physical page in every process that maps it; summed
PSS is the appropriate node-memory comparison.
The highest-value missing coverage is broader end-to-end numerical and
lifecycle comparison against ordinary layerwise offload for DP+SP,
HSDP+SP+no-AllGather, and TP greater than one across additional models and
target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not
change the Phase A direct-mmap TP1 support boundary.
## Recommendations
@@ -150,7 +271,9 @@ on the target CUDA/NCCL or CANN/HCCL hardware.
scaling path.
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
not the goal.
- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
the expectation of higher host-memory use and lower validation confidence.
- Use **no-AllGather** when independent replica execution is required. TP1
direct-mmap deployments can share checkpoint pages per node; other layouts
retain the ordinary loader's private host memory behavior and are outside
the Phase A shared-mmap support scope.
- Prefer **HSDP alone** for production HSDP deployments until the combined
HSDP + DLO no-AllGather path has broader end-to-end coverage.
+11 -3
View File
@@ -76,6 +76,12 @@ equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter.
ModelOpt checkpoint metadata is detected from its method, producer metadata, or
`quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration.
Unless the input is a per-component dictionary, the resolved configuration is
passed unchanged to every quantization-aware component constructed by the
pipeline. It affects only layers supported by that component's quantization
method; it does not rewrite arbitrary `torch.nn` modules. A per-component
dictionary narrows this global scope through runtime layer-prefix routing.
The factory uses Omni-specific overrides for methods whose diffusion behavior
or platform dispatch is not provided by the vLLM registry. The current override
set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/
@@ -154,9 +160,11 @@ quantizable vLLM layers. Each method configuration is responsible for:
Model code must keep precision-sensitive or unsupported modules unquantized.
The shared `safe_quant_config()` helper is used by model integrations for
normalization and modulation layers where a generic FP8 configuration would be
unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers
are not implicitly quantized just because the main transformer is quantized;
they require explicit component routing and method support.
unsafe. Encoders exposed through the quantization factory receive a global
configuration just like the main transformer, but only their supported vLLM
quantizable layers are affected. Ordinary `torch.nn` encoder layers, VAEs,
schedulers, and tokenizers are not rewritten automatically. A component map can
select a narrower scope, but it cannot add quantization support to a layer.
Weight quantization and quantized KV cache are separate paths. The former is
resolved through `quantization_config`; runtime attention/KV-cache quantization
+1 -1
View File
@@ -101,7 +101,7 @@ outputs = omni.generate(
)
# Access custom trajectory data
output = outputs[0].request_output
output = outputs[0]
print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
print(f"Trajectory latents shape: {output.latents.shape}")
```
+2 -3
View File
@@ -47,7 +47,7 @@ if __name__ == "__main__":
omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
prompt = "a cup of coffee on the table"
outputs = omni.generate(prompt)
images = outputs[0].request_output.images
images = outputs[0].images
images[0].save("coffee.png")
```
@@ -75,8 +75,7 @@ if __name__ == "__main__":
]
omni_outputs = omni.generate(prompts)
for i_prompt, prompt_output in enumerate(omni_outputs):
this_request_output = prompt_output.request_output
this_images = this_request_output.images
this_images = prompt_output.images
for i_image, image in enumerate(this_images):
image.save(f"p{i_prompt}-img{i_image}.jpg")
print("saved to", f"p{i_prompt}-img{i_image}.jpg")
+7 -1
View File
@@ -43,7 +43,12 @@ th {
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — |
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
| `LTX2DistilledPipeline` | LTX-2 / LTX-2.3 full-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + `Lightricks/LTX-2.3` upsampler | ✅︎ | ✅︎ | | | — |
| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — |
| `LTX2DistilledTwoStagePipeline` | LTX-2 / LTX-2.3 merged-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + matching Lightricks upsampler | ✅︎ | ✅︎ | | | — |
| `LTX2Pipeline` | LTX-2.5 Full/SFT one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
| `LTX2TwoStagePipeline` | LTX-2.5 Full/SFT two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` + `Lightricks/LTX-2.5` LoRA | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
| `LTX2DistilledOneStagePipeline` | LTX-2.5 merged-distilled one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
| `LTX2DistilledTwoStagePipeline` | LTX-2.5 merged-distilled two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -98,6 +103,7 @@ th {
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | — |
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | | — |
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
The `Recipe` column links the corresponding deployment recipe. `Published` is
@@ -3,7 +3,9 @@
Distributed layerwise offloading (DLO) extends block streaming to multi-device
deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
the host weights and reconstructs each layer at runtime. Without AllGather,
each rank streams its standard-loader rank-local weights independently.
each rank streams a complete block independently. Compatible TP1 deployments
can share checkpoint-backed host pages among processes on the same node;
otherwise DLO streams the ordinary loader's rank-local tensors.
See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
for the implementation contract and compatibility matrix.
@@ -59,21 +61,48 @@ omni = Omni(
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
| `--dlo-no-use-allgather` | Keep standard-loader rank-local weights | `false` |
| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
## mmap weight loading
## Host-weight loading
The DLO plus AllGather path:
The diffusion loader chooses host storage before DLO is enabled. It first
attempts to build a complete, validated direct-checkpoint mmap plan. If names,
coverage, shape, dtype, topology, or loader-callback compatibility cannot be
proven, it runs the ordinary model loader instead. DLO consumes that result and
does not make a second checkpoint-compatibility decision.
The shared-mmap optimization in this phase is supported only with TP1. TP
greater than one falls back before model mutation to ordinary TP-aware loading.
DLO may still consume those TP-local tensors, but this is a compatibility path:
it does not share checkpoint-backed runtime weights across DP replicas and
provides no shared-mmap host-memory guarantee.
The mmap plan skips only dedicated DiT weight sources. Other component sources,
such as a text encoder loaded through the shared diffusion loader, continue to
use their ordinary component loader. A checkpoint source that mixes DiT and
non-DiT weights falls back completely rather than leaving an unplanned
component uninitialized.
With direct checkpoint mmap, the loader:
1. saves non-persistent buffers such as RoPE frequencies;
2. moves the normally created transformer to the meta device;
3. loads checkpoint tensors as mmap views backed by the shared OS page cache;
4. calls model-specific `post_load_weights()` conversions; and
5. restores the saved non-persistent buffers.
4. applies any loader-owned bounded layout adapters while packing blocks;
5. restores saved non-persistent buffers; and
6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle
hooks.
This avoids one full checkpoint RSS copy per rank and does not require
model-specific loading code.
For AllGather with a group larger than one, each process copies only its
persistent shard and then releases the source mapping. For no-AllGather, each
process keeps the mapping open and packs complete blocks through two bounded
pinned staging slots. Processes mapping the same files on one node share the
immutable pages; no-AllGather still performs a complete-block H2D copy in each
process.
When the effective DLO group size is one, `dlo_use_allgather=True` does not
perform a collective and uses the same rank-local transfer behavior.
## Declarative topology
@@ -103,16 +132,22 @@ must enter each collective.
## Limitations
- Online FP8 quantization is rejected with the DLO plus AllGather mmap path.
Use `--dlo-no-use-allgather` or disable online quantization.
- Tensor parallel size greater than one is rejected in the mmap path because
it bypasses TP-aware loader callbacks. The no-AllGather path retains the
standard loader and remains experimental with TP.
- Direct checkpoint mmap currently requires TP1. TP greater than one is
outside the Phase A shared-mmap support scope and falls back before model
mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
but it provides no shared-mmap host-memory benefit or guarantee.
- HSDP plus AllGather is rejected to avoid double sharding. HSDP without
AllGather has limited end-to-end validation.
- Online quantization uses the ordinary loader with no-AllGather. It remains
incompatible with DLO AllGather.
- Resident leading layers require `--dlo-no-use-allgather` and a model
`OffloadPlan` that declares eligible `resident_dit_paths`.
- DP concurrency requires an explicit, identical inference-step count.
Sharing transformed TP or quantized runtime layouts through a normalized mmap
cache is a follow-up design in
[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
the direct-checkpoint path.
See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
for an end-to-end example.
+2 -1
View File
@@ -161,6 +161,7 @@ The following tables show which models support each feature:
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2.5** | ❌ | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -174,7 +175,7 @@ The following tables show which models support each feature:
**Frame Interpolation Support**
- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, Helios, HunyuanVideo-1.5, DreamID-Omni
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni
### AudioGen
@@ -43,7 +43,7 @@ if __name__ == "__main__":
),
],
)
outputs[0].request_output.images[0].save("output.png")
outputs[0].images[0].save("output.png")
```
## Image-to-Image (Image Editing)
@@ -72,7 +72,7 @@ if __name__ == "__main__":
),
],
)
outputs[0].request_output.images[0].save("output.png")
outputs[0].images[0].save("output.png")
```
## Generation Parameters
@@ -55,7 +55,7 @@ if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
prompt = "a cup of coffee on the table"
outputs = omni.generate(prompt)
images = outputs[0].request_output.images
images = outputs[0].images
images[0].save("coffee.png")
```
@@ -179,7 +179,7 @@ if __name__ == "__main__":
]
outputs = omni.generate(prompts)
for i, output in enumerate(outputs):
output.request_output.images[0].save(f"{i}.jpg")
output.images[0].save(f"{i}.jpg")
```
!!! info
@@ -207,7 +207,7 @@ if __name__ == "__main__":
}
])
for i, output in enumerate(outputs):
output.request_output.images[0].save(f"{i}.jpg")
output.images[0].save(f"{i}.jpg")
```
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
+34 -1
View File
@@ -95,7 +95,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
| Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | |
| MiniMax H3 | `MiniMaxAI/MiniMax-H3` | Yes | Not validated | DiT linears except patch, timestep, and final projections; not compatible with layerwise offload | None | |
| MiniMax-H3 | `MiniMaxAI/MiniMax-H3` (`FL2VA` / `Ref2VA`) | Yes | Not validated | `quantization="fp8"` quantizes eligible DiT and text-encoder linears; mixed-precision input/output heads stay FP32 | None | ✅︎ |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
@@ -138,6 +138,39 @@ outputs = omni.generate(
)
```
### Global and per-component scope
A plain, global quantization configuration is passed unchanged to every
quantization-aware component constructed by a pipeline. This includes an
eligible encoder when that encoder is implemented with vLLM quantizable
layers; it does not rewrite arbitrary `torch.nn` modules. A structured
component map is the only way to narrow that scope. Pipeline integrations that
do not yet expose an encoder through the quantization factory remain DiT-only.
For a pipeline that exposes both a transformer and a quantization-aware text
encoder, the scope is:
| Configuration | Transformer | Text encoder | Components without supported quantizable layers |
|---------------|-------------|--------------|-------------------------------------------------|
| `quantization="fp8"` | FP8 | FP8 | checkpoint precision |
| `{"transformer": {"method": "fp8"}}` | FP8 | checkpoint precision | checkpoint precision |
| `{"text_encoder": {"method": "fp8"}}` | checkpoint precision | FP8 | checkpoint precision |
Use `quantization_config` for component-selective Python configuration. For
example, quantize only the text decoder:
```python
omni = Omni(
model="<your-model>",
quantization_config={"text_encoder": {"method": "fp8"}},
)
```
Component keys are runtime prefixes exposed by the pipeline integration; common
keys include `transformer` and `text_encoder`. Entries may be combined, and
`ignored_layers` can keep named eligible layers in checkpoint precision. Check
the model recipe for supported components and their runtime prefixes.
CLI:
```bash
@@ -96,7 +96,7 @@ if __name__ == "__main__":
),
)
from diffusers.utils import export_to_video
frames = outputs[0].request_output.images
frames = outputs[0].images
export_to_video(frames, "quick_test_i2v.mp4", fps=16)
```
@@ -68,6 +68,7 @@ from vllm_omni.model_extras import (
from vllm_omni.model_extras import (
get_extra_body_params,
get_model_class_name,
should_preserve_reference_image_size,
)
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
@@ -454,7 +455,13 @@ def main():
media_inputs: dict[str, Any] = {}
if image is not None:
media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
preserve_image_size = should_preserve_reference_image_size(
model_class_name,
model=args.model,
)
media_inputs["image"] = (
image if preserve_image_size else image.resize((width, height), PIL.Image.Resampling.LANCZOS)
)
if last_image is not None:
media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS)
if mask_image is not None:
@@ -335,7 +335,7 @@ def main(args):
print(f"use_tts (tts_bos): {use_tts}")
for stage_outputs in omni_generator:
output = stage_outputs.request_output
output = stage_outputs
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -1,21 +1,32 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior.
"""Validate MiniMax-H3 request-mode and DLO lifecycle behavior.
The DLO modes accept any positive data-parallel size. With AllGather enabled,
the script submits one complete DP wave and verifies recovery after an invalid
wave. It also checks that every mode shuts down all worker processes cleanly.
Examples:
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
VLLM_OMNI_DLO_DP_WAVE_TIMEOUT=600 \
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode dlo --dp-size 4
# Rank-local mmap storage, with no DLO collective or DP wave scheduling.
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
CUDA_VISIBLE_DEVICES=0,1 \
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode dlo-no-allgather --dp-size 2
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
CUDA_VISIBLE_DEVICES=0,1 \
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode request
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode request --tp-size 2
"""
from __future__ import annotations
@@ -40,14 +51,33 @@ DEFAULT_PROMPTS = (
)
def positive_int(value: str) -> int:
parsed = int(value)
if parsed < 1:
raise argparse.ArgumentTypeError("must be a positive integer")
return parsed
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
parser.add_argument(
"--mode",
choices=("dlo-dp2", "request"),
choices=("dlo", "dlo-no-allgather", "request"),
required=True,
help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
help="DLO (with or without AllGather) or ordinary non-DLO request mode",
)
parser.add_argument(
"--dp-size",
type=positive_int,
default=2,
help="DLO data-parallel size (for example 1, 2, 4, or 8; default: 2)",
)
parser.add_argument(
"--tp-size",
type=positive_int,
default=2,
help="Tensor-parallel size used only by request mode (default: 2)",
)
parser.add_argument("--steps", type=int, default=2)
parser.add_argument("--duration", type=float, default=5.0)
@@ -60,10 +90,11 @@ def parse_args() -> argparse.Namespace:
def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
is_dlo = args.mode in {"dlo", "dlo-no-allgather"}
common: dict[str, Any] = {
"model": args.model,
"trust_remote_code": True,
"num_gpus": 2,
"num_gpus": args.dp_size if is_dlo else args.tp_size,
"usp": 1,
"ring": 1,
"vae_parallel_mode": "tile",
@@ -74,22 +105,22 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
"stage_init_timeout": args.init_timeout,
"init_timeout": args.init_timeout,
}
if args.mode == "dlo-dp2":
if is_dlo:
common.update(
tensor_parallel_size=1,
data_parallel_size=2,
data_parallel_size=args.dp_size,
text_encoder_tp_size=1,
vae_patch_parallel_size=1,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
dlo_use_allgather=args.mode == "dlo",
dlo_resident_layers=0,
)
else:
common.update(
tensor_parallel_size=2,
tensor_parallel_size=args.tp_size,
data_parallel_size=1,
text_encoder_tp_size=2,
vae_patch_parallel_size=2,
text_encoder_tp_size=args.tp_size,
vae_patch_parallel_size=args.tp_size,
enable_distributed_layerwise_offload=False,
)
return common
@@ -172,7 +203,8 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
"engine_kwargs": engine_kwargs(args),
}
try:
if args.mode == "dlo-dp2":
uses_dp_wave = args.mode == "dlo" and args.dp_size > 1
if uses_dp_wave:
started = time.perf_counter()
asymmetric = await asyncio.gather(
generate_one(
@@ -182,12 +214,15 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
prompt="",
seed=1001,
),
generate_one(
engine,
args,
request_id="invalid-peer",
prompt=DEFAULT_PROMPTS[0],
seed=1002,
*(
generate_one(
engine,
args,
request_id=f"invalid-peer-{index}",
prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
seed=1001 + index,
)
for index in range(1, args.dp_size)
),
return_exceptions=True,
)
@@ -195,18 +230,18 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
summary["asymmetric_errors"] = [
f"{type(result).__name__}: {result}" for result in asymmetric if isinstance(result, BaseException)
]
if len(summary["asymmetric_errors"]) != 2:
raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch")
if len(summary["asymmetric_errors"]) != args.dp_size:
raise RuntimeError("Every request in the asymmetric DP wave must fail before dispatch")
started = time.perf_counter()
request_count = 2 if args.mode == "dlo-dp2" else 1
request_count = args.dp_size if uses_dp_wave else 1
outputs = await asyncio.gather(
*(
generate_one(
engine,
args,
request_id=f"recovery-{index}",
prompt=DEFAULT_PROMPTS[index],
prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
seed=2000 + index,
)
for index in range(request_count)
@@ -74,7 +74,7 @@ if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
prompt = "a cup of coffee on the table"
outputs = omni.generate(prompt)
images = outputs[0].request_output.images
images = outputs[0].images
images[0].save("coffee.png")
```
@@ -265,7 +265,7 @@ if __name__ == "__main__":
]
outputs = omni.generate(prompts)
for i, output in enumerate(outputs):
output.request_output.images[0].save(f"{i}.jpg")
output.images[0].save(f"{i}.jpg")
```
!!! info
@@ -301,7 +301,7 @@ if __name__ == "__main__":
}
])
for i, output in enumerate(outputs):
output.request_output.images[0].save(f"{i}.jpg")
output.images[0].save(f"{i}.jpg")
```
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
@@ -16,6 +16,7 @@ list of supported architectures across all modalities, see
|---|---|---|---|---|---|---|
| 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 |
| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz |
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
@@ -395,6 +396,58 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
---
## Gepard-1.0
Single-stage native AR TTS at 22.05 kHz. Pipeline: `Qwen3.5 backbone → 32 FSQ codebook heads (one frame/step) → NeMo NanoCodec`. Backbone runs under vLLM paged attention; the 32-head sampling + learned embedding feedback ride the native-AR runner hooks. `enforce_eager` (CUDA graph is a perf follow-up).
### Prerequisites
The NanoCodec decoder needs NeMo (NVIDIA Open Model License), installed separately:
```bash
uv pip install "nemo_toolkit[tts]==2.7.3" \
--overrides <(printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0")
```
The overrides are what make that command work. `nemo_toolkit[tts]` declares
`transformers~=4.57.0`, so a plain install downgrades the version
`requirements/common.txt` pins and takes `huggingface-hub` back to 0.x with it — after
which `import vllm` fails. Both packages are named here rather than left to resolve on
their own: transformers 5.x happens to require `huggingface-hub` 1.x, but relying on
that would make the second repair an accident of the first. `--constraint` cannot be
used instead, since it narrows a version range rather than overriding a dependency a
package declares, and the two sets never intersect. NeMo 2.7.3 decodes correctly
against transformers 5.x, so nothing is lost by holding it there.
(The same clash is why this is documented here rather than declared in
`pyproject.toml`: an extra resolves together with the base dependencies, so
`vllm-omni[gepard]` would be unsatisfiable rather than opt-in. SoulX-Singer English
SVS documents `nemo_toolkit[asr]` the same way.)
On a host whose CUDA toolkit cannot build kernels for the GPU — no `nvcc`/`ninja`, or a
consumer Blackwell (`sm_120`) card — also set `VLLM_USE_FLASHINFER_SAMPLER=0`. vLLM's
sampler JIT-compiles through FlashInfer on first use, and where that build fails the engine
dies inside `profile_run` long after every import has succeeded, which reads as a model
failure rather than an environment one. The PyTorch sampler it falls back to costs nothing
here.
### Quick start (zero-shot, default voice)
```bash
python examples/offline_inference/text_to_speech/gepard/end2end.py \
--text "Hello, this is a Gepard demo."
```
### Voice cloning
Not yet — PR1 is zero-shot only (the learned `null_prefix` default voice). Reference-audio cloning (`ref_compressor` speaker prefix) is a follow-up PR.
### Notes
- Output: 22.05 kHz mono WAV.
- First run downloads two checkpoints, not one: the model itself, and the NanoCodec decoder it names in `codec_id` (`nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps` by default). Both are fetched automatically; expect the first startup to be correspondingly slower.
- Offline only for now; the `/v1/audio/speech` adapter is a follow-up PR.
- Deploy config: `vllm_omni/deploy/gepard.yaml` (the example's default; copy it and pass `--deploy-config` to change it).
- Generation length and reproducibility are stage settings, not script flags. One output token is one audio frame, so `max_tokens` in the YAML is the frame budget; `seed` makes the in-model 32-head sampling reproducible. The script deliberately passes no `SamplingParams`: one supplied by a caller replaces the stage defaults wholesale rather than merging, which would drop the pipeline's stop token and run every request to `max_tokens`.
- Text length: short texts are repeated internally to match the training layout (the checkpoint's `text_repetition` block); the upper bound is the stage's `max_model_len`, enforced by the engine. Empty text is rejected rather than voiced.
- `VLLM_GEPARD_GREEDY=1` swaps the 32-head Gumbel-max sampling for argmax, for reproducible comparisons that do not depend on a seed.
## VoxCPM2
Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`.
@@ -0,0 +1,139 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline inference example for Gepard-1.0 TTS via vLLM-Omni.
Single-stage native-AR pipeline: a Qwen3.5 backbone samples one 32-code FSQ
frame per step; the NeMo NanoCodec decodes the frames to a 22.05 kHz mono
waveform. Zero-shot: the default learned voice, no reference audio.
Generation length and seed are stage settings rather than flags here. One
output token is one audio frame, so ``max_tokens`` in the deploy YAML is the
frame budget; pass a copy of the YAML via ``--deploy-config`` to change it.
Usage:
python end2end.py --text "Hello, this is Gepard speaking."
"""
from __future__ import annotations
import os
from pathlib import Path
import soundfile as sf
import torch
import vllm_omni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
from vllm_omni import Omni # noqa: E402
MODEL = "nineninesix/gepard-1.0"
SAMPLE_RATE = 22050
# Absolute: the loader hands the path straight to OmegaConf, which resolves a
# bare name against the current working directory, not the package.
DEFAULT_DEPLOY_CONFIG = str(Path(vllm_omni.__file__).resolve().parent / "deploy" / "gepard.yaml")
def build_request(text: str) -> dict:
"""Build an Omni request payload for Gepard (zero-shot, no ref audio).
The model only consumes the [speaker slots, SOT, text, EOT, SOS] layout, so
a bare text prompt would have no speaker slots and no SOS.
"""
from transformers import AutoTokenizer
from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
# From the checkpoint's own sidecar, which is what the worker builds its
# config from too. A layout that disagrees with the trained one does not
# fail, it just degrades the speech.
cfg = GepardConfig.from_checkpoint(MODEL)
tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
prompt_token_ids = build_gepard_prompt_ids(
tokenizer(text, add_special_tokens=False)["input_ids"],
config=cfg,
)
return {
"prompt_token_ids": prompt_token_ids,
"additional_information": {"text": [text]},
}
def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = SAMPLE_RATE) -> None:
sf.write(path, waveform.float().numpy(), sample_rate)
seconds = waveform.numel() / sample_rate
print(
f" Saved {path} ({tuple(waveform.shape)}, {sample_rate} Hz, "
f"{seconds:.2f}s, peak {float(waveform.abs().max()):.3f})"
)
def main(args) -> None:
omni = Omni(
model=MODEL,
deploy_config=args.deploy_config,
stage_init_timeout=args.stage_init_timeout,
init_timeout=args.init_timeout,
)
# No explicit SamplingParams: one supplied here would replace the stage
# defaults rather than merge over them, dropping the pipeline's stop token.
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
print(f"Synthesizing: {args.text!r}")
inputs = build_request(text=args.text)
for i, stage_outputs in enumerate(omni.generate(inputs)):
for j, out in enumerate(stage_outputs.outputs):
mm = out.multimodal_output
if mm is None:
print(f" [req {i}] No audio output.")
continue
# The consolidation path renames "model_outputs" to "audio".
audio = mm.get("audio")
if audio is None:
audio = mm.get("model_outputs")
if isinstance(audio, list):
audio = audio[0] if len(audio) else None
if audio is None:
print(f" [req {i}] No waveform in multimodal_output (keys: {list(mm)}).")
continue
sr_tensor = mm.get("sr")
sr = int(sr_tensor.item()) if hasattr(sr_tensor, "item") else SAMPLE_RATE
save_audio(audio.cpu(), str(output_dir / f"output_{i}_{j}.wav"), sr)
print("Done.")
def parse_args():
parser = TrackingArgumentParser(description="Gepard-1.0 offline TTS inference (zero-shot)")
parser.add_argument("--text", default="Hello, this is Gepard speaking.", help="Text to synthesize.")
parser.add_argument(
"--output-dir",
default=os.path.join(
os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")),
"gepard_output",
),
help="Directory for WAV outputs (default: ~/.cache/gepard_output).",
)
parser.add_argument(
"--deploy-config",
default=DEFAULT_DEPLOY_CONFIG,
help=(
"Path to the deploy YAML (default: the packaged vllm_omni/deploy/gepard.yaml). This must "
"stay set: without the YAML's `pipeline: gepard` pin the architectures fallback routes "
"Qwen3_5ForCausalLM to the diffusion registry. Copy it to change max_tokens or seed."
),
)
# The upstream defaults are too tight for a cold start.
parser.add_argument("--stage-init-timeout", type=int, default=900)
parser.add_argument("--init-timeout", type=int, default=1800)
return parser.parse_args()
if __name__ == "__main__":
main(parse_args())
@@ -570,6 +570,8 @@ def main():
num_frames=args.num_frames,
extra_args=extra_args,
)
if args.frame_rate is not None:
sampling_kwargs["frame_rate"] = args.frame_rate
if args.guidance_scale_high is not None:
sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
+175
View File
@@ -0,0 +1,175 @@
# LTX-2.5
> Text-to-video and first-frame image-to-video generation with synchronized audio
## Overview
vLLM-Omni supports the gated
[`Lightricks/LTX-2.5-Diffusers`](https://huggingface.co/Lightricks/LTX-2.5-Diffusers)
checkpoint for T2V and first-frame I2V. Every output includes synchronized
48 kHz stereo audio.
The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5)
repository is not directly loadable with `--model`. It supplies the official
upsampler and LoRA sidecars used by the Full/SFT two-stage pipeline, so accept
both model licenses and authenticate before first use.
## Choose a pipeline
| Pipeline | Mode | Default size | Schedule |
|---|---|---:|---:|
| `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps |
| `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 steps |
| `LTX2DistilledOneStagePipeline` | Distilled one-stage | 960x544 | 8 steps |
| `LTX2DistilledTwoStagePipeline` | Distilled two-stage | 1920x1088 | 8 + 3 steps |
Both two-stage pipelines generate at half resolution, apply the official x2
latent upsampler, and run a three-step refinement stage. Select the class with
`--model-class-name`; no `--task-type` flag is required. Supplying one initial
image selects I2V, while omitting it selects T2V.
## Prerequisites
```bash
hf auth login
export MODEL=Lightricks/LTX-2.5-Diffusers
```
Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and
`ffprobe` are on `PATH`. I2V requires PyAV backed by an FFmpeg build with the
`libx264` encoder.
## Hardware
| GPU | Status | Recommended scope |
|---|---|---|
| NVIDIA B300 | Verified | All four canonical pipelines |
| NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
| NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
| NVIDIA H100 80 GB | FP8 recipe | Distilled one-stage at 960x544 |
The 1920x1088 two-stage examples require about 114 GB of peak GPU memory.
80 GB GPUs do not have enough safety margin for the canonical two-stage
configuration; use a larger GPU or reduce the output configuration.
### H100 80 GB FP8
Use the distilled one-stage pipeline with FP8 and cuDNN attention:
```bash
vllm serve Lightricks/LTX-2.5-Diffusers \
--omni \
--model-class-name LTX2DistilledOneStagePipeline \
--quantization fp8 \
--diffusion-attention-backend CUDNN_ATTN \
--host 0.0.0.0 \
--port 8000
```
Use the T2V or I2V request below after the server is ready.
## Offline inference
Choose values from the pipeline table. For example, the distilled two-stage
path uses:
```bash
export PIPELINE=LTX2DistilledTwoStagePipeline
export WIDTH=1920
export HEIGHT=1088
export STEPS=8
```
### Text-to-video
```bash
python examples/offline_inference/text_to_video/text_to_video.py \
--model "${MODEL}" \
--model-class-name "${PIPELINE}" \
--prompt "A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot." \
--width "${WIDTH}" \
--height "${HEIGHT}" \
--num-frames 121 \
--num-inference-steps "${STEPS}" \
--frame-rate 24 \
--fps 24 \
--seed 42 \
--output ltx25-t2v.mp4
```
### First-frame image-to-video
```bash
python examples/offline_inference/image_to_video/image_to_video.py \
--model "${MODEL}" \
--model-class-name "${PIPELINE}" \
--image /absolute/path/to/first-frame.png \
--prompt "The red fox walks forward while the camera tracks alongside." \
--width "${WIDTH}" \
--height "${HEIGHT}" \
--num-frames 121 \
--num-inference-steps "${STEPS}" \
--frame-rate 24 \
--fps 24 \
--seed 42 \
--output ltx25-i2v.mp4
```
LTX-2.5 uses the official CRF-18 first-frame conditioning path by default.
## Online serving
Start one server for the selected pipeline:
```bash
vllm serve "${MODEL}" \
--omni \
--model-class-name "${PIPELINE}" \
--host 0.0.0.0 \
--port 8000 \
--stage-init-timeout 900
```
T2V request:
```bash
curl -sS --fail-with-body \
-X POST http://127.0.0.1:8000/v1/videos/sync \
-F 'prompt=A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot.' \
-o ltx25-online-t2v.mp4
```
For I2V, add exactly one first frame:
```bash
export FIRST_FRAME=/absolute/path/to/first-frame.png
curl -sS --fail-with-body \
-X POST http://127.0.0.1:8000/v1/videos/sync \
-F "input_reference=@${FIRST_FRAME};type=image/png" \
-F 'prompt=The red fox walks forward while the camera tracks alongside.' \
-o ltx25-online-i2v.mp4
```
Restart the server after changing `PIPELINE` because each class selects a
different weight layout and execution topology.
## Constraints
- Online `width`, `height`, `num_frames`, `fps`, and `num_inference_steps`
are optional and use the selected pipeline defaults when omitted. `seed` is
optional and random when omitted.
- `num_frames` must be `8k+1` when overridden.
- One-stage width and height must be divisible by 32; two-stage final
dimensions must be divisible by 64.
- Full/SFT pipelines accept negative prompts. Distilled pipelines are
positive-only and reject negative prompts.
- I2V accepts exactly one initial image per prompt.
- The supported `--model` value is `Lightricks/LTX-2.5-Diffusers`.
## References
- [LTX-2.5-Diffusers model card](https://huggingface.co/Lightricks/LTX-2.5-Diffusers)
- [Official LTX-2 repository](https://github.com/Lightricks/LTX-2)
- [Generic T2V example](../../examples/offline_inference/text_to_video/text_to_video.py)
- [Generic I2V example](../../examples/offline_inference/image_to_video/image_to_video.py)
+14 -9
View File
@@ -8,10 +8,12 @@
|---|---|---|
| `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` |
| `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` |
| `LTX2DistilledPipeline` | LTX-2 full-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
| `LTX2DistilledOneStagePipeline` | LTX-2 merged-distilled one-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
| `LTX2DistilledTwoStagePipeline` | LTX-2 merged-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
| `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` |
| `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers`<br>`Lightricks/LTX-2.3` |
| `LTX2DistilledPipeline` | LTX-2.3 full-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |
| `LTX2DistilledOneStagePipeline` | LTX-2.3 merged-distilled one-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers` |
| `LTX2DistilledTwoStagePipeline` | LTX-2.3 merged-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |
Repositories in the table are download units. A full pipeline repository
contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler,
@@ -26,9 +28,12 @@ image selects I2V. Both one-stage repositories declare this class, so
the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable.
`LTX2TwoStagePipeline` samples the regular model at half resolution, upsamples,
then refines with the distilled LoRA. `LTX2DistilledPipeline` uses a fully
merged distilled Transformer in both stages. Both entries support T2V and I2V;
select their class explicitly.
then refines with the distilled LoRA. `LTX2DistilledOneStagePipeline` uses a
merged distilled Transformer without upsampling, while
`LTX2DistilledTwoStagePipeline` uses it in both stages. All entries support
T2V and I2V; select their class explicitly. The deprecated
`LTX2DistilledPipeline` name remains an alias for
`LTX2DistilledTwoStagePipeline`.
## API Migration
@@ -52,8 +57,8 @@ The consolidation also removes these registry names without aliases:
| `LTX23Pipeline` | `LTX2Pipeline`; checkpoint metadata selects LTX-2.3 |
| `LTX2ImageToVideoPipeline` | `LTX2Pipeline` with `image=` |
| `LTX23ImageToVideoPipeline` | `LTX2Pipeline` with `image=`; checkpoint metadata selects LTX-2.3 |
| `LTX2TwoStagesPipeline` | `LTX2DistilledPipeline` |
| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledPipeline` with `image=` |
| `LTX2TwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` |
| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` with `image=` |
Passing any second positional argument now raises `TypeError`. These changes
affect direct Python callers and explicit `--model-class-name` overrides;
@@ -112,10 +117,10 @@ Start a two-stage checkpoint with its explicit class:
```bash
vllm serve rootonchair/LTX-2-19b-distilled --omni \
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
--model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
# LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent
vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
--model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
# Ordinary LTX-2.3 two-stage
vllm serve diffusers/LTX-2.3-Diffusers --omni \
--model-class-name LTX2TwoStagePipeline \
+195
View File
@@ -0,0 +1,195 @@
# MiniMax-H3 on a single NPU 950PR
This recipe runs MiniMax-H3 with online INT8 quantization on one NPU 950PR
NPU (128 GB HiBL 1.0 HBM). It covers the single-card T2VA configuration at
1024x576. For the eight-card Atlas 800I A3 BF16 route at 768P, see
[MiniMax-H3-NPU.md](MiniMax-H3-NPU.md).
## Capacity requirements
| Resource | Requirement |
| --- | ---: |
| NPU | 1x NPU 950PR |
| NPU HBM | 128 GiB (131,072 MiB reported by `npu-smi`) |
| Observed HBM high-water mark | 118,442 MiB (90.4% of capacity) |
| Checkpoint storage | 135 GiB per partition, local disk strongly preferred |
| Container shared memory | 8 GiB, or a writable local filesystem for spill |
The observed high-water mark leaves about 12.3 GiB of headroom, so this card is
sized for exactly one resident task partition at 1024x576. Do not expect a
larger output shape or concurrency greater than one to fit.
Because a single 950PR holds the whole partition in HBM, this route does not
need layerwise offload, and the 200 GiB system-RAM floor carried by the 72 GiB
GPU recipes does not apply. The validated run completed inside a container
limited to 32 GiB of RAM. Passing `--enable-layerwise-offload` on this card is
actively harmful: it stages weights in non-reclaimable host memory and the
container's OOM killer terminates the server with exit code -9.
## Environment
- Host architecture: x86_64
- driver: 25.7.rc1.6 (hal 7.35.23)
- firmware: 9.0.0.105.229
- CANN toolkit: 9.1.0 (`/usr/local//cann-9.1.0`)
- npu-smi: 25.7.rc1.6
- Python: 3.12.13
- PyTorch: 2.10.0+cpu
- torch_npu: 2.10.0.post2
- vLLM: 0.26.0
- vLLM-Omni: 0.26.1.dev103+g584d78c67.npu (commit `584d78c6`)
Install vLLM-Omni from a checkout with MiniMax-H3 support:
```bash
uv venv
source .venv/bin/activate
uv pip install -e .
```
## Start a server
```bash
export MODEL=/path/to/MiniMax-H3/FL2VA
export PORT=8000
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400
vllm serve "${MODEL}" \
--omni \
--host 0.0.0.0 \
--port "${PORT}" \
--trust-remote-code \
--num-gpus 1 \
--tensor-parallel-size 1 \
--usp 1 \
--ring 1 \
--text-encoder-tp-size 1 \
--vae-patch-parallel-size 1 \
--vae-parallel-mode tile \
--vae-use-tiling \
--init-timeout 14400 \
--stage-init-timeout 14400 \
--quantization int8
```
Do not pass `--diffusion-attention-backend CUDNN_ATTN`. cuDNN attention has no
NPU implementation and fails at the first denoise step. Leave the backend
unset: it resolves to `TORCH_SDPA`, or to `FLASH_ATTN` when MindIE-SD is
installed (see [MiniMax-H3-NPU.md § Environment](MiniMax-H3-NPU.md#environment)).
H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
## Validated evidence
Measured on one NPU 950PR with the server command above plus
`--enable-diffusion-pipeline-profiler`, generating a 5 s 1024x576 T2VA clip from
the request in [§ T2VA request example](#t2va-request-example). MiniMax-H3
requested 60 denoise steps and executed 59 denoise updates, so per-step latency
is `denoise / 59`. `ffprobe` confirms every returned MP4 is 1024x576, 24 fps,
124 video frames, 5.175 s, with a 32 kHz audio track of 165,600 samples.
Five requests were issued back to back with no idle time between them. The card
is measurably faster on the first request than in steady state, so both regimes
are reported.
| Stage | First request, cold card | Thermal steady state |
| --- | ---: | ---: |
| End-to-end request | 462.13 s | 506.38 s |
| Text encode | 0.55 s | 0.05 s |
| Denoise (59 updates) | 429.29 s | 471.74 s |
| Per denoise update | 7,276 ms | 7,996 ms |
| VAE decode | 8.43 s | 8.14 s |
| Worker-to-server handoff | 24 s | 27 s |
| MP4 muxing | 0.57 s | 0.50 s |
| Server-reported `denoise_step_latency_ms` | 7,702 ms | 8,440 ms |
| Peak HBM (`npu-smi`) | 118,442 MiB | 118,442 MiB |
| Average NPU power | 566.9 W | 556.9 W |
| Peak NPU temperature | 104 C | 104 C |
Sample counts differ per row. The cold column is a single request. In the steady
column, text encode, denoise and VAE decode are means over the four subsequent
requests — denoise spanned 470.37 to 473.21 s, a 0.60% spread — while
end-to-end, handoff and muxing are means over the two of those four that
survived the handoff timeout described above.
**Report the steady-state column.** The cold card completes denoise 9.9% faster
while drawing only 1.8% more average power, and both regimes reach the same
104 C ceiling, so the first request is buying a short window of higher clocks
before the die saturates. Any benchmark that issues a single request against an
idle card will overstate throughput by about 10%.
**Do not read `denoise_step_latency_ms` as a per-step time.** The server divides
the entire stage wall time by the 60 *requested* steps, so it absorbs text
encoding, VAE decode and the 837 MiB handoff. All three successful requests
satisfy `denoise_step_latency_ms == e2e_stage_wall_time_ms / 60` exactly. The
per-update figure in the table is `diffuse / 59` from the pipeline profiler.
Denoise accounts for 93% of the steady-state request and the handoff accounts
for 5.3%, so keeping the transfer in shared memory should bring end-to-end down
to roughly 480 s.
Peak HBM is sampled externally with `npu-smi` at a 3 s interval. It is a
reserved high-water mark rather than live allocation: it reads 111,272 MiB after
weight loading and before the first request, and does not fall back after
requests complete. Re-measure for longer outputs, a different output shape, or
concurrency greater than one.
### Startup
Weight loading took **31 min 50 s** on the validated host (first log line
08:04:59, `Application startup complete` 08:36:49), with the checkpoint on a
shared GlusterFS network mount. This is why the server command sets
`--init-timeout 14400` and `--stage-init-timeout 14400`; the stock 600 s and
1800 s timeouts both abort mid-load with `TimeoutError` and exit code 143.
Reduce both timeouts when the partition is staged on local disk.
## T2VA request example
```bash
export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
curl -sS --max-time 1800 -X POST "${API_URL}" \
-F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
-F 'width=1024' \
-F 'height=576' \
-F 'aspect_ratio=16:9' \
-F 'fps=24' \
-F 'num_inference_steps=60' \
-F 'flow_shift=12' \
-F 'seed=1101' \
-F 'extra_params={"task":"t2va","duration":5,"audio_flow_shift":3.0}' \
-o t2va.mp4
```
## Known limitations
- INT8 online quantization is validated for T2VA on this card. Use BF16 for
FL2VA and Ref2VA, or re-measure before relying on it.
- Single-card serving loads one task partition at a time. For Ref2VA, stop the
server and restart it against the `Ref2VA` directory.
- Sustained load drives the die to 104 C and costs about 10% of denoise
throughput relative to a cold card. Warm the server with at least one full
request before measuring, and treat steady state as the reportable number.
- Requests can fail during result handoff on hosts where the IPC spill path is
slow. See [§ Result handoff and shared
memory](#result-handoff-and-shared-memory).
- Loading the checkpoint from a network mount takes over 30 minutes and forces
very large init timeouts. Stage the partition on local disk when possible.
- The configuration measured here is 1024x576 at 60 steps, which differs from
the 1344x768 at 50 steps used by the GPU recipes in this directory. Denoise
cost scales superlinearly with token count, so the numbers are not directly
comparable across recipes.
- The image ships a `triton` package whose backend is not built
(`No module named 'triton._C.libtriton.'`). vLLM logs this as an error
at startup and disables Triton; the diffusion path does not need it and the
run completes normally.
## Additional resources
- [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide
- [MiniMax-H3-NPU.md](MiniMax-H3-NPU.md) — eight-card Atlas 800I A3 BF16 guide
- [Int8 quantization](../../docs/user_guide/quantization/int8.md)
- [Supported models](../../docs/models/supported_models.md)
- [Video API](../../docs/serving/videos_api.md)
+48 -11
View File
@@ -357,12 +357,19 @@ No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while
### Online FP8 quantization
MiniMax H3 supports load-time FP8 quantization of the DiT. The checkpoint
remains BF16 on disk; vLLM-Omni quantizes eligible weights while loading and
uses dynamic activation scaling during inference. By default, attention and
MLP linears in the token refiner and main DiT blocks, the condition
projection, and all AdaLN projections use FP8. Patch, timestep, and final
projections remain FP32; the text encoder and VAEs are unchanged.
MiniMax H3 supports online FP8 quantization of both the DiT and the Qwen3-VL
text decoder. The checkpoint remains BF16 on disk; vLLM-Omni creates FP8
weights at runtime and uses dynamic activation scaling during inference. By
default, `--quantization fp8` quantizes eligible attention and MLP linears in
the text decoder, token refiner, and main DiT blocks, as well as the DiT
condition and AdaLN projections. The Qwen vision tower, embeddings, norms,
RoPE, both VAEs, and the model's FP32 patch, timestep, and output projections
keep checkpoint precision.
To select a component, use `--diffusion-quantization-config` with
`{"transformer":{"method":"fp8"}}` for DiT-only FP8 or
`{"text_encoder":{"method":"fp8"}}` for text-decoder-only FP8. The two
entries can be combined. The shorthand below enables both components.
Add this option to an existing H3 server command:
@@ -370,6 +377,28 @@ Add this option to an existing H3 server command:
--quantization fp8
```
#### Single 96 GB GPU, no-offload capacity check
Use the FL2VA-only partition for this capacity test. Loading the combined
service would also load the Ref2VA DiT and would test a different memory
budget. A no-offload capacity check should contain none of
`--enable-cpu-offload`, `--enable-layerwise-offload`, or
`--enable-distributed-layerwise-offload`. VAE tiling changes decode placement
but does not offload model weights to the CPU.
The run passes the capacity check when the server initializes, the request
finishes without CUDA OOM or Xid errors, `peak_used_mib` remains below the
card's reported `memory_total_mib`, and `ffprobe` reports H.264 video plus
32 kHz stereo AAC audio. Report the measured headroom rather than assuming
that every nominal 96 GB SKU exposes the same MiB total.
As a capacity proxy only, the same five-second case on one B300 measured a
92,946 MiB whole-device first-request peak and a 92,146 MiB worker peak. Its
encode, diffuse, decode, and client wall times were 8.664 s, 134.504 s,
6.016 s, and 151.583 s. These numbers suggest that a 96 GB card may fit, but
they are not an RTX PRO 6000 validation: kernels, allocator behavior, and
reported device capacity differ across GPUs.
Use `ignored_layers` to keep any otherwise eligible linear in BF16. H3
resolves the `transformer` component before constructing the DiT, so names do
not start with `transformer.`. Entries are exact runtime linear prefixes; a
@@ -386,10 +415,10 @@ For example, keep the first main block's attention projections in BF16 with:
'{"transformer":{"method":"fp8","ignored_layers":["blocks.0.attn.qkv_proj","blocks.0.attn.out_proj"]}}'
```
The structured option replaces `--quantization fp8`. Online FP8 is currently
incompatible with H3 layerwise offload because the offload path produces a
weight stride rejected by the Cutlass FP8 kernel. Use resident FP8 with tensor
parallelism and VAE tiling instead.
The structured option replaces `--quantization fp8`. Online FP8 can be used
with H3 layerwise offload and with distributed layerwise offload's full-weight
per-rank path (`--dlo-no-use-allgather`). The sharded DLO AllGather path is not
supported for runtime-created FP8 weights.
## AMD ROCm (gfx942 / gfx950)
@@ -750,6 +779,13 @@ audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The
resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22%
reduction.
For direct human inspection, see the external
[BF16 versus global-FP8 comparison](https://lishunyang12.github.io/vllm-omni-rankings/scripts/minimax_h3_global_fp8_vs_bf16/),
which includes matched five-second T2VA, I2VA, and Ref2VA videos. The
[comparison sources](https://github.com/lishunyang12/vllm-omni-rankings/tree/main/scripts/minimax_h3_global_fp8_vs_bf16)
also record per-task fidelity metrics and provenance without storing generated
media in this repository.
## TeaCache acceleration
TeaCache reuses DiT block residuals across denoising steps when consecutive
@@ -834,7 +870,8 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
H3 native `tile` mode only.
- A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
pure Ulysses.
- Online FP8 is not compatible with layerwise offload.
- Runtime-created FP8 weights do not support the sharded DLO AllGather path;
use `--dlo-no-use-allgather` when combining online FP8 with DLO.
- TeaCache and Cache-DiT cannot be enabled on the same server.
- Image+audio Ref2VA accepts exactly one image and one audio reference.
- Video Ref2VA accepts one or more video files, but not an additional standalone
+1
View File
@@ -55,6 +55,7 @@ recipes/
| [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) |
| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
| [`LTX/LTX-2.5.md`](./LTX/LTX-2.5.md) | LTX-2.5-Diffusers: Full/SFT one-stage and distilled two-stage T2V/I2V with synchronized audio | NVIDIA B300; cuDNN-qualified |
| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
| [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB |
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
+5 -5
View File
@@ -1,15 +1,15 @@
# Common dependencies for all platforms
# transformers: kept exactly aligned with the upstream vLLM constraint.
# vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0".
# The previous extra "< 5.9.0" exclusion was removed when re-aligning with
# upstream: transformers 5.9 changed _adjust_missing_and_unexpected_keys from
# transformers: required unified multimodal model classes are unavailable
# before 5.10.1; 5.10.0 is yanked, and 5.15 has a known construction
# regression. transformers 5.9 also changed
# _adjust_missing_and_unexpected_keys from
# list+list to set|set, breaking models that declare
# _keys_to_ignore_on_load_unexpected as a list (e.g. OpenMOSS-Team/MOSS-TTS-Nano)
# with "TypeError: unsupported operand type(s) for |: 'list' and 'set'".
# That breakage is now handled in code instead of by pinning:
# MossTTSNanoForGeneration wraps weight loading in
# transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py).
transformers >= 5.5.3
transformers >= 5.10.1, < 5.15
av>=14.0.0
omegaconf>=2.3.0
diffusers==0.38.0
+36 -21
View File
@@ -39,8 +39,6 @@ def _render(changed_files: list[str]) -> str:
def test_bootstrap_if_injected_by_step_key() -> None:
import yaml
rendered = _render_bootstrap_pipeline(
BOOTSTRAP_STEPS_TEMPLATE,
decision=resolve_ci_decision([]),
@@ -142,12 +140,7 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None:
]
COVERAGE_PILOT_LABELS = ("Diffusion · Bagel Test", "TTS · Qwen3-TTS Base Test")
MERGE_PIPELINE = Path(".buildkite/cuda/test-merge.yml")
def _surviving_labels(changed_files: list[str]) -> set[str]:
doc = yaml.safe_load(MERGE_PIPELINE.read_text(encoding="utf-8"))
def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]:
rendered = _render_test_pipeline(doc, changed_files=changed_files)
labels: set[str] = set()
@@ -163,6 +156,29 @@ def _surviving_labels(changed_files: list[str]) -> set[str]:
return labels
# Synthetic coverage-style job: shared inputs that change what the split measures.
_COVERAGE_SHARED_INPUTS_DOC = {
"steps": [
{
"label": "Coverage Pilot",
"source_file_dependencies": [
"tests/e2e/online_serving/test_example.py",
".buildkite/common/scripts/run_cov_split.sh",
"pyproject.toml",
],
"commands": [".buildkite/common/scripts/run_cov_split.sh --model-id example"],
},
{
"label": "Unrelated Model Test",
"source_file_dependencies": [
"tests/e2e/online_serving/test_other.py",
],
"commands": ["pytest -sv tests/e2e/online_serving/test_other.py"],
},
],
}
@pytest.mark.parametrize(
"changed_file",
[
@@ -170,18 +186,17 @@ def _surviving_labels(changed_files: list[str]) -> set[str]:
"pyproject.toml",
],
)
def test_coverage_pilots_selected_by_their_shared_inputs(changed_file: str) -> None:
"""The split helper and the coverage config both change what these jobs measure.
Without them in source_file_dependencies the jobs are filtered out of normal
PR builds, so a regression would only surface in a later nightly.
"""
labels = _surviving_labels([changed_file])
for pilot in COVERAGE_PILOT_LABELS:
assert pilot in labels, f"{pilot} dropped when {changed_file} changed"
def test_coverage_shared_inputs_select_dependent_job(changed_file: str) -> None:
"""Jobs that list coverage shared inputs must stay selected when those files change."""
labels = _surviving_labels(_COVERAGE_SHARED_INPUTS_DOC, [changed_file])
assert "Coverage Pilot" in labels
assert "Unrelated Model Test" not in labels
def test_coverage_pilots_not_selected_by_unrelated_change() -> None:
labels = _surviving_labels(["vllm_omni/entrypoints/openai/serving_chat.py"])
for pilot in COVERAGE_PILOT_LABELS:
assert pilot not in labels
def test_coverage_shared_inputs_ignored_for_unrelated_change() -> None:
labels = _surviving_labels(
_COVERAGE_SHARED_INPUTS_DOC,
["vllm_omni/entrypoints/openai/serving_chat.py"],
)
assert "Coverage Pilot" not in labels
assert "Unrelated Model Test" not in labels
+89 -1
View File
@@ -14,7 +14,8 @@ import pytest
from transformers import PretrainedConfig, Qwen3OmniMoeConfig
from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage
from vllm_omni.config.config_factory import StageConfigFactory
from vllm_omni.config import config_factory as config_factory_module
from vllm_omni.config.config_factory import StageConfigFactory, _materialize_object_storage_configs
from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
from vllm_omni.config.omni_config import VllmOmniConfig
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline, resolve_pipeline_config
@@ -55,6 +56,7 @@ def clear_config_factory_caches():
yield
StageConfigFactory.get_hf_config.cache_clear()
StageConfigFactory.try_infer_model_type.cache_clear()
_materialize_object_storage_configs.cache_clear()
Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True)
@@ -2765,3 +2767,89 @@ class TestPipelineConfigResolvers:
pass
assert resolver(NotTheRightHfConfig()) is None
class TestObjectStorageConfigResolution:
"""Object-storage URIs must be materialized locally before any HF-style reads.
Regression coverage for review on the Run:AI PR: ``get_config`` and the
config.json/model_index.json fallbacks crashed with ``HFValidationError``
for ``s3://`` URIs, and the name-match fallback scanned the whole URI, so
a bucket named after another pipeline could hijack pipeline selection.
"""
@pytest.fixture
def fake_object_storage(self, monkeypatch, tmp_path):
"""Replace ObjectStorageModel with a fake that "pulls" into a tmp dir.
Returns a recorder: calls to the returned callable write files into
the materialized directory, and ``recorder.pulls`` records every
``pull_files`` invocation.
"""
pulls: list[tuple[str, list[str] | None, list[str] | None]] = []
materialized_files: list[tuple[str, str]] = []
class FakeObjectStorageModel:
def __init__(self, url: str):
self.dir = str(tmp_path)
def pull_files(self, model_path, allow_pattern=None, ignore_pattern=None):
pulls.append((model_path, allow_pattern, ignore_pattern))
for name, content in materialized_files:
(tmp_path / name).write_text(content)
monkeypatch.setattr(config_factory_module, "ObjectStorageModel", FakeObjectStorageModel)
class Recorder:
def set_files(self, files: list[tuple[str, str]]) -> None:
materialized_files.clear()
materialized_files.extend(files)
recorder = Recorder()
recorder.pulls = pulls
return recorder
def test_passthrough_for_non_uri(self, fake_object_storage):
assert _materialize_object_storage_configs("org/model") == "org/model"
assert _materialize_object_storage_configs("/local/model") == "/local/model"
assert fake_object_storage.pulls == []
def test_materialize_pulls_configs_once_per_uri(self, fake_object_storage):
uri = "s3://bucket/model"
first = _materialize_object_storage_configs(uri)
second = _materialize_object_storage_configs(uri)
assert first == second
assert len(fake_object_storage.pulls) == 1
pulled_uri, allow_pattern, ignore_pattern = fake_object_storage.pulls[0]
assert pulled_uri == uri
assert "*.json" in allow_pattern
assert ignore_pattern is None
def test_try_infer_model_type_reads_materialized_config(self, fake_object_storage):
"""A misleading bucket name must not beat the real config.json content."""
fake_object_storage.set_files(
[
(
"config.json",
'{"model_type": "qwen3_omni_moe", "architectures": ["Qwen3OmniMoeForConditionalGeneration"]}',
)
]
)
uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct"
model = StageConfigFactory.try_infer_model_type(model=uri, trust_remote_code=False)
assert model == "qwen3_omni_moe"
assert fake_object_storage.pulls # materialization happened
def test_name_match_fallback_scans_basename_only(self, fake_object_storage):
# Empty config.json (CosyVoice3 style) forces the name-match fallback.
fake_object_storage.set_files([("config.json", "{}")])
deceptive = "s3://qwen3-tts-models/plain-checkpoint"
assert StageConfigFactory.try_infer_model_type(model=deceptive, trust_remote_code=False) is None
matching = "s3://any-bucket/my-cosyvoice3-model"
assert StageConfigFactory.try_infer_model_type(model=matching, trust_remote_code=False) == "cosyvoice3"
+8 -8
View File
@@ -424,10 +424,10 @@
"baseline": {
"H100": {
"median_audio_ttfp_ms": [
82.1886
82.6344
],
"median_audio_rtf": [
0.1241
0.1277
]
}
}
@@ -453,13 +453,13 @@
"baseline": {
"H100": {
"median_audio_ttfp_ms": [
160.9168
170.3737
],
"median_audio_rtf": [
0.1907
0.2076
],
"audio_throughput": [
39.2332
37.0834
]
}
}
@@ -485,13 +485,13 @@
"baseline": {
"H100": {
"median_audio_ttfp_ms": [
100.1265
101.0293
],
"median_audio_rtf": [
0.1437
0.1498
],
"audio_throughput": [
11.7605
10.1186
]
}
}
@@ -5,6 +5,7 @@ from contextlib import nullcontext
import pytest
import torch
from torch.nn.attention import SDPBackend
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
@@ -12,6 +13,48 @@ from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@pytest.mark.parametrize(
("query_sequence_length", "key_sequence_length", "expected_backend"),
[
(1, 1, SDPBackend.MATH),
(8, 1, SDPBackend.MATH),
(1, 8, SDPBackend.MATH),
(8, 8, SDPBackend.CUDNN_ATTENTION),
],
)
def test_cudnn_routes_single_token_attention_to_math(
monkeypatch, query_sequence_length, key_sequence_length, expected_backend
):
selected_backends = []
def fake_sdpa_kernel(backends):
selected_backends.append(backends)
return nullcontext()
monkeypatch.setattr(
"vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
fake_sdpa_kernel,
)
monkeypatch.setattr(
torch.nn.functional,
"scaled_dot_product_attention",
lambda query, _key, _value, **_kwargs: query,
)
impl = CuDNNAttentionImpl(
num_heads=2,
head_size=4,
softmax_scale=0.5,
)
query = torch.randn(1, query_sequence_length, 2, 4)
key = torch.randn(1, key_sequence_length, 2, 4)
value = torch.randn_like(key)
output = impl.forward_cuda(query, key, value)
assert output.shape == query.shape
assert selected_backends == [[expected_backend]]
def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
observed = {}
+7 -1
View File
@@ -36,7 +36,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
SEPARATE_CFG_TRANSFORMERS = [
DreamIdOmniModel,
LTX2VideoTransformer3DModel,
HeliosTransformer3DModel,
LongCatImageTransformer2DModel,
Cosmos3VFMTransformer,
@@ -124,6 +123,13 @@ def test_cache_dit_configs_have_separate_cfg(transformer_model):
assert transformer_model._cache_dit_adapter_config.has_separate_cfg is True
def test_ltx2_cache_dit_uses_one_forward_per_denoise_step():
"""LTX batches all guidance passes into one Transformer invocation."""
adapter_config = LTX2VideoTransformer3DModel._cache_dit_adapter_config
assert adapter_config.has_separate_cfg is False
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit")
def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter):
@@ -53,14 +53,10 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
def _extract_single_image(outputs) -> Image.Image:
first_output = outputs[0]
assert first_output.final_output_type == "image"
if not hasattr(first_output, "request_output") or not first_output.request_output:
raise ValueError("No request_output found in OmniRequestOutput")
if not isinstance(first_output, OmniRequestOutput):
raise ValueError(f"Expected an OmniRequestOutput, got {type(first_output).__name__}")
req_out = first_output.request_output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
images = req_out.images
images = first_output.images
if images is None or len(images) != 1:
raise ValueError(f"Expected 1 image, got {0 if images is None else len(images)}")
return images[0]
@@ -72,8 +72,7 @@ def _to_float_array(data: Any) -> np.ndarray:
def _extract_output_array(output: Any) -> np.ndarray:
req_out = output.request_output
payload = req_out.images[0]
payload = output.images[0]
return _to_float_array(payload)
@@ -16,6 +16,11 @@ from vllm.config.load import LoadConfig
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.host_weight_plan import (
HostWeightPlan,
HostWeightPlanResult,
TensorBinding,
)
from vllm_omni.diffusion.models.helios import HeliosPipeline
from vllm_omni.diffusion.registry import initialize_model
@@ -110,6 +115,85 @@ def test_empty_source_prefix_keeps_full_model_strict_check():
loader.load_weights(model)
def test_stream_online_quant_weights_offloads_layers_after_processing():
from vllm.model_executor.model_loader.reload.layerwise import (
get_layerwise_info,
)
events: list[str] = []
class _OnlineQuantMethod:
uses_meta_device = True
class _TrackedLayer(nn.Linear):
def __init__(self, name: str):
super().__init__(2, 2, bias=False)
self.name = name
self.quant_method = _OnlineQuantMethod()
get_layerwise_info(self).load_numel_total = self.weight.numel()
def to(self, *args, **kwargs):
events.append(self.name)
return super().to(*args, **kwargs)
class _StreamingModel(nn.Module):
def __init__(self):
super().__init__()
self.first = _TrackedLayer("first")
self.second = _TrackedLayer("second")
model = _StreamingModel()
weights = iter(
[
("first.weight", torch.zeros((2, 2))),
("second.weight", torch.zeros((2, 2))),
]
)
streamed = DiffusersPipelineLoader._stream_online_quant_weights_to_cpu(model, weights)
assert next(streamed)[0] == "first.weight"
get_layerwise_info(model.first).reset()
assert next(streamed)[0] == "second.weight"
assert events == ["first"]
get_layerwise_info(model.second).reset()
with pytest.raises(StopIteration):
next(streamed)
assert events == ["first", "second"]
def test_process_weights_skips_completed_online_quant_layer(monkeypatch):
from unittest.mock import Mock
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
from vllm.model_executor.model_loader.reload import layerwise
class _TrackedLayer(nn.Linear):
def __init__(self):
super().__init__(2, 2, bias=False)
self.to_calls: list[object] = []
def to(self, *args, **kwargs):
self.to_calls.append(args[0] if args else kwargs.get("device"))
return super().to(*args, **kwargs)
model = nn.Module()
model.layer = _TrackedLayer()
quant_method = Mock(spec=QuantizeMethodBase)
quant_method.uses_meta_device = True
model.layer.quant_method = quant_method
model.layer._already_called_process_weights_after_loading = True
finalize = Mock()
monkeypatch.setattr(layerwise, "finalize_layerwise_processing", finalize)
loader = _make_loader_with_weights([])
loader._process_weights_after_loading(model, torch.device("cuda"))
finalize.assert_called_once_with(model, model_config=None)
quant_method.process_weights_after_loading.assert_not_called()
assert model.layer.to_calls == []
class _ConfigAwareModel(nn.Module):
def __init__(self, *, od_config):
super().__init__()
@@ -179,6 +263,137 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
assert get_current_diffusion_config_or_none() is None
def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
od_config = SimpleNamespace(
dtype=torch.float32,
parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
model="unused",
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
plan = HostWeightPlan(
backing_kind="checkpoint_mmap",
bindings={},
)
loaded_ordinary_weights = False
def load_weights(_model):
nonlocal loaded_ordinary_weights
loaded_ordinary_weights = True
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
loader.load_weights = load_weights # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(plan),
)
assert loader.load_model(load_device="cpu") is model
assert not loaded_ordinary_weights
assert loader.take_host_weight_plan() is plan
assert loader.take_host_weight_plan() is None
def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
od_config = SimpleNamespace(
dtype=torch.float32,
parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
model="unused",
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
class MixedSourceModel(nn.Module):
def __init__(self):
super().__init__()
self.transformer = nn.Linear(2, 2, bias=False)
self.text_encoder = nn.Linear(2, 2, bias=False)
self.weights_sources = (
DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="transformer."),
DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="text_encoder."),
)
self.loaded_weight_names: list[str] = []
def load_weights(self, weights):
self.loaded_weight_names = [name for name, _ in weights]
return set(self.loaded_weight_names)
model = MixedSourceModel()
plan = HostWeightPlan(
backing_kind="checkpoint_mmap",
bindings={
"transformer.weight": TensorBinding(
checkpoint_key="weight",
file_path="unused",
)
},
planned_source_prefixes=frozenset({"transformer."}),
)
requested_prefixes: list[str] = []
def get_weights(source, model=None):
del model
requested_prefixes.append(source.prefix)
yield source.prefix + "weight", torch.ones(2, 2)
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
loader._get_weights_iterator = get_weights # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(plan),
)
assert loader.load_model(load_device="cpu") is model
assert requested_prefixes == ["text_encoder."]
assert model.loaded_weight_names == ["text_encoder.weight"]
assert loader.take_host_weight_plan() is plan
def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
od_config = SimpleNamespace(
dtype=torch.float32,
parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
model="unused",
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
calls: list[str] = []
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
loader.load_weights = lambda _model: calls.append("load") # type: ignore[method-assign]
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"),
)
assert loader.load_model(load_device="cpu") is model
assert calls == ["load", "process"]
assert loader.take_host_weight_plan() is None
def test_hsdp_processes_quantized_weights_before_sharding(mocker):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
from vllm_omni.diffusion.offloader.module_collector import PipelineModules
@@ -0,0 +1,72 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import contextlib
import huggingface_hub
import pytest
from vllm_omni.diffusion.model_loader import hub_prefetch
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def test_prefetch_subfolders_propagates_revision(monkeypatch):
calls = []
def fake_snapshot_download(**kwargs):
calls.append(kwargs)
monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext())
hub_prefetch.prefetch_subfolders(
"org/model",
["transformer"],
revision="pinned-revision",
)
assert calls == [
{
"repo_id": "org/model",
"revision": "pinned-revision",
"allow_patterns": ["transformer/*", "transformer/**", "*.json", "*.txt"],
}
]
def test_from_pretrained_retry_preserves_revision(monkeypatch):
factory_calls = []
prefetch_calls = []
def fake_factory(model, **kwargs):
factory_calls.append((model, kwargs))
if len(factory_calls) == 1:
raise OSError("partial cache")
return "loaded"
def fake_prefetch(model, subfolders, **kwargs):
prefetch_calls.append((model, tuple(subfolders), kwargs))
monkeypatch.setattr(hub_prefetch, "prefetch_subfolders", fake_prefetch)
monkeypatch.setattr(hub_prefetch.time, "sleep", lambda _seconds: None)
result = hub_prefetch.from_pretrained_with_prefetch(
fake_factory,
"org/model",
subfolder="transformer",
prefetch_list=["transformer", "vae"],
revision="pinned-revision",
max_attempts=2,
)
assert result == "loaded"
assert len(factory_calls) == 2
assert all(call[1]["revision"] == "pinned-revision" for call in factory_calls)
assert prefetch_calls == [
(
"org/model",
("transformer", "vae"),
{"local_files_only": False, "revision": "pinned-revision"},
)
]
@@ -7,17 +7,17 @@ from vllm_omni.diffusion.models.flux2.pipeline_flux2 import (
Flux2Pipeline,
_resolve_component_quant_config,
)
from vllm_omni.quantization import ComponentQuantizationConfig
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_resolve_component_quant_config_routes_components():
text_encoder_config = object()
component_config = Mock()
component_config.resolve.side_effect = {
"text_encoder": text_encoder_config,
"transformer": None,
}.get
component_config = ComponentQuantizationConfig(
component_configs={"text_encoder": text_encoder_config},
default_config=None,
)
assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config
assert _resolve_component_quant_config(component_config, "transformer") is None
@@ -0,0 +1,881 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Focused LTX-2.5 pipeline correctness tests."""
import json
from dataclasses import replace
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.models.ltx2 import ltx2_components, ltx2_latents, ltx2_runtime
from vllm_omni.diffusion.models.ltx2.ltx2_components import (
LTX25_DISTILLED_COMPONENT_PROFILE,
LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
LTX25_FULL_COMPONENT_PROFILE,
LTX25_TWO_STAGE_COMPONENT_PROFILE,
detect_ltx_model_version,
resolve_ltx_checkpoint_kind,
resolve_ltx_component_profile,
)
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
LTXPhaseResult,
_first_frame_keyframes_mask,
_official_ltx_sigmas,
)
from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
LTX_GUIDANCE_EXECUTOR,
LTXGuidanceSpec,
LTXModalityGuidance,
velocity_from_x0,
x0_from_velocity,
)
from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
LTX2_DISTILLED_TWO_STAGE_RECIPE,
LTX2_ONE_STAGE_RECIPE,
LTX23_ONE_STAGE_RECIPE,
LTX25_DISTILLED_ONE_STAGE_RECIPE,
LTX25_DISTILLED_TWO_STAGE_RECIPE,
LTX25_FULL_RECIPE,
LTX25_TWO_STAGE_RECIPE,
LTX_DISTILLED_SIGMAS,
LTX_POSITIVE_ONLY_RECIPE,
resolve_ltx_pipeline_recipe,
)
from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs
from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2DistilledOneStagePipeline,
LTX2Pipeline,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
LTX2DistilledTwoStagePipeline,
LTX2TwoStagePipeline,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.mark.parametrize("model_version", ["2", "2.3", "2.5"])
def test_ltx_timestep_adaln_expands_per_token_like_official_without_sp(model_version):
pipe = object.__new__(LTX2Pipeline)
pipe.model_version = model_version
pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1))
ts = torch.tensor([750.0, 500.0])
kwargs = LTXRuntime._denoise_timestep_kwargs(
pipe,
ts,
SimpleNamespace(),
SimpleNamespace(),
video_token_count=4,
audio_token_count=3,
)
torch.testing.assert_close(kwargs["timestep"], ts[:, None].expand(-1, 4))
torch.testing.assert_close(kwargs["audio_timestep"], ts[:, None].expand(-1, 3))
assert kwargs["sigma"] is ts
assert kwargs["audio_sigma"] is ts
def test_ltx_ancestral_step_clears_sp_audio_padding_after_scheduler_update(monkeypatch):
state = LTXAVState(
video=torch.zeros(1, 2, 2),
audio=torch.zeros(1, 4, 2),
)
forward_ctx = SimpleNamespace(
original_audio_num_frames=3,
sampler="euler_ancestral",
)
denoise_ctx = SimpleNamespace(latents=None, audio_latents=None)
pipe = SimpleNamespace(
_predict_noise_for_step=lambda *_args: (torch.zeros_like(state.video), torch.zeros_like(state.audio))
)
def fake_step(_pipeline, actual_forward_ctx, _denoise_ctx, *_args):
assert actual_forward_ctx.sampler == "euler_ancestral"
updated_audio = torch.ones_like(state.audio)
updated_audio[:, 3:] = 99.0
return torch.ones_like(state.video), updated_audio
monkeypatch.setattr(ltx2_runtime, "step_denoised_latents", fake_step)
actual = LTXRuntime._denoise_step(
pipe,
0,
torch.tensor(1.0),
state,
forward_ctx,
denoise_ctx,
)
torch.testing.assert_close(actual.audio[:, :3], torch.ones_like(actual.audio[:, :3]))
torch.testing.assert_close(actual.audio[:, 3:], torch.zeros_like(actual.audio[:, 3:]))
def test_ltx25_first_causal_latent_frame_is_marked_for_keyframe_embedding():
reference = torch.zeros(2, 12, 8)
mask = _first_frame_keyframes_mask(reference, latent_num_frames=3)
assert mask.shape == (2, 12, 1)
torch.testing.assert_close(mask[:, :4], torch.ones(2, 4, 1))
torch.testing.assert_close(mask[:, 4:], torch.zeros(2, 8, 1))
def _make_ltx_request_pipe(cls):
pipe = object.__new__(cls)
torch.nn.Module.__init__(pipe)
pipe.device = torch.device("cpu")
pipe.tokenizer_max_length = 99
pipe.vae_spatial_compression_ratio = 32
pipe.vae_temporal_compression_ratio = 8
return pipe
def _resolve_request_inputs_for_test(
pipe,
req,
*,
guidance_scale=4.0,
negative_prompt=None,
):
return pipe._resolve_request_inputs(
req,
prompt=None,
negative_prompt=negative_prompt,
height=None,
width=None,
num_frames=None,
frame_rate=None,
num_inference_steps=None,
guidance_scale=guidance_scale,
num_videos_per_prompt=1,
generator=None,
latents=None,
audio_latents=None,
prompt_embeds=None,
negative_prompt_embeds=None,
prompt_attention_mask=None,
negative_prompt_attention_mask=None,
decode_timestep=0.0,
decode_noise_scale=None,
output_type="np",
max_sequence_length=None,
)
def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch):
pipe = SimpleNamespace(component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None))
od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16)
monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
monkeypatch.setattr(ltx2_components, "prefetch_subfolders", lambda *_args, **_kwargs: None)
monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", lambda *_args, **_kwargs: object())
with pytest.raises(ImportError) as exc_info:
ltx2_components.initialize_pipeline_components(pipe, od_config)
assert str(exc_info.value) == (
"LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15."
)
def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
revision = "pinned-revision"
calls = {"components": []}
profile = replace(
LTX25_DISTILLED_COMPONENT_PROFILE,
text_encoder_cls=object,
video_vae_cls=object,
vocoder_cls=object,
vocoder_fallback_cls=None,
scheduler_use_dynamic_shifting=False,
)
pipeline = SimpleNamespace(component_profile=profile, pipeline_kind="distilled_two_stage")
od_config = SimpleNamespace(
model="org/converted-ltx25",
revision=revision,
dtype=torch.bfloat16,
quantization_config=None,
)
monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
monkeypatch.setattr(ltx2_components, "_install_connector_attention", lambda _connectors, **_kwargs: None)
monkeypatch.setattr(ltx2_components, "_place_aux_components", lambda _pipeline: None)
def fake_prefetch(model, subfolders, **kwargs):
calls["prefetch"] = (model, tuple(subfolders), kwargs)
def fake_tokenizer(*args, **kwargs):
calls["tokenizer"] = (args, kwargs)
return SimpleNamespace(model_max_length=1024)
def fake_component(component_cls, model, subfolder, **kwargs):
calls["components"].append((component_cls, model, subfolder, kwargs))
if subfolder == "vae":
return SimpleNamespace(spatial_compression_ratio=32, temporal_compression_ratio=8)
if subfolder == "audio_vae":
return SimpleNamespace(
mel_compression_ratio=4,
temporal_compression_ratio=4,
config=SimpleNamespace(sample_rate=16_000, mel_hop_length=160),
)
return object()
def fake_transformer_config(model, subfolder, local_files_only, *, revision):
calls["transformer_config"] = (model, subfolder, local_files_only, revision)
return {"component": "transformer"}
monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch)
monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer)
monkeypatch.setattr(ltx2_components, "_load_component", fake_component)
monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config)
monkeypatch.setattr(
ltx2_components,
"create_transformer_from_config",
lambda *_args, **_kwargs: SimpleNamespace(config=SimpleNamespace(patch_size=1, patch_size_t=1)),
)
def fake_scheduler(*args, **kwargs):
calls["scheduler"] = (args, kwargs)
return SimpleNamespace(config={"use_dynamic_shifting": False, "shift_terminal": None})
monkeypatch.setattr(ltx2_components.FlowMatchEulerDiscreteScheduler, "from_pretrained", fake_scheduler)
ltx2_components.initialize_pipeline_components(pipeline, od_config)
assert pipeline.weights_sources[0].revision == revision
assert calls["prefetch"][2]["revision"] == revision
assert calls["tokenizer"][1]["revision"] == revision
assert {call[2] for call in calls["components"]} == {
"text_encoder",
"connectors",
"vae",
"audio_vae",
"vocoder",
"latent_upsampler",
}
assert all(call[3]["revision"] == revision for call in calls["components"])
assert calls["transformer_config"] == (
od_config.model,
profile.transformer_subfolder,
False,
revision,
)
assert calls["scheduler"][1]["revision"] == revision
def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path):
(tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"}))
(tmp_path / "transformer").mkdir()
(tmp_path / "transformer" / "config.json").write_text(json.dumps({"ff_bias": False}))
assert detect_ltx_model_version(str(tmp_path)) == "2.3"
def test_ltx25_checkpoint_selects_full_one_stage_profile(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(
json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
)
def stub_components(pipe, od_config):
pipe.od_config = od_config
pipe.vae_spatial_compression_ratio = 32
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
pipe = LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False))
assert pipe.model_version == "2.5"
assert pipe.component_profile is LTX25_FULL_COMPONENT_PROFILE
assert pipe.pipeline_recipe is LTX25_FULL_RECIPE
assert pipe.preserve_sp_padded_audio_duration
assert pipe.reports_stage_durations
def test_ltx25_full_recipe_matches_official_sft_defaults():
(phase,) = LTX25_FULL_RECIPE.phases
assert (LTX25_FULL_RECIPE.height, LTX25_FULL_RECIPE.width) == (544, 960)
assert LTX25_FULL_RECIPE.num_frames == 121
assert LTX25_FULL_RECIPE.frame_rate == 24.0
assert LTX25_FULL_RECIPE.num_inference_steps == 30
assert phase.guidance.video == LTXModalityGuidance(
cfg_scale=3.0,
stg_scale=1.0,
modality_scale=3.0,
rescale_scale=0.7,
stg_blocks=(28,),
)
assert phase.guidance.audio == LTXModalityGuidance(
cfg_scale=7.0,
stg_scale=1.0,
modality_scale=3.0,
rescale_scale=0.7,
stg_blocks=(28,),
)
assert phase.sigmas is None
assert phase.noise_scale == 1.0
assert phase.use_official_sigma_schedule
assert LTX25_FULL_COMPONENT_PROFILE.transformer_subfolder == "transformer_full"
assert LTX25_FULL_COMPONENT_PROFILE.scheduler_use_dynamic_shifting
assert LTX25_FULL_COMPONENT_PROFILE.scheduler_shift_terminal == 0.1
def test_ltx25_full_sigma_schedule_uses_official_default_sequence_anchor():
scheduler = SimpleNamespace(
config={
"base_image_seq_len": 1024,
"max_image_seq_len": 4096,
"base_shift": 0.95,
"max_shift": 2.05,
"shift_terminal": 0.1,
}
)
sigmas = _official_ltx_sigmas(
scheduler,
steps=30,
device=torch.device("cpu"),
)
torch.testing.assert_close(
sigmas[[0, 1, -2, -1]],
torch.tensor([1.0, 0.9949570, 0.1, 0.0]),
rtol=1e-5,
atol=1e-6,
)
def test_ltx25_checkpoint_selects_distilled_two_stage_profile(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(
json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
)
def stub_components(pipe, od_config):
pipe.od_config = od_config
pipe.vae_spatial_compression_ratio = 32
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
pipe = LTX2DistilledTwoStagePipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
enable_diffusion_pipeline_profiler=False,
)
)
assert pipe.component_profile is LTX25_DISTILLED_COMPONENT_PROFILE
assert pipe.pipeline_recipe is LTX25_DISTILLED_TWO_STAGE_RECIPE
def test_ltx25_full_config_enrichment_uses_selected_transformer_subfolder(tmp_path):
from vllm_omni.diffusion.data import OmniDiffusionConfig
(tmp_path / "model_index.json").write_text(json.dumps({"_class_name": "LTX2Pipeline", "model_version": "2.5"}))
for subfolder, marker in (("transformer", "distilled"), ("transformer_full", "full")):
path = tmp_path / subfolder
path.mkdir()
(path / "config.json").write_text(json.dumps({"marker": marker}))
config = OmniDiffusionConfig(model=str(tmp_path), model_class_name="LTX2Pipeline")
config.enrich_config()
assert config.tf_model_config.get("marker") == "full"
def test_ltx25_four_public_pipeline_semantics_are_disjoint():
assert resolve_ltx_component_profile("one_stage", "2.5") is LTX25_FULL_COMPONENT_PROFILE
assert resolve_ltx_pipeline_recipe("one_stage", "2.5") is LTX25_FULL_RECIPE
assert resolve_ltx_component_profile("two_stage", "2.5") is LTX25_TWO_STAGE_COMPONENT_PROFILE
assert resolve_ltx_pipeline_recipe("two_stage", "2.5") is LTX25_TWO_STAGE_RECIPE
assert resolve_ltx_component_profile("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
assert resolve_ltx_pipeline_recipe("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_RECIPE
assert resolve_ltx_component_profile("distilled_two_stage", "2.5") is LTX25_DISTILLED_COMPONENT_PROFILE
assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.5") is LTX25_DISTILLED_TWO_STAGE_RECIPE
assert LTX2Pipeline.pipeline_kind == "one_stage"
assert LTX2TwoStagePipeline.pipeline_kind == "two_stage"
assert LTX2DistilledOneStagePipeline.pipeline_kind == "distilled_one_stage"
assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage"
@pytest.mark.parametrize(
("pipeline_kind", "expected"),
[
("one_stage", "regular"),
("two_stage", "regular"),
("distilled_one_stage", "distilled"),
("distilled_two_stage", "distilled"),
("dmd2", None),
],
)
def test_ltx_checkpoint_kind_is_derived_from_pipeline_kind(pipeline_kind, expected):
assert resolve_ltx_checkpoint_kind(pipeline_kind) == expected
def test_ltx_checkpoint_kind_rejects_unknown_pipeline():
with pytest.raises(ValueError, match="Unsupported LTX pipeline kind"):
resolve_ltx_checkpoint_kind("unknown")
def test_ltx25_distilled_two_stage_recipe_matches_model_card_resolution():
stage1, stage2 = LTX25_DISTILLED_TWO_STAGE_RECIPE.phases
assert (LTX25_DISTILLED_TWO_STAGE_RECIPE.height, LTX25_DISTILLED_TWO_STAGE_RECIPE.width) == (1088, 1920)
assert stage1.spatial_downscale == 2
assert stage1.sigmas == LTX_DISTILLED_SIGMAS
assert stage1.sampler == "euler_ancestral"
assert stage2.sampler == "euler"
assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0)
assert stage2.input_transform == "spatial_upsample"
assert LTX25_DISTILLED_TWO_STAGE_RECIPE.negative_prompt == ""
assert not LTX25_DISTILLED_TWO_STAGE_RECIPE.allow_negative_prompt
def test_ltx25_distilled_one_stage_uses_native_low_resolution():
(phase,) = LTX25_DISTILLED_ONE_STAGE_RECIPE.phases
assert (LTX25_DISTILLED_ONE_STAGE_RECIPE.height, LTX25_DISTILLED_ONE_STAGE_RECIPE.width) == (544, 960)
assert phase.spatial_downscale == 1
assert phase.sigmas == LTX_DISTILLED_SIGMAS
assert phase.sampler == "euler_ancestral"
assert not LTX25_DISTILLED_ONE_STAGE_RECIPE.supports_cache_dit
def test_ltx25_regular_two_stage_uses_full_transformer_then_distilled_lora():
stage1, stage2 = LTX25_TWO_STAGE_RECIPE.phases
assert (LTX25_TWO_STAGE_RECIPE.height, LTX25_TWO_STAGE_RECIPE.width) == (1088, 1920)
assert LTX25_TWO_STAGE_RECIPE.num_inference_steps == 30
assert stage1.spatial_downscale == 2
assert stage1.guidance == LTX25_FULL_RECIPE.request_guidance
assert stage2.adapter_slot == "ltx_distilled"
assert LTX25_TWO_STAGE_RECIPE.allow_request_phase_sigmas
assert LTX25_TWO_STAGE_COMPONENT_PROFILE.transformer_subfolder == "transformer_full"
assert LTX25_TWO_STAGE_COMPONENT_PROFILE.distilled_lora_filename is not None
@pytest.mark.parametrize(
"recipe",
[
LTX2_ONE_STAGE_RECIPE,
LTX23_ONE_STAGE_RECIPE,
LTX_POSITIVE_ONLY_RECIPE,
],
)
def test_ltx_one_stage_recipes_declare_cache_dit_supported(recipe):
assert recipe.supports_cache_dit
@pytest.mark.parametrize(
"recipe",
[
LTX25_FULL_RECIPE,
LTX25_DISTILLED_ONE_STAGE_RECIPE,
LTX2_DISTILLED_TWO_STAGE_RECIPE,
LTX25_DISTILLED_TWO_STAGE_RECIPE,
],
)
def test_ltx25_and_multistage_recipes_declare_cache_dit_unsupported(recipe):
assert not recipe.supports_cache_dit
def test_ltx_positive_only_guidance_preserves_official_x0_roundtrip():
sample = torch.full((1, 1), 1e8, dtype=torch.float32)
velocity = torch.ones((1, 1), dtype=torch.bfloat16)
sigma = torch.tensor(0.5)
actual = LTX_GUIDANCE_EXECUTOR._guide_modality(
sample,
{"cond": velocity},
sigma,
LTXModalityGuidance(),
)
expected = velocity_from_x0(sample, x0_from_velocity(sample, velocity, sigma), sigma)
assert actual.dtype == torch.float32
torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
def test_ltx_velocity_from_x0_materializes_official_scalar_sigma():
item_calls = []
class TrackingTensor(torch.Tensor):
@staticmethod
def __new__(cls):
return torch.Tensor._make_subclass(cls, torch.tensor(0.725), False)
def item(self):
item_calls.append(True)
return super().item()
sample = torch.tensor([[1.0]], dtype=torch.bfloat16)
x0 = torch.tensor([[0.5]], dtype=torch.bfloat16)
velocity_from_x0(sample, x0, TrackingTensor())
assert item_calls == [True]
def test_ltx_ancestral_positive_only_guidance_preserves_raw_velocity():
sample = torch.full((1, 1), 1e8, dtype=torch.float32)
velocity = torch.ones((1, 1), dtype=torch.bfloat16)
actual = LTX_GUIDANCE_EXECUTOR._guide_modality(
sample,
{"cond": velocity},
torch.tensor(0.5),
LTXModalityGuidance(),
preserve_positive_velocity=True,
)
assert actual is velocity
@pytest.mark.parametrize("model_version", ["2.5", "2.3"])
def test_ltx_t2v_denoise_state_uses_official_bfloat16_dtype(model_version):
expected_dtype = torch.bfloat16
captured = {}
def prepare_latents(**kwargs):
captured["video_dtype"] = kwargs["dtype"]
return torch.empty(1, 1, 1, dtype=kwargs["dtype"])
def prepare_audio_latents(*_args, **kwargs):
captured["audio_dtype"] = kwargs["dtype"]
return torch.empty(1, 1, 1, dtype=kwargs["dtype"]), 1, 1
pipeline = SimpleNamespace(
model_version=model_version,
transformer=SimpleNamespace(config=SimpleNamespace(in_channels=128)),
prepare_latents=prepare_latents,
prepare_audio_latents=prepare_audio_latents,
audio_sampling_rate=48_000,
audio_hop_length=160,
audio_vae_temporal_compression_ratio=4,
audio_vae_mel_compression_ratio=4,
audio_vae=SimpleNamespace(config=SimpleNamespace(mel_bins=64, latent_channels=8)),
_resolve_audio_latent_length=lambda length, _latents: length,
)
request = SimpleNamespace(
num_videos_per_prompt=1,
height=64,
width=64,
num_frames=1,
frame_rate=24.0,
generator=None,
latents=None,
audio_latents=None,
audio_latents_normalized=False,
)
prompt = SimpleNamespace(
batch_size=1,
positive_connector_prompt_embeds=torch.empty(1, dtype=torch.bfloat16),
positive_connector_audio_prompt_embeds=torch.empty(1, dtype=torch.bfloat16),
)
LTXRuntime._prepare_video_latents_stage(
pipeline,
request,
prompt,
device=torch.device("cpu"),
noise_scale=1.0,
)
LTXRuntime._prepare_audio_latents_stage(
pipeline,
request,
prompt,
device=torch.device("cpu"),
noise_scale=1.0,
)
assert captured == {"video_dtype": expected_dtype, "audio_dtype": expected_dtype}
def test_ltx_seeded_latents_match_official_packed_rng_layout():
pipeline = SimpleNamespace(
model_version="2.5",
vae_spatial_compression_ratio=2,
vae_temporal_compression_ratio=2,
transformer_spatial_patch_size=1,
transformer_temporal_patch_size=1,
audio_vae_mel_compression_ratio=2,
od_config=SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)),
)
device = torch.device("cpu")
actual_generator = torch.Generator(device=device).manual_seed(123)
expected_generator = torch.Generator(device=device).manual_seed(123)
actual_video = ltx2_latents.prepare_video_latents(
pipeline,
batch_size=1,
num_channels_latents=2,
height=4,
width=6,
num_frames=3,
dtype=torch.float32,
device=device,
generator=actual_generator,
)
actual_audio, original_length, padded_length = ltx2_latents.prepare_audio_latents(
pipeline,
batch_size=1,
num_channels_latents=2,
audio_latent_length=3,
num_mel_bins=4,
dtype=torch.float32,
device=device,
generator=actual_generator,
)
expected_video = torch.randn((1, 12, 2), generator=expected_generator)
expected_audio = torch.randn((1, 3, 4), generator=expected_generator)
torch.testing.assert_close(actual_video, expected_video, rtol=0.0, atol=0.0)
torch.testing.assert_close(actual_audio, expected_audio, rtol=0.0, atol=0.0)
assert (original_length, padded_length) == (3, 3)
def test_ltx25_distilled_two_stage_executes_custom_phase_schedules():
request_inputs = LTXRequestInputs(
prompt="prompt",
negative_prompt="",
height=64,
width=64,
num_frames=1,
frame_rate=24.0,
num_inference_steps=8,
guidance=LTXGuidanceSpec.positive_only(),
num_videos_per_prompt=1,
generator=None,
latents=None,
audio_latents=None,
prompt_embeds=None,
negative_prompt_embeds=None,
prompt_attention_mask=None,
negative_prompt_attention_mask=None,
decode_timestep=0.25,
decode_noise_scale=0.5,
output_type="np",
max_sequence_length=16,
)
prompt_context = object()
phase_calls = []
def resolve_request_inputs(req, **kwargs):
return request_inputs
def run_phase(req, inputs, *, prompt_context=None, **kwargs):
phase_recipe = kwargs["phase_recipe"]
assert kwargs["image"] is source_image
phase_calls.append((phase_recipe, inputs, prompt_context))
if len(phase_calls) == 1:
assert prompt_context is None
assert inputs.decode_timestep == 0.25
assert inputs.decode_noise_scale == 0.5
context = prompt_context_sentinel
assert phase_recipe.name == "generate_lowres"
assert (inputs.height, inputs.width) == (32, 32)
assert inputs.num_inference_steps == 2
assert kwargs["sigmas"] == stage_1_sigmas
assert kwargs["noise_scale"] == 1.0
video = torch.ones(1, 128, 1, 1, 1)
audio = torch.full((1, 8, 1, 2), 2.0)
audio_for_next_phase = torch.full((1, 8, 1, 2), 5.0)
else:
assert prompt_context is prompt_context_sentinel
assert phase_recipe.name == "refine"
assert (inputs.height, inputs.width) == (64, 64)
assert inputs.latents.shape == (1, 128, 1, 2, 2)
torch.testing.assert_close(inputs.audio_latents, torch.full((1, 8, 1, 2), 5.0))
assert inputs.audio_latents_normalized
assert inputs.guidance_scale == 1.0
assert inputs.num_inference_steps == 2
assert kwargs["sigmas"] == stage_2_sigmas
assert kwargs["noise_scale"] == stage_2_sigmas[0]
assert inputs.decode_timestep == 0.0
assert inputs.decode_noise_scale is None
context = prompt_context
video = torch.full((1, 128, 1, 2, 2), 3.0)
audio = torch.full((1, 8, 1, 2), 4.0)
audio_for_next_phase = None
return LTXPhaseResult(
forward_context=SimpleNamespace(prompt_context=context),
video=video,
audio=audio,
audio_for_next_phase=audio_for_next_phase,
)
def decode_phase(phase):
return DiffusionOutput(output=(phase.video, phase.audio))
class FakeUpsampler(torch.nn.Module):
dtype = torch.float32
def forward(self, latents):
return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
prompt_context_sentinel = prompt_context
stage_1_sigmas = [1.0, 0.5, 0.0]
stage_2_sigmas = [0.8, 0.2, 0.0]
pipeline = object.__new__(LTX2TwoStagePipeline)
torch.nn.Module.__init__(pipeline)
pipeline.pipeline_recipe = LTX25_DISTILLED_TWO_STAGE_RECIPE
pipeline.device = torch.device("cpu")
pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8
pipeline.latent_upsampler = FakeUpsampler()
object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
object.__setattr__(pipeline, "run_phase", run_phase)
object.__setattr__(pipeline, "decode_phase", decode_phase)
source_image = object()
output = pipeline.forward(
SimpleNamespace(sampling_params_list=[], prompts=[]),
image=source_image,
stage_1_sigmas=stage_1_sigmas,
stage_2_sigmas=stage_2_sigmas,
)
assert len(phase_calls) == 2
assert phase_calls[1][2] is prompt_context_sentinel
torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0))
torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0))
class TestLTXRequestParsing:
def test_request_resolves_custom_sigmas_from_extra_args(self):
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
sigmas = [1.0, 0.5, 0.0]
req = DiffusionRequestBatch(
[
OmniDiffusionRequest(
prompt="prompt",
sampling_params=OmniDiffusionSamplingParams(extra_args={"sigmas": sigmas}),
request_id="ltx-custom-sigmas",
)
]
)
assert _make_ltx_request_pipe(LTX2Pipeline)._resolve_request_sigmas(req, None) == sigmas
def test_request_resolves_per_stage_sigmas_from_extra_args(self):
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
stage_1_sigmas = [1.0, 0.4, 0.0]
req = DiffusionRequestBatch(
[
OmniDiffusionRequest(
prompt="prompt",
sampling_params=OmniDiffusionSamplingParams(extra_args={"stage_1_sigmas": stage_1_sigmas}),
request_id="ltx-custom-phase-sigmas",
)
]
)
resolved = _make_ltx_request_pipe(LTX2TwoStagePipeline)._resolve_request_phase_sigmas(
req, [1.0, 0.0], [0.8, 0.0]
)
assert resolved == (stage_1_sigmas, [0.8, 0.0])
@pytest.mark.parametrize("value", [-1, 52, 1.5, True])
def test_request_rejects_invalid_image_crf(self, value):
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
req = DiffusionRequestBatch(
[
OmniDiffusionRequest(
prompt="prompt",
sampling_params=OmniDiffusionSamplingParams(extra_args={"image_crf": value}),
request_id="ltx-invalid-image-crf",
)
]
)
with pytest.raises(ValueError, match="image_crf"):
_resolve_request_inputs_for_test(_make_ltx_request_pipe(LTX2Pipeline), req)
class TestLTXForwardStages:
def test_gemma_tokenization_matches_official_bos_and_left_padding(self):
pipe = _make_ltx_request_pipe(LTX2Pipeline)
class FakeTokenizer:
bos_token_id = 2
eos_token_id = 1
eos_token = "<eos>"
pad_token_id = 0
pad_token = "<pad>"
padding_side = "right"
def __init__(self):
self.calls = []
def __call__(self, text, **kwargs):
self.calls.append((text, kwargs))
ids = [7, 8, 9, self.eos_token_id] if text == "long" else [6, self.eos_token_id]
return SimpleNamespace(input_ids=ids)
def pad(self, encoded, *, padding, max_length, return_tensors, return_attention_mask):
assert padding == "max_length"
assert return_tensors == "pt"
assert return_attention_mask is True
rows = []
masks = []
for ids in encoded["input_ids"]:
pad = max_length - len(ids)
rows.append([self.pad_token_id] * pad + ids)
masks.append([0] * pad + [1] * len(ids))
return SimpleNamespace(input_ids=torch.tensor(rows), attention_mask=torch.tensor(masks))
class FakeTextEncoder(torch.nn.Module):
dtype = torch.float32
def __init__(self):
super().__init__()
self.seen_input_ids = None
self.seen_attention_mask = None
def forward(self, *, input_ids, attention_mask, output_hidden_states):
assert output_hidden_states is True
self.seen_input_ids = input_ids.clone()
self.seen_attention_mask = attention_mask.clone()
return SimpleNamespace(hidden_states=(input_ids.unsqueeze(-1).float(),))
tokenizer = FakeTokenizer()
text_encoder = FakeTextEncoder()
object.__setattr__(pipe, "tokenizer", tokenizer)
object.__setattr__(pipe, "text_encoder", text_encoder)
pipe._get_gemma_prompt_embeds(
[" long ", "short"],
max_sequence_length=4,
device=torch.device("cpu"),
dtype=torch.float32,
)
torch.testing.assert_close(text_encoder.seen_input_ids, torch.tensor([[2, 7, 8, 9], [0, 2, 6, 1]]))
torch.testing.assert_close(text_encoder.seen_attention_mask, torch.tensor([[1, 1, 1, 1], [0, 1, 1, 1]]))
assert tokenizer.padding_side == "left"
assert [call[0] for call in tokenizer.calls] == ["long", "short"]
@@ -3,6 +3,7 @@
"""Unit tests for LTX image-to-video input and conditioning behavior."""
import sys
from types import SimpleNamespace
import pytest
@@ -23,7 +24,7 @@ def _make_ltx23_request_pipe(cls):
class TestLTXImageToVideoForwardStages:
def test_i2v_pil_preprocessing_matches_official_center_crop(self):
def test_legacy_i2v_pil_preprocessing_preserves_aspect_ratio_and_center_crops(self):
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _preprocess_i2v_pil_images
pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
@@ -35,6 +36,136 @@ class TestLTXImageToVideoForwardStages:
torch.testing.assert_close(actual, expected)
def test_ltx25_i2v_pil_preprocessing_preserves_aspect_ratio_after_crf(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning
pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
pixels[:, :, 0] = torch.arange(8, dtype=torch.uint8)
image = Image.fromarray(pixels.numpy())
monkeypatch.setattr(conditioning, "_apply_image_conditioning_crf", lambda image_array, _crf: image_array)
actual = conditioning._preprocess_i2v_pil_images(image, height=4, width=4, crf=18)
expected = pixels[:, 2:6].permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1.0
torch.testing.assert_close(actual, expected)
def test_ltx25_i2v_applies_crf18_before_resize(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
captured_crfs = []
captured_prepare_kwargs = []
def fake_preprocess(image, *, height, width, crf, device, dtype):
del image, device
captured_crfs.append(crf)
return torch.zeros(1, 3, height, width, dtype=dtype)
def fake_prepare_latents(**kwargs):
captured_prepare_kwargs.append(kwargs)
return kwargs["image"], None
monkeypatch.setattr(conditioning, "_preprocess_i2v_pil_images", fake_preprocess)
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.model_version = "2.5"
object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4)))
object.__setattr__(pipe, "prepare_latents", fake_prepare_latents)
request_inputs = SimpleNamespace(
height=4,
width=4,
num_frames=3,
num_videos_per_prompt=1,
generator=None,
latents=None,
)
prompt_context = SimpleNamespace(
batch_size=1,
positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16),
)
pipe._prepare_video_latents_stage(
request_inputs,
prompt_context,
device=torch.device("cpu"),
noise_scale=0.0,
image=Image.new("RGB", (8, 4)),
)
assert captured_crfs == [18]
assert captured_prepare_kwargs[0]["dtype"] is torch.bfloat16
assert captured_prepare_kwargs[0]["image"].dtype is torch.bfloat16
def test_ltx25_i2v_rejects_inputs_that_cannot_receive_crf18(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.model_version = "2.5"
with pytest.raises(ValueError, match="requires PIL images"):
pipe._prepare_video_latents_stage(
SimpleNamespace(),
SimpleNamespace(),
device=torch.device("cpu"),
noise_scale=0.0,
image=torch.zeros(3, 4, 4),
)
def test_ltx25_i2v_tensor_input_can_opt_out_of_crf18(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
captured_prepare_kwargs = []
def fake_prepare_latents(**kwargs):
captured_prepare_kwargs.append(kwargs)
return kwargs["image"], None
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.model_version = "2.5"
object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4)))
object.__setattr__(pipe, "prepare_latents", fake_prepare_latents)
request_inputs = SimpleNamespace(
height=4,
width=4,
num_frames=3,
num_videos_per_prompt=1,
generator=None,
latents=None,
image_crf=0,
)
prompt_context = SimpleNamespace(
batch_size=1,
positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16),
)
pipe._prepare_video_latents_stage(
request_inputs,
prompt_context,
device=torch.device("cpu"),
noise_scale=0.0,
image=torch.zeros(3, 4, 4),
)
assert captured_prepare_kwargs[0]["image"].shape == (1, 3, 4, 4)
def test_ltx25_i2v_crf_preserves_single_pixel_dimension(self, monkeypatch):
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf
image = torch.zeros(1, 4, 3, dtype=torch.uint8).numpy()
monkeypatch.setitem(sys.modules, "av", None)
assert _apply_image_conditioning_crf(image, 18) is image
def test_ltx25_i2v_reports_missing_pyav(self, monkeypatch):
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf
monkeypatch.setitem(sys.modules, "av", None)
with pytest.raises(ImportError, match="PyAV with a libx264 encoder"):
_apply_image_conditioning_crf(torch.zeros(4, 4, 3, dtype=torch.uint8).numpy(), 18)
def test_forward_resolves_request_image_and_delegates_to_shared_recipe_runtime(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -98,7 +229,7 @@ class TestLTXImageToVideoForwardStages:
with pytest.raises(ValueError, match="cannot mix text-to-video and image-to-video"):
pipe._resolve_request_image(mixed_req, None, request_inputs)
def test_denoise_timestep_kwargs_masks_video_only(self):
def test_denoise_timestep_kwargs_masks_video_and_keeps_audio_per_token(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
@@ -189,13 +320,18 @@ class TestLTXImageToVideoConditioning:
torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [20.0, 21.0], [30.0, 31.0]]]))
def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
@pytest.mark.parametrize(
("model_version", "sampled_shape"),
[("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))],
)
def test_i2v_5d_latents_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape):
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.model_version = model_version
pipe.vae_spatial_compression_ratio = 1
pipe.vae_temporal_compression_ratio = 1
pipe.transformer_spatial_patch_size = 1
@@ -237,14 +373,20 @@ class TestLTXImageToVideoConditioning:
torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [1.0, 1.0], [1.0, 1.0]]]))
assert sampled_shapes == [(1, 3, 2)]
assert sampled_shapes == [sampled_shape]
def test_ltx23_i2v_image_noise_is_sampled_after_packing(self, monkeypatch):
@pytest.mark.parametrize(
("model_version", "sampled_shape"),
[("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))],
)
def test_i2v_image_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape):
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.model_version = model_version
pipe.vae_spatial_compression_ratio = 1
pipe.vae_temporal_compression_ratio = 1
pipe.transformer_spatial_patch_size = 1
@@ -266,6 +408,7 @@ class TestLTXImageToVideoConditioning:
return torch.ones(shape, device=device, dtype=dtype)
monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor)
out, conditioning_mask = pipe.prepare_latents(
image=torch.zeros(1, 3, 1, 1),
@@ -279,10 +422,54 @@ class TestLTXImageToVideoConditioning:
device=torch.device("cpu"),
)
assert sampled_shapes == [(1, 3, 2)]
assert sampled_shapes == [sampled_shape]
torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]]))
def test_ltx25_re_noise_uses_official_scalar_noise_scale(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.model_version = "2.5"
pipe.vae_spatial_compression_ratio = 1
pipe.vae_temporal_compression_ratio = 1
pipe.transformer_spatial_patch_size = 1
pipe.transformer_temporal_patch_size = 1
pipe.vae = SimpleNamespace(
latents_mean=torch.zeros(2),
latents_std=torch.ones(2),
config=SimpleNamespace(scaling_factor=1.0),
)
object.__setattr__(
pipe,
"_encode_i2v_image_latents",
lambda *_args, **_kwargs: torch.zeros(1, 2, 1, 1, 1, dtype=torch.bfloat16),
)
captured = {}
def fake_create_noised_state(latents, noise_scale, generator=None):
del generator
captured["noise_scale"] = noise_scale
return latents
monkeypatch.setattr(ltx2_latents, "create_noised_state", fake_create_noised_state)
pipe.prepare_latents(
image=torch.zeros(1, 3, 1, 1),
batch_size=1,
num_channels_latents=2,
height=1,
width=1,
num_frames=3,
noise_scale=0.909375,
dtype=torch.bfloat16,
device=torch.device("cpu"),
latents=torch.zeros(1, 2, 3, 1, 1, dtype=torch.bfloat16),
)
assert captured["noise_scale"] == 0.909375
def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
@@ -385,6 +572,99 @@ class TestLTXImageToVideoConditioning:
torch.testing.assert_close(audio, torch.tensor([[[-2.0]]]))
assert pipeline.step_indices == [0, 1]
def test_ltx25_ancestral_adapter_matches_official_seeded_step_and_preserves_i2v_mask(self):
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
device = torch.device("cpu")
sigmas = torch.tensor([1.0, 0.5, 0.0])
pipeline = SimpleNamespace(device=device, scheduler=SimpleNamespace(sigmas=sigmas))
mask = torch.tensor([[1.0, 0.0]])
adapter = LTXVideoAudioStepAdapter(
pipeline,
SimpleNamespace(sigmas=sigmas.clone()),
latent_num_frames=1,
latent_height=1,
latent_width=1,
image_conditioned=True,
sampler="euler_ancestral",
generator=torch.Generator(device=device).manual_seed(42),
conditioning_mask=mask,
)
video = torch.tensor([[[3.0], [2.0]]])
audio = torch.tensor([[[4.0]]])
video_velocity = torch.tensor([[[7.0], [0.25]]])
audio_velocity = torch.tensor([[[0.5]]])
reference_generator = torch.Generator(device=device).manual_seed(10042)
video_noise = torch.randn(video.shape, generator=reference_generator)
audio_noise = torch.randn(audio.shape, generator=reference_generator)
def official_step(sample, velocity, noise):
sigma = sigmas[0]
sigma_next = sigmas[1]
denoised = sample.float() - velocity.float() * sigma
sigma_down = sigma_next * (sigma_next / sigma)
ratio = sigma_down / sigma
deterministic = ratio * sample.float() + (1.0 - ratio) * denoised
alpha_next = 1.0 - sigma_next
alpha_down = 1.0 - sigma_down
coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt()
return (alpha_next / alpha_down) * deterministic + noise * coeff
expected_video = official_step(video, video_velocity, video_noise)
expected_video[:, :1] = video[:, :1]
expected_audio = official_step(audio, audio_velocity, audio_noise)
((actual_video, actual_audio),) = adapter.step(
(video_velocity, audio_velocity),
(sigmas[0], sigmas[0]),
(video, audio),
)
torch.testing.assert_close(actual_video, expected_video)
torch.testing.assert_close(actual_audio, expected_audio)
terminal_video_velocity = torch.full_like(actual_video, 0.125)
terminal_audio_velocity = torch.full_like(actual_audio, 0.25)
((terminal_video, terminal_audio),) = adapter.step(
(terminal_video_velocity, terminal_audio_velocity),
(sigmas[1], sigmas[1]),
(actual_video, actual_audio),
)
expected_terminal_video = actual_video - terminal_video_velocity * sigmas[1]
expected_terminal_video[:, :1] = actual_video[:, :1]
expected_terminal_audio = actual_audio - terminal_audio_velocity * sigmas[1]
torch.testing.assert_close(terminal_video, expected_terminal_video)
torch.testing.assert_close(terminal_audio, expected_terminal_audio)
def test_ltx25_ancestral_quantizes_x0_to_model_dtype_before_fp32_step(self):
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import _ancestral_euler_step_from_velocity
sample = torch.tensor([[[3.140625, -1.234375]]], dtype=torch.bfloat16)
velocity = torch.tensor([[[0.333984375, 2.71875]]], dtype=torch.bfloat16)
sigmas = torch.tensor([0.73, 0.41, 0.0], dtype=torch.float32)
generator = torch.Generator(device="cpu").manual_seed(10042)
reference_generator = torch.Generator(device="cpu").manual_seed(10042)
sigma, sigma_next = sigmas[:2]
denoised = (sample.float() - velocity.float() * sigma).to(sample.dtype).float()
sigma_down = sigma_next * (sigma_next / sigma)
sigma_down_ratio = sigma_down / sigma
expected = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised
alpha_next = 1.0 - sigma_next
alpha_down = 1.0 - sigma_down
coefficient = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt()
noise = torch.randn(
sample.shape,
generator=reference_generator,
dtype=sample.dtype,
device=sample.device,
)
expected = ((alpha_next / alpha_down) * expected + noise.float() * coefficient).to(sample.dtype)
actual = _ancestral_euler_step_from_velocity(sample, velocity, sigmas, 0, generator)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def test_i2v_guidance_uses_zero_sigma_for_conditioned_tokens(self):
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
@@ -202,3 +202,68 @@ def test_prepare_audio_latents_rejects_incompatible_provided_length():
device=torch.device("cpu"),
latents=latents,
)
def test_create_noised_state_matches_official_fp32_lerp():
latents = torch.linspace(-2, 2, 4096, dtype=torch.bfloat16).reshape(1, 32, 128)
expected_generator = torch.Generator().manual_seed(42)
noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype)
expected = torch.lerp(latents.float(), noise.float(), torch.tensor(0.15, dtype=torch.bfloat16).float()).to(
latents.dtype
)
actual = latent_ops.create_noised_state(
latents,
torch.tensor(0.15, dtype=torch.bfloat16),
torch.Generator().manual_seed(42),
)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def test_create_conditioned_noised_state_matches_official_two_lerps():
latents = torch.linspace(-2, 2, 24, dtype=torch.bfloat16).reshape(1, 3, 8)
clean_latents = latents.clone()
clean_latents[:, 0] = 3.0
denoise_mask = torch.tensor([[[0.0], [1.0], [1.0]]])
expected_generator = torch.Generator().manual_seed(42)
noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype)
noised = torch.lerp(latents.float(), noise.float(), 0.909375)
expected = torch.lerp(clean_latents.float(), noised, denoise_mask).to(latents.dtype)
actual = latent_ops.create_conditioned_noised_state(
latents,
clean_latents,
denoise_mask,
0.909375,
torch.Generator().manual_seed(42),
)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def test_prepare_supplied_video_latents_uses_official_token_major_layout():
pipeline = _make_pipeline(LTX2Pipeline)
supplied = torch.arange(512, dtype=torch.float32).reshape(1, 32, 16)
actual_video = pipeline.prepare_latents(
batch_size=1,
num_channels_latents=4,
height=64,
width=64,
num_frames=9,
dtype=torch.float32,
device=torch.device("cpu"),
latents=supplied,
)
torch.testing.assert_close(actual_video, supplied, rtol=0, atol=0)
assert actual_video.stride()[1:] == (1, actual_video.shape[1])
def test_clear_audio_padding_keeps_padding_outside_sampler_state():
updated = torch.tensor([[[1.0, 2.0], [3.0, 4.0], [99.0, -99.0]]])
actual = latent_ops.clear_audio_padding(updated, 2)
torch.testing.assert_close(actual[:, :2], updated[:, :2])
torch.testing.assert_close(actual[:, 2:], torch.zeros_like(actual[:, 2:]))
@@ -28,6 +28,15 @@ def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker):
class TestCFGParallelHelpers:
"""Test LTX-2.3 CFG helper math without loading model weights."""
def test_single_guidance_pass_preserves_official_token_layout(self):
from vllm_omni.diffusion.models.ltx2.ltx2_guidance import _repeat_batch
tokens = torch.randn(1, 16, 32).transpose(1, 2).contiguous().transpose(1, 2)
repeated = _repeat_batch(tokens, 1)
assert repeated is tokens
assert repeated.stride() == tokens.stride()
def test_combine_cfg_noise_matches_x0_space_formula(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
@@ -20,6 +20,7 @@ from vllm_omni.diffusion.models.ltx2 import (
from vllm_omni.diffusion.models.ltx2.ltx2_components import (
LTX2_COMPONENT_PROFILE,
LTX2_DISTILLED_COMPONENT_PROFILE,
LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
LTX2_TWO_STAGE_COMPONENT_PROFILE,
LTX23_COMPONENT_PROFILE,
LTX23_DISTILLED_COMPONENT_PROFILE,
@@ -47,6 +48,7 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
)
from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
LTX2_DISTILLED_ONE_STAGE_RECIPE,
LTX2_DISTILLED_TWO_STAGE_RECIPE,
LTX2_ONE_STAGE_RECIPE,
LTX2_TWO_STAGE_RECIPE,
@@ -64,12 +66,14 @@ from vllm_omni.diffusion.models.ltx2.ltx2_request import (
)
from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2DistilledOneStagePipeline,
LTX2I2VDMD2Pipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
LTX2DistilledPipeline,
LTX2DistilledTwoStagePipeline,
LTX2TwoStagePipeline,
)
@@ -125,15 +129,21 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries():
assert issubclass(LTX2Pipeline, LTXRuntime)
assert issubclass(LTX2TwoStagePipeline, LTXRuntime)
assert issubclass(LTX2DistilledOneStagePipeline, LTXRuntime)
assert issubclass(LTX2DistilledTwoStagePipeline, LTXRuntime)
assert issubclass(LTX2DistilledPipeline, LTXRuntime)
assert LTX2DistilledPipeline is LTX2DistilledTwoStagePipeline
assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE
assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE
assert LTX2TwoStagePipeline.component_profile is LTX2_TWO_STAGE_COMPONENT_PROFILE
assert LTX2TwoStagePipeline.pipeline_recipe is LTX2_TWO_STAGE_RECIPE
assert LTX2DistilledOneStagePipeline.component_profile is LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
assert LTX2DistilledOneStagePipeline.pipeline_recipe is LTX2_DISTILLED_ONE_STAGE_RECIPE
assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE
for pipeline_cls, profile in (
(LTX2Pipeline, LTX2_COMPONENT_PROFILE),
(LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE),
(LTX2DistilledOneStagePipeline, LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE),
(LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE),
):
assert pipeline_cls._dit_modules == list(profile.dit_modules)
@@ -254,7 +264,7 @@ def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_pa
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
with pytest.raises(ValueError, match="only one-stage LTX recipes"):
with pytest.raises(ValueError, match="not qualified for this LTX recipe"):
LTX2TwoStagePipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
@@ -709,15 +719,48 @@ def test_ltx_component_loader_installs_omni_connector_attention(monkeypatch):
assert len(installed) == 2
assert all(processor.__class__.__name__ == "_LTXConnectorAttnProcessor" for processor in installed)
assert all(processor.has_learned_registers for processor in installed)
assert all(kernel["role"] == "ltx2.connector" for kernel in kernels)
assert all(kernel["role_category"] == "self" for kernel in kernels)
assert all(kernel["skip_sequence_parallel"] for kernel in kernels)
assert all(kernel["disable_kv_quant"] for kernel in kernels)
assert all(processor.has_learned_registers for processor in installed)
assert not any(processor.preserve_learned_register_mask for processor in installed)
@pytest.mark.parametrize("has_learned_registers", [False, True])
def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_registers):
def test_ltx25_connector_attention_preserves_learned_register_mask(monkeypatch):
installed = []
class OmniAttention:
def __init__(self, **_kwargs):
pass
class Attention:
heads = 32
head_dim = 128
inner_kv_dim = 4096
def set_processor(self, processor):
installed.append(processor)
monkeypatch.setattr(ltx2_components, "OmniAttention", OmniAttention)
connectors = SimpleNamespace(
video_connector=SimpleNamespace(
learnable_registers=object(),
transformer_blocks=[SimpleNamespace(attn1=Attention())],
),
audio_connector=SimpleNamespace(
learnable_registers=object(),
transformer_blocks=[SimpleNamespace(attn1=Attention())],
),
)
_install_connector_attention(connectors, preserve_learned_register_mask=True)
assert all(processor.has_learned_registers for processor in installed)
assert all(processor.preserve_learned_register_mask for processor in installed)
def test_ltx_connector_attention_dispatches_through_omni_kernel():
calls = []
class Backend:
@@ -732,12 +775,16 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
calls.append((query, key, value, metadata))
return query
class UpcastingNorm(torch.nn.Module):
def forward(self, tensor):
return tensor.float()
attention = SimpleNamespace(
heads=2,
head_dim=4,
inner_kv_dim=8,
norm_q=torch.nn.Identity(),
norm_k=torch.nn.Identity(),
norm_q=UpcastingNorm(),
norm_k=UpcastingNorm(),
to_q=torch.nn.Identity(),
to_k=torch.nn.Identity(),
to_v=torch.nn.Identity(),
@@ -746,24 +793,29 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
rope_type="split",
omni_attention=OmniAttention(),
)
hidden_states = torch.randn(2, 3, 8)
hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16)
additive_mask = torch.zeros(2, 1, 3, 3)
rotary = (
torch.randn(2, 2, 3, 2, dtype=torch.float32),
torch.randn(2, 2, 3, 2, dtype=torch.float32),
)
output = ltx2_components._LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers)(
output = ltx2_components._LTXConnectorAttnProcessor()(
attention,
hidden_states,
attention_mask=additive_mask,
query_rotary_emb=rotary,
)
query, key, value, metadata = calls[0]
assert query.shape == key.shape == value.shape == (2, 3, 2, 4)
if has_learned_registers:
assert metadata is None
else:
assert metadata.attn_mask.shape == (2, 3)
assert metadata.attn_mask.dtype == torch.bool
assert metadata.attn_mask.all()
torch.testing.assert_close(output, hidden_states)
assert query.dtype == key.dtype == value.dtype == torch.bfloat16
assert metadata.attn_mask.shape == (2, 3)
assert metadata.attn_mask.dtype == torch.bool
assert metadata.attn_mask.all()
expected_rotary = tuple(component.to(torch.bfloat16) for component in rotary)
expected = ltx2_components.apply_split_rotary_emb(hidden_states, expected_rotary, head_dim=4)
torch.testing.assert_close(output, expected, rtol=0, atol=0)
def test_ltx2_distilled_pipeline_shares_recipe_runtime():
@@ -1707,6 +1759,7 @@ class TestLTXForwardStages:
assert seen["request_inputs"].output_type == "latent"
assert seen["kwargs"] == {
"request_sigmas": request_sigmas,
"request_phase_sigmas": None,
"image": None,
}
@@ -1824,6 +1877,16 @@ class TestRegistryIntegration:
"pipeline_ltx2_two_stage",
"LTX2TwoStagePipeline",
)
assert _DIFFUSION_MODELS["LTX2DistilledOneStagePipeline"] == (
"ltx2",
"pipeline_ltx2",
"LTX2DistilledOneStagePipeline",
)
assert _DIFFUSION_MODELS["LTX2DistilledTwoStagePipeline"] == (
"ltx2",
"pipeline_ltx2_two_stage",
"LTX2DistilledTwoStagePipeline",
)
assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == (
"ltx2",
"pipeline_ltx2_two_stage",
@@ -1848,6 +1911,8 @@ class TestRegistryIntegration:
expected = [
"LTX2Pipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledOneStagePipeline",
"LTX2DistilledTwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
@@ -1861,6 +1926,8 @@ class TestRegistryIntegration:
[
"LTX2Pipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledOneStagePipeline",
"LTX2DistilledTwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
@@ -7,18 +7,54 @@ import pytest
import torch
from torch import nn
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
from vllm_omni.diffusion.models.ltx2 import ltx2_transformer
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import (
LTX2AudioVideoAttnProcessor,
LTX2FeedForward,
LTX2VideoTransformer3DModel,
_LTX2ParallelAttention,
LTX2VideoTransformerBlock,
_make_rms_norm,
apply_interleaved_rotary_emb,
apply_keyframes_absolute_embedding,
apply_split_rotary_emb,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def test_ltx_interleaved_rope_matches_official_model_dtype_arithmetic():
torch.manual_seed(0)
x = torch.randn(2, 3, 8, dtype=torch.bfloat16)
cos = torch.randn(2, 3, 8, dtype=x.dtype)
sin = torch.randn(2, 3, 8, dtype=x.dtype)
x_real, x_imag = x.unflatten(2, (4, 2)).unbind(-1)
rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
expected = x * cos + rotated * sin
actual = apply_interleaved_rotary_emb(x, (cos, sin))
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def test_ltx_split_rope_matches_official_model_dtype_arithmetic():
torch.manual_seed(1)
x = torch.randn(2, 3, 8, dtype=torch.bfloat16)
cos = torch.randn(2, 3, 4, dtype=x.dtype)
sin = torch.randn(2, 3, 4, dtype=x.dtype)
split_x = x.reshape(2, 3, 2, 4)
first_x, second_x = split_x.chunk(2, dim=-2)
cos_u = cos.unsqueeze(-2)
sin_u = sin.unsqueeze(-2)
expected = split_x * cos_u
expected[..., :1, :].addcmul_(-sin_u, second_x)
expected[..., 1:, :].addcmul_(sin_u, first_x)
expected = expected.reshape_as(x)
actual = apply_split_rotary_emb(x, (cos, sin), head_dim=8)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def test_ltx_rms_norm_no_affine_matches_official_torch_module():
norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)
@@ -36,11 +72,53 @@ def test_ltx_rms_norm_affine_weight_remains_parameter():
assert "weight" in norm.state_dict()
def test_ltx_transformer_has_separate_cfg_cache_dit_config():
def test_ltx_attention_restores_projection_dtype_after_upcasting_qk_norms():
calls = []
class UpcastingNorm(nn.Module):
def forward(self, tensor):
return tensor.float()
class RecordingAttention:
def __call__(self, query, key, value, metadata):
calls.append((query, key, value, metadata))
return query
attention = SimpleNamespace(
heads=2,
head_dim=4,
to_qkv=None,
to_q=nn.Identity(),
to_k=nn.Identity(),
to_v=nn.Identity(),
norm_q=UpcastingNorm(),
norm_k=UpcastingNorm(),
to_gate_logits=None,
attn=RecordingAttention(),
to_out=(nn.Identity(), nn.Identity()),
rope_type="split",
)
hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16)
rotary = (
torch.ones(2, 2, 3, 2, dtype=torch.float32),
torch.zeros(2, 2, 3, 2, dtype=torch.float32),
)
processor = ltx2_transformer.LTX2AudioVideoAttnProcessor()
processor._slice_rope_for_tp = lambda rope, _: rope
output = processor(attention, hidden_states, query_rotary_emb=rotary)
query, key, value, metadata = calls[0]
assert query.dtype == key.dtype == value.dtype == torch.bfloat16
assert metadata is None
torch.testing.assert_close(output, hidden_states)
def test_ltx_transformer_has_fused_cfg_cache_dit_config():
adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
assert isinstance(adapter_config, CacheDiTAdapterConfig)
assert adapter_config.has_separate_cfg
assert not adapter_config.has_separate_cfg
def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
@@ -76,8 +154,9 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):
monkeypatch.setattr(LTX2AudioVideoAttnProcessor, "_is_sp_enabled", staticmethod(lambda: True))
processor = LTX2AudioVideoAttnProcessor()
processor_cls = ltx2_transformer.LTX2AudioVideoAttnProcessor
monkeypatch.setattr(processor_cls, "_is_sp_enabled", staticmethod(lambda: True))
processor = processor_cls()
hidden_states = torch.zeros(1, 3, 4)
encoder_hidden_states = torch.zeros(1, 4, 4)
key_padding_mask = torch.tensor([[True, True, True, False]])
@@ -113,7 +192,7 @@ def _fake_ltx_parallel_attention():
return forward
layer = object.__new__(_LTX2ParallelAttention)
layer = object.__new__(ltx2_transformer._LTX2ParallelAttention)
nn.Module.__init__(layer)
layer.attn_backend = Backend
layer.attention = SimpleNamespace(forward=run("native"))
@@ -123,33 +202,281 @@ def _fake_ltx_parallel_attention():
@pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)])
def test_ltx_masked_cross_attention_uses_sdpa(query_length, key_length):
def test_ltx_masked_attention_uses_sdpa(query_length, key_length):
layer, calls = _fake_ltx_parallel_attention()
query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16)
key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16)
key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]])
metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
metadata = ltx2_transformer.AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
layer._run_local_attention(query, key, value, metadata)
assert calls == ["sdpa"]
def test_ltx_masked_self_attention_uses_sdpa():
layer, calls = _fake_ltx_parallel_attention()
query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
key_padding_mask = torch.tensor([[True, True, True, False]])
metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
layer._run_local_attention(query, key, value, metadata)
assert calls == ["sdpa"]
def test_ltx_unmasked_cross_attention_keeps_native_backend():
def test_ltx_unmasked_attention_keeps_native_backend():
layer, calls = _fake_ltx_parallel_attention()
query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
layer._run_local_attention(query, key, value, None)
assert calls == ["native"]
def test_ltx25_feed_forward_bias_flag_controls_both_linears(monkeypatch):
class FakeParallelLinear(nn.Linear):
def __init__(self, dim_in, dim_out, *, bias=True, **kwargs):
super().__init__(dim_in, dim_out, bias=bias)
monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeParallelLinear)
monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeParallelLinear)
without_bias = LTX2FeedForward(8, inner_dim=16, bias=False)
with_bias = LTX2FeedForward(8, inner_dim=16, bias=True)
assert without_bias.net[0].proj.bias is None
assert without_bias.net[2].bias is None
assert with_bias.net[0].proj.bias is not None
assert with_bias.net[2].bias is not None
def _build_tiny_ltx_transformer(monkeypatch, **overrides):
class FakeBlock(nn.Module):
def __init__(self, **kwargs):
super().__init__()
self.init_kwargs = kwargs
def forward(self, hidden_states, audio_hidden_states, *_args, **_kwargs):
return hidden_states, audio_hidden_states
monkeypatch.setattr(ltx2_transformer, "LTX2VideoTransformerBlock", FakeBlock)
kwargs = {
"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": 8,
"base_width": 8,
"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,
"use_prompt_embeddings": False,
"cross_attn_mod": True,
"audio_cross_attn_mod": True,
}
kwargs.update(overrides)
return LTX2VideoTransformer3DModel(**kwargs)
def test_ltx25_keyframe_absolute_embedding_matches_official():
hidden_states = torch.arange(32, dtype=torch.bfloat16).view(1, 4, 8)
keyframes_mask = torch.tensor([[[1.0], [1.0], [0.0], [0.0]]])
embedding = torch.linspace(-0.5, 0.5, 8, dtype=torch.bfloat16).view(1, 8)
actual = apply_keyframes_absolute_embedding(hidden_states, keyframes_mask, embedding)
expected = hidden_states + (keyframes_mask > 0).to(torch.bfloat16) * embedding
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
assert apply_keyframes_absolute_embedding(hidden_states, None, embedding) is hidden_states
def test_ltx25_constructor_flags_and_native_weight_loader(monkeypatch):
legacy_model = _build_tiny_ltx_transformer(monkeypatch)
assert legacy_model.config.ff_bias is True
assert legacy_model.config.audio_ff_bias is True
assert legacy_model.config.use_prompt_adaln_single is True
assert legacy_model.config.use_keyframes_abs_pos_embedding is False
assert hasattr(legacy_model, "prompt_adaln")
assert hasattr(legacy_model, "audio_prompt_adaln")
assert not hasattr(legacy_model, "keyframes_abs_pos_embedding")
assert legacy_model.transformer_blocks[0].init_kwargs["ff_bias"] is True
assert legacy_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True
ltx25_model = _build_tiny_ltx_transformer(
monkeypatch,
ff_bias=False,
audio_ff_bias=True,
use_prompt_adaln_single=False,
use_keyframes_abs_pos_embedding=True,
)
assert ltx25_model.keyframes_abs_pos_embedding.shape == (1, 8)
torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, torch.zeros(1, 8))
assert not hasattr(ltx25_model, "prompt_adaln")
assert not hasattr(ltx25_model, "audio_prompt_adaln")
assert ltx25_model.transformer_blocks[0].init_kwargs["ff_bias"] is False
assert ltx25_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True
monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 1)
expected = torch.arange(8, dtype=torch.float32).view(1, 8)
loaded = ltx25_model.load_weights([("keyframes_abs_pos_embedding", expected)])
assert loaded == {"keyframes_abs_pos_embedding"}
torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, expected)
def test_ltx_cross_attention_modulations_use_official_timesteps(monkeypatch):
class CaptureAda(nn.Module):
def __init__(self):
super().__init__()
self.calls = []
def forward(self, timestep, *, batch_size, hidden_dtype):
self.calls.append(timestep.detach().clone())
embedding = torch.zeros(timestep.numel(), 8, dtype=hidden_dtype)
return embedding, embedding
class CaptureRope(nn.Module):
def __init__(self):
super().__init__()
self.output_dtypes = []
def forward(self, _coords, *, device, out_dtype=None):
self.output_dtypes.append(out_dtype)
empty = torch.empty(0, device=device, dtype=out_dtype)
return empty, empty
model = _build_tiny_ltx_transformer(monkeypatch)
rope = CaptureRope()
model.rope = model.audio_rope = model.cross_attn_rope = model.cross_attn_audio_rope = rope
model.prompt_modulation = False
video_scale_shift = CaptureAda()
video_gate = CaptureAda()
audio_scale_shift = CaptureAda()
audio_gate = CaptureAda()
model.av_cross_attn_video_scale_shift = video_scale_shift
model.av_cross_attn_video_a2v_gate = video_gate
model.av_cross_attn_audio_scale_shift = audio_scale_shift
model.av_cross_attn_audio_v2a_gate = audio_gate
model(
hidden_states=torch.zeros(1, 2, 4),
audio_hidden_states=torch.zeros(1, 1, 4),
encoder_hidden_states=torch.zeros(1, 1, 8),
audio_encoder_hidden_states=torch.zeros(1, 1, 8),
timestep=torch.tensor([[0.0, 2.0]]),
audio_timestep=torch.tensor([[5.0]]),
sigma=torch.tensor([2.0]),
audio_sigma=torch.tensor([5.0]),
video_coords=torch.zeros(1, 3, 2, 2),
audio_coords=torch.zeros(1, 1, 1, 2),
)
torch.testing.assert_close(video_scale_shift.calls[0], torch.tensor([0.0, 2.0]))
torch.testing.assert_close(video_gate.calls[0], torch.tensor([5.0]))
torch.testing.assert_close(audio_scale_shift.calls[0], torch.tensor([5.0]))
torch.testing.assert_close(audio_gate.calls[0], torch.tensor([2.0]))
assert len(rope.output_dtypes) == 4
assert all(dtype is torch.float32 for dtype in rope.output_dtypes)
def test_ltx25_static_prompt_table_is_used_without_prompt_adaln():
class CaptureAttention(nn.Module):
def __init__(self):
super().__init__()
self.encoder_hidden_states = None
def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
if encoder_hidden_states is not None:
self.encoder_hidden_states = encoder_hidden_states.detach().clone()
return torch.zeros_like(hidden_states)
class ZeroOutput(nn.Module):
def forward(self, hidden_states):
return torch.zeros_like(hidden_states)
block = object.__new__(LTX2VideoTransformerBlock)
nn.Module.__init__(block)
block.video_cross_attn_adaln = True
block.audio_cross_attn_adaln = True
block.cross_attn_adaln = True
block.scale_shift_table = nn.Parameter(torch.zeros(9, 2))
block.audio_scale_shift_table = nn.Parameter(torch.zeros(9, 2))
block.prompt_scale_shift_table = nn.Parameter(torch.tensor([[1.0, 2.0], [3.0, 4.0]]))
block.audio_prompt_scale_shift_table = nn.Parameter(torch.tensor([[0.5, 1.0], [1.5, 2.0]]))
block.norm1 = nn.Identity()
block.audio_norm1 = nn.Identity()
block.norm2 = nn.Identity()
block.audio_norm2 = nn.Identity()
block.norm3 = nn.Identity()
block.audio_norm3 = nn.Identity()
block.attn1 = CaptureAttention()
block.audio_attn1 = CaptureAttention()
block.attn2 = CaptureAttention()
block.audio_attn2 = CaptureAttention()
block.ff = ZeroOutput()
block.audio_ff = ZeroOutput()
hidden_states = torch.ones(1, 1, 2)
audio_hidden_states = torch.ones(1, 1, 2)
encoder_hidden_states = torch.ones(1, 1, 2)
audio_encoder_hidden_states = torch.ones(1, 1, 2)
block(
hidden_states,
audio_hidden_states,
encoder_hidden_states,
audio_encoder_hidden_states,
temb=torch.zeros(1, 1, 18),
temb_audio=torch.zeros(1, 1, 18),
temb_ca_scale_shift=torch.empty(0),
temb_ca_audio_scale_shift=torch.empty(0),
temb_ca_gate=torch.empty(0),
temb_ca_audio_gate=torch.empty(0),
temb_prompt=None,
temb_prompt_audio=None,
use_a2v_cross_attention=False,
use_v2a_cross_attention=False,
)
torch.testing.assert_close(block.attn2.encoder_hidden_states, torch.tensor([[[5.0, 7.0]]]))
torch.testing.assert_close(block.audio_attn2.encoder_hidden_states, torch.tensor([[[3.0, 4.0]]]))
def test_ltx_gated_attention_uses_tp_sharded_projection(monkeypatch):
column_calls = []
class FakeColumnParallelLinear(nn.Linear):
def __init__(self, input_size, output_size, *, bias=True, **kwargs):
column_calls.append((output_size, kwargs))
super().__init__(input_size, output_size, bias=bias)
class FakeRowParallelLinear(nn.Linear):
def __init__(self, input_size, output_size, *, bias=True, **_kwargs):
super().__init__(input_size, output_size, bias=bias)
class FakeParallelAttention(nn.Module):
def __init__(self, **_kwargs):
super().__init__()
monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumnParallelLinear)
monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRowParallelLinear)
monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeParallelAttention)
monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2)
attention = ltx2_transformer.LTX2Attention(
query_dim=8,
heads=4,
kv_heads=4,
dim_head=2,
qk_norm="rms_norm",
apply_gated_attention=True,
pack_qkv=False,
)
gate_output_size, gate_kwargs = column_calls[-1]
assert gate_output_size == 4
assert gate_kwargs["gather_output"] is False
assert gate_kwargs["return_bias"] is False
assert attention.query_num_heads == 2
@@ -5,7 +5,7 @@ import json
import sys
from multiprocessing.reduction import ForkingPickler
from types import SimpleNamespace
from unittest.mock import Mock
from unittest.mock import Mock, patch
import numpy as np
import pytest
@@ -363,8 +363,18 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
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(tmp_path / partition_name) for partition_name in source_partitions),
str(component_path),
]
assert [source.subfolder for source in pipeline.weights_sources] == [
*("transformer" for _ in source_partitions),
"text_encoder",
]
expected_prefixes = ["transformer."]
if expected_dits == 2:
expected_prefixes.append("transformers_ref.")
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"
@@ -936,6 +946,66 @@ def test_text_encoder_stub_constructs_without_group_or_weights():
assert list(encoder.named_parameters()) == []
def test_global_quant_config_is_shared_by_dit_and_encoder():
from vllm_omni.quantization import build_quant_config, resolve_component_quant_config
global_config = build_quant_config("fp8")
assert resolve_component_quant_config(global_config, "transformer") is global_config
assert resolve_component_quant_config(global_config, "text_encoder") is global_config
def test_minimax_h3_fp8_component_selection():
from vllm_omni.quantization import build_quant_config, resolve_component_quant_config
dit_only = build_quant_config({"transformer": {"method": "fp8"}})
encoder_only = build_quant_config({"text_encoder": {"method": "fp8"}})
assert resolve_component_quant_config(dit_only, "transformer").get_name() == "fp8"
assert resolve_component_quant_config(dit_only, "text_encoder") is None
assert resolve_component_quant_config(encoder_only, "transformer") is None
assert resolve_component_quant_config(encoder_only, "text_encoder").get_name() == "fp8"
def test_text_encoder_rejects_serialized_fp8():
from vllm_omni.diffusion.models.minimax_h3.encoder import _validate_encoder_quant_config
config = SimpleNamespace(
get_name=lambda: "fp8",
is_checkpoint_fp8_serialized=True,
)
with pytest.raises(ValueError, match="online FP8 only"):
_validate_encoder_quant_config(config)
def test_text_encoder_linear_delegates_quantization_to_vllm_factory():
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
import vllm_omni.diffusion.models.minimax_h3.encoder as encoder_module
group = SimpleNamespace(rank_in_group=0, world_size=1)
method = UnquantizedLinearMethod()
quant_config = Mock()
quant_config.get_quant_method.return_value = method
prefix = "text_encoder.text_model.layers.0.mlp.gate_up_proj"
with (
patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
):
linear = encoder_module.MiniMaxH3Qwen3VLMergedColumnParallelLinear(
group,
input_size=4,
intermediate_size=8,
dtype=torch.bfloat16,
quant_config=quant_config,
prefix=prefix,
)
assert linear.quant_method is method
quant_config.get_quant_method.assert_called_once_with(linear, prefix=prefix)
def test_no_offload_keeps_text_encoder_resident():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
@@ -978,15 +1048,21 @@ def test_model_offload_uses_hooked_text_encoder_call():
pipeline.text_encoder.load_to_device.assert_not_called()
def test_layerwise_offload_releases_text_encoder():
@pytest.mark.parametrize(
"offload_flag",
["enable_layerwise_offload", "enable_distributed_layerwise_offload"],
)
def test_layerwise_offload_releases_text_encoder(offload_flag):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=True,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=False,
)
setattr(pipeline.od_config, offload_flag, True)
pipeline.text_encoder = Mock()
expected = torch.ones(2, 3)
pipeline.text_encoder.encode_ids.return_value = expected
@@ -1988,21 +2064,25 @@ def _one_layer_text_encoder():
group = _SingleRankEncoderGroup(0)
layer = nn.Module()
layer.self_attn = nn.Module()
layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear(
group,
hidden_size=_ENCODER_HIDDEN,
num_heads=_ENCODER_NUM_HEADS,
num_kv_heads=_ENCODER_NUM_KV_HEADS,
head_dim=_ENCODER_HEAD_DIM,
dtype=torch.float32,
)
layer.mlp = nn.Module()
layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear(
group,
input_size=_ENCODER_HIDDEN,
intermediate_size=_ENCODER_INTERMEDIATE,
dtype=torch.float32,
)
with (
patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
):
layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear(
group,
hidden_size=_ENCODER_HIDDEN,
num_heads=_ENCODER_NUM_HEADS,
num_kv_heads=_ENCODER_NUM_KV_HEADS,
head_dim=_ENCODER_HEAD_DIM,
dtype=torch.float32,
)
layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear(
group,
input_size=_ENCODER_HIDDEN,
intermediate_size=_ENCODER_INTERMEDIATE,
dtype=torch.float32,
)
layer.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(_ENCODER_HIDDEN)
encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
@@ -2035,6 +2115,7 @@ def _write_text_encoder_checkpoint(path, *, omit=()):
safetensors.torch.save_file(tensors, str(path / "model.safetensors"))
index = {"weight_map": dict.fromkeys(tensors, "model.safetensors")}
(path / "model.safetensors.index.json").write_text(json.dumps(index))
return tensors
@pytest.mark.parametrize(
@@ -2056,10 +2137,10 @@ def test_encoder_load_reports_missing_fused_source_shards(
tmp_path, omitted_sources, expected_detail, unreported_target
):
encoder = _one_layer_text_encoder()
_write_text_encoder_checkpoint(tmp_path, omit=omitted_sources)
weights = _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources)
with pytest.raises(RuntimeError) as excinfo:
encoder._load_weights(str(tmp_path))
encoder.load_weights(weights.items())
message = str(excinfo.value)
assert expected_detail in message
@@ -2068,17 +2149,17 @@ def test_encoder_load_reports_missing_fused_source_shards(
def test_encoder_load_rejects_unloaded_plain_param(tmp_path):
encoder = _one_layer_text_encoder()
_write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",))
weights = _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",))
with pytest.raises(RuntimeError, match=r"weights not loaded.*input_layernorm"):
encoder._load_weights(str(tmp_path))
encoder.load_weights(weights.items())
def test_encoder_load_places_every_source_in_its_own_rows(tmp_path):
encoder = _one_layer_text_encoder()
_write_text_encoder_checkpoint(tmp_path)
weights = _write_text_encoder_checkpoint(tmp_path)
encoder._load_weights(str(tmp_path))
encoder.load_weights(weights.items())
params = dict(encoder.named_parameters())
q_rows = _ENCODER_NUM_HEADS * _ENCODER_HEAD_DIM
@@ -36,7 +36,13 @@ def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp():
assert MiniMaxH3DiTModel._layerwise_offload_blocks_attrs == ["blocks"]
assert MiniMaxH3DiTModel._cache_dit_adapter_config.block_forward_patterns["blocks"] == ForwardPattern.Pattern_3
assert not MiniMaxH3DiTModel._cache_dit_adapter_config.has_separate_cfg
assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"}
assert set(MiniMaxH3DiTModel._sp_plan) == {
"sp_prepare",
"local_sp_prepare",
"sp_gather",
}
assert set(MiniMaxH3DiTModel._sp_plan["sp_prepare"]) == {0, 1, 2}
assert set(MiniMaxH3DiTModel._sp_plan["local_sp_prepare"]) == {2}
model = object.__new__(MiniMaxH3DiTModel)
nn.Module.__init__(model)
@@ -89,6 +95,279 @@ def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
torch.testing.assert_close(actual[..., 96:], x[..., 96:], atol=0, rtol=0)
def test_h3_rope_table_materializes_local_rows_in_fused_kernel_layout():
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
_build_rope_table,
)
local_freqs_half = torch.randn(2, 48)
local_freqs = torch.cat((local_freqs_half, local_freqs_half), dim=-1)
actual = _build_rope_table(local_freqs)
expected = torch.cat(
(torch.cos(local_freqs_half), torch.sin(local_freqs_half)),
dim=-1,
).to(torch.bfloat16)
assert actual.shape == (2, 96)
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
def test_strict_sp_local_span_uses_rank_owned_rows(monkeypatch):
from vllm_omni.diffusion import forward_context
from vllm_omni.diffusion.distributed import parallel_state
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True)
monkeypatch.setattr(forward_context, "get_ulysses_mode", lambda **kwargs: "strict")
monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 8)
monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 3)
monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 8)
monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
monkeypatch.setattr(parallel_state, "get_allgather_parallel_world_size", lambda: 1)
assert h3._sequence_parallel_local_span(16, hooks_applied=True) == (6, 2)
@pytest.mark.parametrize(
("hooks_applied", "mode", "seq_len", "ulysses", "ring", "allgather"),
[
(False, "strict", 16, 8, 1, 1),
(True, "advanced_uaa", 16, 8, 1, 1),
(True, "strict", 17, 8, 1, 1),
(True, "strict", 16, 4, 2, 1),
(True, "strict", 16, 1, 1, 8),
],
)
def test_local_span_falls_back_when_local_embedding_is_unsafe(
monkeypatch,
hooks_applied,
mode,
seq_len,
ulysses,
ring,
allgather,
):
from vllm_omni.diffusion import forward_context
from vllm_omni.diffusion.distributed import parallel_state
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True)
monkeypatch.setattr(forward_context, "get_ulysses_mode", lambda **kwargs: mode)
monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 8)
monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 3)
monkeypatch.setattr(
parallel_state,
"get_ulysses_parallel_world_size",
lambda: ulysses,
)
monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: ring)
monkeypatch.setattr(
parallel_state,
"get_allgather_parallel_world_size",
lambda: allgather,
)
assert h3._sequence_parallel_local_span(seq_len, hooks_applied=hooks_applied) == (
0,
seq_len,
)
def test_local_embedding_spans_reassemble_multimodal_rows():
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
class _IdentityProjection(nn.Module):
def forward(self, x):
return x, None
class _IdentityRefiner(nn.Module):
def forward(self, x, **kwargs):
return x
model = object.__new__(h3.MiniMaxH3DiTModel)
nn.Module.__init__(model)
model.hidden_size = 2
model.video_patch_proj = _IdentityProjection()
model.audio_patch_proj = _IdentityProjection()
model.condition_proj = _IdentityProjection()
model.token_refiner = _IdentityRefiner()
model.time_embedder = nn.Identity()
seq_len = 16
x = torch.arange(seq_len * 2, dtype=torch.float32).reshape(1, seq_len, 2)
audio_x = x + 100
text = torch.tensor(
[[200.0, 201.0], [202.0, 203.0], [204.0, 205.0], [206.0, 207.0]],
dtype=torch.bfloat16,
)
common = {
"x": x,
"audio_x": audio_x,
"text_embeddings_selected": text,
"unique_timesteps": torch.tensor([0.2, 0.3]),
"img_pos": torch.tensor([1, 2, 6, 7, 11, 12]),
"audio_pos": torch.tensor([3, 4, 8, 9, 13, 14]),
"text_pos": torch.tensor([0, 5, 10, 15]),
"refiner_cu_seqlens": torch.tensor([0, 4, 4], dtype=torch.int32),
"refiner_max_seqlen": 4,
"seq_len": seq_len,
"device": torch.device("cpu"),
"local_span": (0, seq_len),
}
full, full_t_emb = model._embed(**common)
local_spans = []
for rank in range(8):
local_span = (rank * (seq_len // 8), seq_len // 8)
local, local_t_emb = model._embed(
**{
**common,
"local_span": local_span,
}
)
assert local.shape == (2, 2)
torch.testing.assert_close(local_t_emb, full_t_emb)
local_spans.append(local)
torch.testing.assert_close(torch.cat(local_spans), full, atol=0, rtol=0)
@pytest.mark.parametrize(
"local_projection",
[False, True],
ids=["fallback", "strict-local"],
)
def test_sp_output_boundary_preserves_logits_and_minimizes_gather(
monkeypatch,
local_projection,
):
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
class _StaticRope(nn.Module):
input_rows = None
def forward(self, position_ids):
self.input_rows = position_ids.shape[1]
return torch.zeros(position_ids.shape[1], 6)
class _Prepare(nn.Module):
def __init__(self, shard):
super().__init__()
self.shard = shard
def forward(self, hidden, rope, combined):
local_len = hidden.shape[0] // 2 if self.shard else hidden.shape[0]
return hidden[:local_len], rope[:local_len], combined[:local_len]
class _Gather(nn.Module):
def __init__(self, remote_rows, replace):
super().__init__()
self.remote_rows = remote_rows
self.replace = replace
self.input_width = None
def forward(self, tensor):
self.input_width = tensor.shape[-1]
if self.replace:
return self.remote_rows
return torch.cat((tensor, self.remote_rows), dim=0)
class _FinalLayer(nn.Module):
def forward(self, hidden, *, t_emb, inverse_indices):
del t_emb
inverse = inverse_indices.to(torch.float32).unsqueeze(-1)
return (
hidden[:, :2].to(torch.float32) + inverse,
hidden[:, 2:3].to(torch.float32) - inverse,
)
seq_len = 8
full_hidden = torch.arange(seq_len * 4, dtype=torch.bfloat16).reshape(seq_len, 4)
inverse_indices = torch.tensor([0, 1, 0, 1, 1, 0, 1, 0])
final_layer = _FinalLayer()
remote_video, remote_audio = final_layer(
full_hidden[seq_len // 2 :],
t_emb=torch.empty(0),
inverse_indices=inverse_indices[seq_len // 2 :],
)
if local_projection:
remote_rows = torch.cat((remote_video, remote_audio), dim=-1)
else:
remote_rows = full_hidden
gather = _Gather(remote_rows, replace=not local_projection)
model = object.__new__(h3.MiniMaxH3DiTModel)
nn.Module.__init__(model)
model.arch = h3.MiniMaxH3DiTArchConfig(
latents_dim=2,
audio_latents_dim=1,
patch_size=(1, 1, 1),
)
rope = _StaticRope()
model.rope = rope
model.sp_prepare = _Prepare(shard=True)
model.local_sp_prepare = _Prepare(shard=False)
model.sp_gather = gather
model.blocks = nn.ModuleList()
model.final_layer = final_layer
monkeypatch.setattr(
model,
"_embed",
lambda **kwargs: (
full_hidden[: seq_len // 2] if local_projection else full_hidden,
torch.tensor([[0.5]], dtype=torch.float32),
),
)
def _local_span(*args, **kwargs):
del args, kwargs
return (0, seq_len // 2) if local_projection else (0, seq_len)
monkeypatch.setattr(
h3,
"_sequence_parallel_local_span",
_local_span,
)
img_pos = torch.tensor([0, 2, 4, 6])
audio_pos = torch.tensor([1, 3, 5, 7])
video, audio = model(
x=torch.zeros(1, seq_len, 2),
audio_x=torch.zeros(1, seq_len, 1),
img_position_ids=torch.zeros(1, seq_len, 3, dtype=torch.long),
unique_timesteps=torch.tensor([0.5]),
inverse_indices=inverse_indices,
update_mask=torch.ones(img_pos.numel()),
token_tags=torch.zeros(seq_len, dtype=torch.long),
prompt_embeds=torch.empty(0, 2),
img_pos_info={"position_ids": img_pos},
audio_pos_info={"position_ids": audio_pos},
text_pos_info={"position_ids": torch.empty(0, dtype=torch.long)},
img_pos_for_infer_output_info={"position_ids": img_pos},
packed_seq_params={
"cu_seqlens_q": torch.tensor([0, seq_len, seq_len], dtype=torch.int32),
"max_seqlen_q": seq_len,
},
refiner_packed_seq_params={
"cu_seqlens_q": torch.tensor([0, 0, 0], dtype=torch.int32),
"max_seqlen_q": 0,
},
)
expected_video, expected_audio = final_layer(
full_hidden,
t_emb=torch.empty(0),
inverse_indices=inverse_indices,
)
torch.testing.assert_close(video, expected_video.index_select(0, img_pos))
torch.testing.assert_close(audio, expected_audio.index_select(0, audio_pos))
expected_gather_width = 3 if local_projection else full_hidden.shape[-1]
expected_rope_rows = seq_len // 2 if local_projection else seq_len
assert gather.input_width == expected_gather_width
assert rope.input_rows == expected_rope_rows
@pytest.mark.parametrize(
("tp_size", "message"),
[
@@ -191,6 +191,49 @@ def test_model_load_weights_transforms_before_calling_vllm_loader():
]
def test_loader_adapter_declares_equivalent_direct_mmap_transform(monkeypatch):
from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
get_direct_mmap_adapter,
)
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear)
monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear)
monkeypatch.setattr(h3, "Attention", _FakeAttention)
arch = h3.MiniMaxH3DiTArchConfig(
hidden_size=1,
num_attention_heads=2,
attention_head_dim=1,
rope_inv_freq_len=1,
)
attention = h3.MiniMaxH3Attention(
arch,
quant_config=None,
prefix="blocks.0.attn",
)
transformer = object.__new__(h3.MiniMaxH3DiTModel)
nn.Module.__init__(transformer)
transformer.blocks = nn.ModuleList([nn.Module()])
transformer.blocks[0].attn = attention
pipeline = nn.Module()
pipeline.transformer = transformer
adapter = get_direct_mmap_adapter(pipeline)
assert adapter is not None
policy = adapter.policy_for(
"transformer.blocks.0.attn.qkv_proj.weight",
attention.qkv_proj.weight,
)
assert policy is not None
assert policy.allow_custom_loader
assert policy.transform is not None
checkpoint_weight = torch.arange(6, dtype=torch.float32).reshape(6, 1)
assert policy.transform(checkpoint_weight)[:, 0].tolist() == [0, 3, 1, 4, 2, 5]
assert not hasattr(attention.qkv_proj.weight, "mmap_weight_transform")
def test_pipeline_resolves_transformer_component_quant_config():
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_resolve_component_quant_config,
@@ -211,3 +254,18 @@ def test_pipeline_resolves_transformer_component_quant_config():
assert transformer_config.ignored_layers == ignored_layers
assert _resolve_component_quant_config(component_config, "transformer") is transformer_config
assert _resolve_component_quant_config(transformer_config, "transformer") is transformer_config
def test_pipeline_strips_prequantized_text_encoder_config():
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _resolve_minimax_h3_text_encoder_quant_config
from vllm_omni.quantization import ComponentQuantizationConfig
prequantized = Mock()
prequantized.get_name.return_value = "modelopt"
online_fp8 = Mock()
online_fp8.get_name.return_value = "fp8"
component_config = ComponentQuantizationConfig({"text_encoder": prequantized})
assert _resolve_minimax_h3_text_encoder_quant_config(online_fp8) is online_fp8
assert _resolve_minimax_h3_text_encoder_quant_config(prequantized) is None
assert _resolve_minimax_h3_text_encoder_quant_config(component_config) is None
@@ -142,12 +142,6 @@ def _generate_joint_output(omni, config: QualityTestConfig):
first = outputs[0]
worker_peak_memory_mb = float(getattr(first, "peak_memory_mb", 0.0) or 0.0)
if hasattr(first, "request_output") and isinstance(first.request_output, list):
first = first.request_output[0]
worker_peak_memory_mb = max(
worker_peak_memory_mb,
float(getattr(first, "peak_memory_mb", 0.0) or 0.0),
)
if not hasattr(first, "images") or not first.images:
raise ValueError("Could not extract video frames from H3 output")
frames = first.images[0]
@@ -8,15 +8,28 @@ import torch
from PIL import Image
from torch import nn
from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
Wan22I2VPipeline,
get_wan22_i2v_post_process_func,
get_wan22_i2v_pre_process_func,
)
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def test_wan22_i2v_postprocess_honors_request_output_type() -> None:
video = torch.zeros(1, 4, 1, 2, 2)
output = get_wan22_i2v_post_process_func(SimpleNamespace())(
video,
sampling_params=SimpleNamespace(output_type="latent"),
)
assert output is video
def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
pipeline = object.__new__(Wan22I2VPipeline)
nn.Module.__init__(pipeline)
@@ -31,6 +44,29 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
return pipeline
def _make_i2v_sampling(**overrides):
values = {
"height": 16,
"width": 16,
"num_frames": 5,
"num_inference_steps": 1,
"guidance_scale_provided": True,
"guidance_scale": 1.0,
"guidance_scale_2": None,
"guidance_scale_2_provided": False,
"boundary_ratio": None,
"generator": None,
"seed": None,
"num_outputs_per_prompt": 1,
"max_sequence_length": 8,
"latents": None,
"output_type": "latent",
"extra_args": {},
}
values.update(overrides)
return SimpleNamespace(**values)
def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
request = SimpleNamespace(
@@ -53,6 +89,88 @@ def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None:
preprocess(missing_image)
def _make_i2v_preprocess_request(last_image):
return SimpleNamespace(
prompt={
"prompt": "p",
"multi_modal_data": {
"image": Image.new("RGB", (320, 160), "red"),
"last_image": last_image,
},
},
sampling_params=SimpleNamespace(height=16, width=16),
)
def test_i2v_preprocess_treats_empty_last_image_list_as_absent() -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
result = preprocess(_make_i2v_preprocess_request([]))
assert result.prompt["multi_modal_data"]["last_image"] is None
assert result.batch_compatibility_key == ("wan22_i2v_last_image", False)
def test_i2v_preprocess_unwraps_single_last_image_list() -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
last_image = Image.new("RGB", (16, 16), "blue")
result = preprocess(_make_i2v_preprocess_request([last_image]))
assert result.prompt["multi_modal_data"]["last_image"] is last_image
assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
@pytest.mark.parametrize(
"last_image",
[
Image.new("RGB", (16, 16), "blue"),
torch.zeros(3, 16, 16),
],
)
def test_i2v_preprocess_preserves_supported_last_image(last_image) -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
result = preprocess(_make_i2v_preprocess_request(last_image))
assert result.prompt["multi_modal_data"]["last_image"] is last_image
assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
def test_i2v_preprocess_loads_last_image_path(tmp_path) -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
path = tmp_path / "last.png"
Image.new("RGB", (16, 16), "blue").save(path)
result = preprocess(_make_i2v_preprocess_request(str(path)))
last_image = result.prompt["multi_modal_data"]["last_image"]
assert isinstance(last_image, Image.Image)
assert last_image.mode == "RGB"
assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
def test_i2v_preprocess_rejects_multiple_last_images() -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
with pytest.raises(ValueError, match="at most one last_image"):
preprocess(
_make_i2v_preprocess_request(
[
Image.new("RGB", (16, 16), "blue"),
Image.new("RGB", (16, 16), "green"),
]
)
)
def test_i2v_preprocess_rejects_unsupported_last_image_type() -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
with pytest.raises(TypeError, match="Unsupported last_image format"):
preprocess(_make_i2v_preprocess_request({"not": "an image"}))
def test_i2v_diffuse_selects_stage_guidance_and_expands_timesteps() -> None:
pipeline = _make_i2v_pipeline(expand_timesteps=True)
latents = torch.zeros(1, 4, 2, 4, 4)
@@ -126,3 +244,169 @@ def test_i2v_prepare_latents_builds_expand_condition_and_first_frame_mask() -> N
assert first_frame_mask.shape == (1, 1, 2, 2, 2)
assert first_frame_mask[:, :, 0].sum() == 0
assert first_frame_mask[:, :, 1].sum() == 4
def test_i2v_prepare_latents_preserves_batched_image_conditions() -> None:
pipeline = _make_i2v_pipeline(expand_timesteps=True)
generators = [
torch.Generator(device="cpu").manual_seed(1),
torch.Generator(device="cpu").manual_seed(2),
]
latents, condition, first_frame_mask = pipeline.prepare_latents(
image=torch.zeros(2, 3, 16, 16),
batch_size=2,
num_channels_latents=4,
height=16,
width=16,
num_frames=5,
dtype=torch.float32,
device=torch.device("cpu"),
generator=generators,
)
assert latents.shape == (2, 4, 2, 2, 2)
assert condition.shape == (2, 4, 1, 2, 2)
assert first_frame_mask.shape == (2, 1, 2, 2, 2)
def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -> None:
monkeypatch.setattr(
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
SimpleNamespace(is_available=lambda: False),
)
pipeline = _make_i2v_pipeline(expand_timesteps=True)
pipeline.scheduler = StubScheduler([9])
pipeline.od_config = SimpleNamespace(flow_shift=5.0)
pipeline._sample_solver = "unipc"
pipeline._flow_shift = 5.0
pipeline.boundary_ratio = 0.875
pipeline.has_image_encoder = False
pipeline._guidance_scale = None
pipeline._guidance_scale_2 = None
pipeline._num_timesteps = None
pipeline._current_timestep = None
pipeline.check_inputs = lambda **kwargs: None
prepare_call = {}
def _fake_encode_prompt(**kwargs):
batch_size = len(kwargs["prompt"])
n = batch_size * kwargs["num_videos_per_prompt"]
return torch.zeros(n, 2, 3), None
def _fake_prepare_latents(**kwargs):
prepare_call.update(kwargs)
batch_size = kwargs["batch_size"]
return (
kwargs["latents"],
torch.zeros(batch_size, 4, 1, 2, 2),
torch.ones(batch_size, 1, 2, 2, 2),
)
pipeline.encode_prompt = _fake_encode_prompt # type: ignore[method-assign]
pipeline.prepare_latents = _fake_prepare_latents # type: ignore[method-assign]
pipeline.diffuse = lambda **kwargs: kwargs["latents"] # type: ignore[method-assign]
gen_a = torch.Generator(device="cpu").manual_seed(1)
gen_b = torch.Generator(device="cpu").manual_seed(2)
latents_a = torch.zeros(2, 4, 2, 2, 2)
latents_b = torch.ones(2, 4, 2, 2, 2)
image_a = torch.zeros(1, 3, 16, 16)
image_b = torch.ones(1, 3, 16, 16)
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={"prompt": "first", "multi_modal_data": {"image": image_a}},
sampling_params=_make_i2v_sampling(
generator=gen_a,
latents=latents_a,
num_outputs_per_prompt=2,
),
),
SimpleNamespace(
request_id="b",
prompt={"prompt": "second", "multi_modal_data": {"image": image_b}},
sampling_params=_make_i2v_sampling(
generator=gen_b,
latents=latents_b,
num_outputs_per_prompt=2,
),
),
]
)
outputs = pipeline.forward(batch)
assert prepare_call["batch_size"] == 4
assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
torch.testing.assert_close(
prepare_call["image"],
torch.cat([image_a, image_a, image_b, image_b]),
)
assert len(outputs) == 2
torch.testing.assert_close(outputs[0].output, latents_a)
torch.testing.assert_close(outputs[1].output, latents_b)
def test_i2v_forward_rejects_mismatched_tensor_condition_shapes(monkeypatch) -> None:
monkeypatch.setattr(
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
SimpleNamespace(is_available=lambda: False),
)
pipeline = _make_i2v_pipeline(expand_timesteps=True)
pipeline.scheduler = StubScheduler([9])
pipeline.od_config = SimpleNamespace(flow_shift=5.0)
pipeline._sample_solver = "unipc"
pipeline._flow_shift = 5.0
pipeline.boundary_ratio = 0.875
pipeline.has_image_encoder = False
pipeline._guidance_scale = None
pipeline._guidance_scale_2 = None
pipeline._num_timesteps = None
pipeline._current_timestep = None
pipeline.check_inputs = lambda **kwargs: None
pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign]
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={"prompt": "first", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 16)}},
sampling_params=_make_i2v_sampling(),
),
SimpleNamespace(
request_id="b",
prompt={"prompt": "second", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 8)}},
sampling_params=_make_i2v_sampling(),
),
]
)
with pytest.raises(ValueError, match="image condition"):
pipeline.forward(batch)
def test_i2v_forward_rejects_mixed_last_image_presence() -> None:
pipeline = _make_i2v_pipeline(expand_timesteps=True)
image = torch.zeros(1, 3, 16, 16)
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={
"prompt": "first",
"multi_modal_data": {"image": image, "last_image": image},
},
sampling_params=_make_i2v_sampling(),
),
SimpleNamespace(
request_id="b",
prompt={"prompt": "second", "multi_modal_data": {"image": image}},
sampling_params=_make_i2v_sampling(),
),
]
)
with pytest.raises(ValueError, match="mix of provided and missing last_image"):
pipeline.forward(batch)
@@ -93,6 +93,29 @@ def _make_pipeline() -> Wan22Pipeline:
return pipeline
def _make_sampling(**overrides):
values = {
"height": None,
"width": None,
"num_frames": 1,
"num_inference_steps": 2,
"guidance_scale_provided": True,
"guidance_scale": 1.0,
"guidance_scale_2": None,
"guidance_scale_2_provided": False,
"boundary_ratio": None,
"generator": None,
"seed": None,
"num_outputs_per_prompt": 1,
"max_sequence_length": 32,
"latents": None,
"output_type": "latent",
"extra_args": {},
}
values.update(overrides)
return SimpleNamespace(**values)
@pytest.mark.parametrize(
("sampling_params_kwargs", "expected_low", "expected_high"),
[
@@ -129,9 +152,10 @@ def test_forward_delegates_denoising_to_diffuse(
)
batch = DiffusionRequestBatch(requests=[mock_req])
output = pipeline.forward(batch)
outputs = pipeline.forward(batch)
assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8)))
assert len(outputs) == 1
assert torch.equal(outputs[0].output, torch.ones((1, 4, 1, 8, 8)))
assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8))
assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps)
assert captured["guidance_low"] == expected_low
@@ -142,6 +166,134 @@ def test_forward_delegates_denoising_to_diffuse(
assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))]
def test_forward_batches_text_generators_latents_and_splits_outputs() -> None:
pipeline = _make_pipeline()
encode_call = {}
prepare_call = {}
def _fake_encode_prompt(**kwargs):
encode_call.update(kwargs)
batch_size = len(kwargs["prompt"])
n = batch_size * kwargs["num_videos_per_prompt"]
return torch.arange(n, dtype=torch.float32).view(n, 1, 1), torch.zeros(n, 1, 1)
def _fake_prepare_latents(**kwargs):
prepare_call.update(kwargs)
return kwargs["latents"]
pipeline.encode_prompt = _fake_encode_prompt # type: ignore[method-assign]
pipeline.prepare_latents = _fake_prepare_latents # type: ignore[method-assign]
pipeline.diffuse = lambda **kwargs: kwargs["latents"] # type: ignore[method-assign]
gen_a = torch.Generator(device="cpu").manual_seed(1)
gen_b = torch.Generator(device="cpu").manual_seed(2)
latents_a = torch.zeros(2, 4, 1, 2, 2)
latents_b = torch.ones(2, 4, 1, 2, 2)
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={"prompt": "first", "negative_prompt": "bad first"},
sampling_params=_make_sampling(
generator=gen_a,
latents=latents_a,
num_outputs_per_prompt=2,
),
),
SimpleNamespace(
request_id="b",
prompt={"prompt": "second", "negative_prompt": "bad second"},
sampling_params=_make_sampling(
generator=gen_b,
latents=latents_b,
num_outputs_per_prompt=2,
),
),
]
)
outputs = pipeline.forward(batch)
assert encode_call["prompt"] == ["first", "second"]
assert encode_call["negative_prompt"] == ["bad first", "bad second"]
assert prepare_call["batch_size"] == 4
assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
assert len(outputs) == 2
torch.testing.assert_close(outputs[0].output, latents_a)
torch.testing.assert_close(outputs[1].output, latents_b)
def test_forward_batches_precomputed_prompt_embeddings() -> None:
pipeline = _make_pipeline()
diffuse_call = {}
pipeline.encode_prompt = lambda **kwargs: pytest.fail("text encoder must not run") # type: ignore[method-assign]
pipeline.prepare_latents = lambda **kwargs: torch.zeros(kwargs["batch_size"], 4, 1, 2, 2) # type: ignore[method-assign]
def _fake_diffuse(**kwargs):
diffuse_call.update(kwargs)
return kwargs["latents"]
pipeline.diffuse = _fake_diffuse # type: ignore[method-assign]
embeds_a = torch.zeros(3, 4)
embeds_b = torch.ones(3, 4)
negative_a = torch.full((3, 4), 2.0)
negative_b = torch.full((3, 4), 3.0)
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={"prompt_embeds": embeds_a, "negative_prompt_embeds": negative_a},
sampling_params=_make_sampling(guidance_scale=4.0),
),
SimpleNamespace(
request_id="b",
prompt={"prompt_embeds": embeds_b, "negative_prompt_embeds": negative_b},
sampling_params=_make_sampling(guidance_scale=4.0),
),
]
)
outputs = pipeline.forward(batch)
torch.testing.assert_close(diffuse_call["prompt_embeds"], torch.stack([embeds_a, embeds_b]))
torch.testing.assert_close(diffuse_call["negative_prompt_embeds"], torch.stack([negative_a, negative_b]))
assert len(outputs) == 2
def test_prepare_latents_with_request_generators_matches_single_generation() -> None:
pipeline = _make_pipeline()
kwargs = {
"num_channels_latents": 4,
"height": 16,
"width": 16,
"num_frames": 5,
"dtype": torch.float32,
"device": torch.device("cpu"),
}
batched = Wan22Pipeline.prepare_latents(
pipeline,
batch_size=2,
generator=[torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)],
**kwargs,
)
singles = torch.cat(
[
Wan22Pipeline.prepare_latents(
pipeline,
batch_size=1,
generator=torch.Generator().manual_seed(seed),
**kwargs,
)
for seed in (1, 2)
]
)
torch.testing.assert_close(batched, singles)
assert not torch.equal(batched[0], batched[1])
def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
pipeline = _make_pipeline()
latents = torch.zeros((1, 1, 1, 2, 2), dtype=torch.float32)
@@ -10,6 +10,7 @@ import torch
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
create_transformer_from_config,
get_wan22_post_process_func,
load_transformer_config,
retrieve_latents,
)
@@ -26,6 +27,17 @@ class _LatentDist:
return torch.tensor([2.0])
def test_wan22_postprocess_honors_request_output_type() -> None:
video = torch.zeros(1, 4, 1, 2, 2)
output = get_wan22_post_process_func(SimpleNamespace())(
video,
sampling_params=SimpleNamespace(output_type="latent"),
)
assert output is video
def test_retrieve_latents_supports_sample_mode_argmax_and_direct_latents() -> None:
generator = torch.Generator(device="cpu")
@@ -1,14 +1,307 @@
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import numpy as np
import PIL.Image
import pytest
import torch
import torch.nn as nn
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import (
Wan22S2VPipeline,
_make_clip_generators,
)
from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_s2v_sampling(**overrides):
values = {
"height": 16,
"width": 16,
"num_frames": 8,
"num_inference_steps": 1,
"guidance_scale_provided": True,
"guidance_scale": 1.0,
"generator": None,
"seed": None,
"num_outputs_per_prompt": 1,
"max_sequence_length": 8,
"latents": None,
"output_type": "latent",
"extra_args": {},
}
values.update(overrides)
return SimpleNamespace(**values)
def _make_s2v_validation_pipeline() -> Wan22S2VPipeline:
pipeline = object.__new__(Wan22S2VPipeline)
nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
pipeline.transformer = SimpleNamespace(dtype=torch.float32)
pipeline.vae_scale_factor_spatial = 8
pipeline.resolution_divisor = 16
pipeline.motion_frames = 7
pipeline.drop_first_motion = True
pipeline._DEFAULT_INFER_FRAMES = 8
pipeline._guidance_scale = None
pipeline.check_inputs = lambda *args, **kwargs: None
pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign]
return pipeline
def test_s2v_predict_noise_keeps_batch_dimension() -> None:
pipeline = object.__new__(Wan22S2VPipeline)
nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
transformer = MagicMock()
transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)])
transformer.return_value = (torch.zeros(1, 16, 2, 2, 2),)
pipeline.transformer = transformer
result = pipeline.predict_noise(hidden_states=torch.zeros(1, 16, 2, 2, 2))
assert result.shape == (1, 16, 2, 2, 2)
def test_s2v_clip_generators_preserve_main_seed_per_clip_behavior() -> None:
explicit_generator = torch.Generator(device="cpu").manual_seed(999)
generators = _make_clip_generators(
seeds=[1234, 5678, None, None],
request_generators=[
torch.Generator(device="cpu").manual_seed(1234),
torch.Generator(device="cpu").manual_seed(5678),
explicit_generator,
None,
],
clip_index=2,
device=torch.device("cpu"),
)
assert [generator.initial_seed() for generator in generators if generator is not None] == [1236, 5680, 999, 2]
assert generators[2] is explicit_generator
def test_s2v_forward_rejects_different_num_repeat() -> None:
pipeline = _make_s2v_validation_pipeline()
pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 2, 16) # type: ignore[method-assign]
image = PIL.Image.new("RGB", (16, 16))
audio = np.zeros(16000, dtype=np.float32)
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={
"prompt": "first",
"multi_modal_data": {"image": image, "audio": audio},
"additional_information": {"num_repeat": 1},
},
sampling_params=_make_s2v_sampling(),
),
SimpleNamespace(
request_id="b",
prompt={
"prompt": "second",
"multi_modal_data": {"image": image, "audio": audio},
"additional_information": {"num_repeat": 2},
},
sampling_params=_make_s2v_sampling(),
),
]
)
with pytest.raises(ValueError, match="num_repeat"):
pipeline.forward(batch)
def test_s2v_forward_rejects_different_audio_time_lengths() -> None:
pipeline = _make_s2v_validation_pipeline()
pipeline.encode_audio = lambda audio, **kwargs: ( # type: ignore[method-assign]
torch.zeros(1, 1, 2, 16 if audio[0] == 0 else 24),
1,
16 if audio[0] == 0 else 24,
)
image = PIL.Image.new("RGB", (16, 16))
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={
"prompt": "first",
"multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)},
},
sampling_params=_make_s2v_sampling(),
),
SimpleNamespace(
request_id="b",
prompt={
"prompt": "second",
"multi_modal_data": {"image": image, "audio": np.ones(16000, dtype=np.float32)},
},
sampling_params=_make_s2v_sampling(),
),
]
)
with pytest.raises(ValueError, match="audio embedding time lengths"):
pipeline.forward(batch)
def test_s2v_forward_rejects_different_init_first_frame() -> None:
pipeline = _make_s2v_validation_pipeline()
image = PIL.Image.new("RGB", (16, 16))
audio = np.zeros(16000, dtype=np.float32)
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={
"prompt": "first",
"multi_modal_data": {"image": image, "audio": audio},
"additional_information": {"init_first_frame": True},
},
sampling_params=_make_s2v_sampling(),
),
SimpleNamespace(
request_id="b",
prompt={
"prompt": "second",
"multi_modal_data": {"image": image, "audio": audio},
"additional_information": {"init_first_frame": False},
},
sampling_params=_make_s2v_sampling(),
),
]
)
with pytest.raises(ValueError, match="init_first_frame"):
pipeline.forward(batch)
def test_s2v_forward_rejects_different_raw_audio_shapes() -> None:
pipeline = _make_s2v_validation_pipeline()
pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 1, 16) # type: ignore[method-assign]
image = PIL.Image.new("RGB", (16, 16))
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={
"prompt": "first",
"multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)},
},
sampling_params=_make_s2v_sampling(),
),
SimpleNamespace(
request_id="b",
prompt={
"prompt": "second",
"multi_modal_data": {"image": image, "audio": np.zeros(8000, dtype=np.float32)},
},
sampling_params=_make_s2v_sampling(),
),
]
)
with pytest.raises(ValueError, match="raw audio shapes and sample rates"):
pipeline.forward(batch)
def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
pipeline = object.__new__(Wan22S2VPipeline)
nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
pipeline.transformer = MagicMock()
pipeline.transformer.dtype = torch.float32
pipeline.transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.float32)])
pipeline.transformer.casual_audio_encoder = None
pipeline.transformer.encode_audio.side_effect = lambda audio, _motion: {"audio_emb": audio}
pipeline.vae = MagicMock()
pipeline.vae.dtype = torch.float32
pipeline.vae.decode.return_value = (torch.zeros(4, 3, 8, 16, 16),)
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
parallel_config=SimpleNamespace(use_hsdp=False),
)
pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0]))
pipeline.vae_scale_factor_spatial = 8
pipeline.resolution_divisor = 16
pipeline.motion_frames = 7
pipeline.drop_first_motion = True
pipeline._DEFAULT_INFER_FRAMES = 8
pipeline._guidance_scale = None
pipeline._num_timesteps = None
pipeline.check_inputs = lambda *args, **kwargs: None
pipeline.encode_prompt = MagicMock(return_value=(torch.zeros(4, 2, 3), torch.ones(4, 2, 3)))
pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 1, 8))
pipeline.encode_ref_image = MagicMock(return_value=torch.zeros(1, 16, 1, 2, 2))
pipeline.prepare_motion_latents = MagicMock(
side_effect=lambda pixels, **_: torch.zeros(pixels.shape[0], 16, 2, 2, 2)
)
pipeline._denormalize_latents = lambda latents: latents
pipeline.diffuse = MagicMock(side_effect=lambda **kwargs: kwargs["latents"])
image = PIL.Image.new("RGB", (16, 16))
audio_a = np.zeros(16000, dtype=np.float32)
audio_b = np.ones(16000, dtype=np.float32)
latents_a = torch.zeros(2, 16, 2, 2, 2)
latents_b = torch.ones(2, 16, 2, 2, 2)
generators_a = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)]
generators_b = [torch.Generator().manual_seed(3), torch.Generator().manual_seed(4)]
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={
"prompt": "first",
"negative_prompt": "negative first",
"multi_modal_data": {"image": image, "audio": audio_a},
},
sampling_params=_make_s2v_sampling(
num_outputs_per_prompt=2,
generator=generators_a,
latents=latents_a,
),
),
SimpleNamespace(
request_id="b",
prompt={
"prompt": "second",
"negative_prompt": "negative second",
"multi_modal_data": {"image": image, "audio": audio_b},
},
sampling_params=_make_s2v_sampling(
num_outputs_per_prompt=2,
generator=generators_b,
latents=latents_b,
),
),
]
)
with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
platform.is_available.return_value = False
outputs = pipeline.forward(batch)
assert len(outputs) == 2
assert outputs[0].output[0].shape[0] == 2
assert outputs[1].output[0].shape[0] == 2
assert outputs[0].output[1].shape == (16000,)
assert outputs[1].output[1].shape == (16000,)
np.testing.assert_array_equal(outputs[0].output[1], audio_a)
np.testing.assert_array_equal(outputs[1].output[1], audio_b)
assert np.shares_memory(outputs[0].output[1], audio_a)
assert np.shares_memory(outputs[1].output[1], audio_b)
torch.testing.assert_close(pipeline.diffuse.call_args.kwargs["latents"], torch.cat([latents_a, latents_b]))
assert pipeline.diffuse.call_args.kwargs["clip_generator"] == generators_a + generators_b
assert pipeline.encode_prompt.call_args.kwargs["prompt"] == ["first", "second"]
assert pipeline.encode_prompt.call_args.kwargs["negative_prompt"] == ["negative first", "negative second"]
def test_s2v_exposes_hsdp_shard_conditions_for_transformer_blocks():
model = object.__new__(WanS2VTransformer3DModel)
nn.Module.__init__(model)
@@ -260,31 +553,37 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
pipeline._prompt_clean = Wan22S2VPipeline._prompt_clean
# -- Build request --
from vllm_omni.diffusion.request import OmniDiffusionRequest
sampling_params = MagicMock()
sampling_params.height = 704
sampling_params.width = 1024
sampling_params.num_frames = 80
sampling_params.num_inference_steps = 5
sampling_params.guidance_scale = 4.5
sampling_params.guidance_scale_provided = True
sampling_params.generator = None
sampling_params.seed = 42
sampling_params = _make_s2v_sampling(
height=704,
width=1024,
num_frames=80,
num_inference_steps=5,
guidance_scale=4.5,
generator=torch.Generator(device="cpu").manual_seed(42),
max_sequence_length=512,
)
ref_image = PIL.Image.new("RGB", (1024, 704))
audio_data = np.zeros(16000, dtype=np.float32)
req = MagicMock(spec=OmniDiffusionRequest)
req.prompts = [
{
"prompt": "test prompt",
"negative_prompt": "bad quality",
"multi_modal_data": {"image": ref_image, "audio": audio_data},
"additional_information": {"audio_path": audio_data, "pose_video": None, "init_first_frame": False},
}
]
req.sampling_params = sampling_params
req = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="test",
prompt={
"prompt": "test prompt",
"negative_prompt": "bad quality",
"multi_modal_data": {"image": ref_image, "audio": audio_data},
"additional_information": {
"audio_path": audio_data,
"pose_video": None,
"init_first_frame": False,
},
},
sampling_params=sampling_params,
)
]
)
# -- Mock methods that use complex internal state --
pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 25, 64, 80), 1, 80))
@@ -295,7 +594,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
pipeline.progress_bar.return_value.__exit__ = MagicMock(return_value=False)
# Mock predict_noise_maybe_with_cfg from CFGParallelMixin
pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(16, 20, 88, 128))
pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(1, 16, 20, 88, 128))
pipeline.predict_noise = Wan22S2VPipeline.predict_noise.__get__(pipeline)
# Mock platform methods
@@ -306,7 +605,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
with patch(
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.load_audio", return_value=(audio_data, 16000)
):
result = Wan22S2VPipeline.forward(pipeline, req=req)
result = Wan22S2VPipeline.forward(pipeline, req=req)[0]
# -- Assertions --
# Text encoder was called
@@ -8,13 +8,14 @@ import torch
from PIL import Image
from torch import nn
from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
Wan22VACEPipeline,
create_vace_transformer_from_config,
get_wan22_vace_pre_process_func,
)
from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -32,6 +33,27 @@ def _make_vace_pipeline() -> Wan22VACEPipeline:
return pipeline
def _make_vace_sampling(**overrides):
values = {
"height": 16,
"width": 16,
"num_frames": 5,
"num_inference_steps": 1,
"guidance_scale_provided": True,
"guidance_scale": 1.0,
"boundary_ratio": None,
"generator": None,
"seed": None,
"num_outputs_per_prompt": 1,
"max_sequence_length": 8,
"latents": None,
"output_type": "latent",
"extra_args": {},
}
values.update(overrides)
return SimpleNamespace(**values)
def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None:
preprocess = get_wan22_vace_pre_process_func(SimpleNamespace())
ref = Image.new("RGB", (320, 160), "green")
@@ -138,3 +160,189 @@ def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None:
assert calls[0]["positive_kwargs"]["vace_context"] is vace_context
assert calls[0]["negative_kwargs"]["vace_context_scale"] == 0.75
torch.testing.assert_close(result, torch.ones_like(latents))
def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> None:
monkeypatch.setattr(
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
SimpleNamespace(is_available=lambda: False),
)
pipeline = _make_vace_pipeline()
pipeline.transformer.vace_patch_embedding = None
pipeline.transformer_2 = None
pipeline.scheduler = StubScheduler([9])
pipeline.od_config = SimpleNamespace(flow_shift=3.0)
pipeline._sample_solver = "unipc"
pipeline._flow_shift = 3.0
pipeline.boundary_ratio = None
pipeline._guidance_scale = None
pipeline._num_timesteps = None
pipeline._current_timestep = None
pipeline.check_inputs = lambda **kwargs: None
prepare_call = {}
def _fake_encode_prompt(**kwargs):
batch_size = len(kwargs["prompt"])
n = batch_size * kwargs["num_videos_per_prompt"]
return torch.zeros(n, 2, 3), None
def _fake_prepare_latents(**kwargs):
prepare_call.update(kwargs)
return kwargs["latents"]
pipeline.encode_prompt = _fake_encode_prompt # type: ignore[method-assign]
pipeline.prepare_latents = _fake_prepare_latents # type: ignore[method-assign]
pipeline.diffuse = lambda **kwargs: kwargs["latents"] # type: ignore[method-assign]
gen_a = torch.Generator(device="cpu").manual_seed(1)
gen_b = torch.Generator(device="cpu").manual_seed(2)
latents_a = torch.zeros(2, 4, 2, 2, 2)
latents_b = torch.ones(2, 4, 2, 2, 2)
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={"prompt": "first"},
sampling_params=_make_vace_sampling(
generator=gen_a,
latents=latents_a,
num_outputs_per_prompt=2,
),
),
SimpleNamespace(
request_id="b",
prompt={"prompt": "second"},
sampling_params=_make_vace_sampling(
generator=gen_b,
latents=latents_b,
num_outputs_per_prompt=2,
),
),
]
)
outputs = pipeline.forward(batch)
assert prepare_call["batch_size"] == 4
assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
assert len(outputs) == 2
torch.testing.assert_close(outputs[0].output, latents_a)
torch.testing.assert_close(outputs[1].output, latents_b)
def test_vace_forward_rejects_mismatched_condition_shapes(monkeypatch) -> None:
monkeypatch.setattr(
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
SimpleNamespace(is_available=lambda: False),
)
pipeline = _make_vace_pipeline()
pipeline.transformer.vace_patch_embedding = object()
pipeline.transformer_2 = None
pipeline.scheduler = StubScheduler([9])
pipeline.od_config = SimpleNamespace(flow_shift=3.0)
pipeline._sample_solver = "unipc"
pipeline._flow_shift = 3.0
pipeline.boundary_ratio = None
pipeline._guidance_scale = None
pipeline._num_timesteps = None
pipeline._current_timestep = None
pipeline.check_inputs = lambda **kwargs: None
pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign]
def _fake_preprocess(video, **kwargs):
del kwargs
mask = torch.ones_like(video)
return video, mask, [[]]
pipeline.preprocess_conditions = _fake_preprocess # type: ignore[method-assign]
pipeline.prepare_video_latents = ( # type: ignore[method-assign]
lambda video, mask, refs, generator, device: torch.zeros(1, 8, video.shape[2], 2, 2)
)
pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, mask.shape[2], 2, 2) # type: ignore[method-assign]
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={
"prompt": "first",
"additional_information": {"source_video": torch.zeros(1, 3, 5, 16, 16)},
},
sampling_params=_make_vace_sampling(),
),
SimpleNamespace(
request_id="b",
prompt={
"prompt": "second",
"additional_information": {"source_video": torch.zeros(1, 3, 3, 16, 16)},
},
sampling_params=_make_vace_sampling(),
),
]
)
with pytest.raises(ValueError, match="VACE condition"):
pipeline.forward(batch)
def test_vace_forward_rejects_mismatched_reference_image_counts(monkeypatch) -> None:
monkeypatch.setattr(
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
SimpleNamespace(is_available=lambda: False),
)
pipeline = _make_vace_pipeline()
pipeline.transformer.vace_patch_embedding = object()
pipeline.transformer_2 = None
pipeline.scheduler = StubScheduler([9])
pipeline.od_config = SimpleNamespace(flow_shift=3.0)
pipeline._sample_solver = "unipc"
pipeline._flow_shift = 3.0
pipeline.boundary_ratio = None
pipeline._guidance_scale = None
pipeline._num_timesteps = None
pipeline._current_timestep = None
pipeline.check_inputs = lambda **kwargs: None
pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign]
def _fake_preprocess(video, mask, reference_images, **kwargs):
del kwargs
return video, mask, [reference_images]
pipeline.preprocess_conditions = _fake_preprocess # type: ignore[method-assign]
pipeline.prepare_video_latents = ( # type: ignore[method-assign]
lambda video, mask, refs, generator, device: torch.zeros(1, 8, 5, 2, 2)
)
pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, 5, 2, 2) # type: ignore[method-assign]
video = torch.zeros(1, 3, 5, 16, 16)
mask = torch.ones_like(video)
reference = torch.zeros(3, 16, 16)
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={
"prompt": "first",
"additional_information": {
"source_video": video,
"mask": mask,
"reference_images": [reference],
},
},
sampling_params=_make_vace_sampling(),
),
SimpleNamespace(
request_id="b",
prompt={
"prompt": "second",
"additional_information": {
"source_video": video,
"mask": mask,
"reference_images": [reference, reference],
},
},
sampling_params=_make_vace_sampling(),
),
]
)
with pytest.raises(ValueError, match="same number of reference images"):
pipeline.forward(batch)
@@ -18,6 +18,11 @@ from torch import nn
from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
from vllm_omni.diffusion.model_loader.host_weight_plan import (
HostWeightPlan,
TensorBinding,
build_checkpoint_mmap_plan,
)
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
from vllm_omni.diffusion.offloader.block_discovery import (
get_blocks_attr_names,
@@ -30,7 +35,10 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
PinnedResidentLayerGroup,
)
from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan
from vllm_omni.diffusion.offloader.offload_plan import (
OffloadPlan,
get_offload_plan,
)
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -48,6 +56,9 @@ class DummyEvent:
def record(self, _stream) -> None:
return None
def synchronize(self) -> None:
return None
@contextmanager
def dummy_stream(_stream):
@@ -139,6 +150,41 @@ class TestDistributedLayerwiseOffloadHook:
assert shard.numel() == 4
assert torch.equal(shard, _make_values(10.0))
def test_prefetch_preserves_transposed_weight_stride(
self,
dist_group,
patched_offload_runtime,
):
"""Online-FP8 Cutlass weights must retain their transposed layout."""
class StridedBlock(nn.Module):
def __init__(self, start: float):
super().__init__()
base = torch.arange(start, start + 12).reshape(3, 4)
self.weight = nn.Parameter(base.t(), requires_grad=False)
current_block = StridedBlock(1.0)
next_block = StridedBlock(20.0)
expected = next_block.weight.detach().clone()
expected_stride = next_block.weight.stride()
assert expected_stride == (1, 4)
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
device=torch.device("cpu"),
dp_group=None,
dp_size=1,
rank=0,
copy_stream=DummyStream(),
comm_stream=DummyStream(),
pin_memory=False,
)
hook.initialize_hook(current_block)
hook.prefetch_layer(slot=0, non_blocking=False)
assert next_block.weight.stride() == expected_stride
assert torch.equal(next_block.weight, expected)
def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
"""DTensor wrapper should be preserved through prefetch/offload cycle."""
current_block = TinyBlock(_make_values(1.0))
@@ -316,8 +362,9 @@ class TestDistributedLayerwiseOffloadHook:
block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
next_block = nn.Module()
next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
next_block.weight.mmap_weight_transform_pending = True
def transform(tensor):
return tensor.flip(0)
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
@@ -326,17 +373,17 @@ class TestDistributedLayerwiseOffloadHook:
dp_size=2,
rank=0,
pin_memory=False,
tensor_transforms={id(next_block.weight): transform},
)
hook.initialize_hook(block)
assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([3.0, 2.0]))
def test_regular_loader_weight_does_not_reapply_mmap_transform(self, dist_group, patched_offload_runtime):
def test_transform_is_not_applied_without_a_loader_plan(self, dist_group, patched_offload_runtime):
block = nn.Module()
block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
next_block = nn.Module()
next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
@@ -350,6 +397,82 @@ class TestDistributedLayerwiseOffloadHook:
assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0]))
def test_rank_local_mmap_retains_source_and_uses_staging(self, dist_group, patched_offload_runtime):
current_block = nn.Linear(2, 2, bias=False)
next_block = nn.Linear(2, 2, bias=False)
next_block.weight.data.copy_(torch.arange(4, dtype=torch.float32).view(2, 2))
source_storage = next_block.weight.untyped_storage().data_ptr()
def transform(tensor):
return tensor.t()
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
device=torch.device("cpu"),
dp_group=None,
dp_size=1,
rank=0,
pin_memory=False,
rank_local_mmap=True,
tensor_transforms={id(next_block.weight): transform},
)
hook.initialize_hook(current_block)
source = hook.cpu_sources[torch.float32][0]["tensor"]
assert source.untyped_storage().data_ptr() == source_storage
assert hook.cpu_shards == {}
assert next_block.weight.numel() == 0
hook.cpu_staging_buffers = [
{torch.float32: torch.empty(4)},
{torch.float32: torch.empty(4)},
]
hook.prefetch_layer(slot=0, non_blocking=False)
assert torch.equal(
next_block.weight,
torch.tensor([[0.0, 2.0], [1.0, 3.0]]),
)
assert next_block.weight.stride() == (1, 2)
assert torch.equal(source, torch.arange(4, dtype=torch.float32).view(2, 2))
def test_rank_local_mmap_staging_is_bounded_by_largest_block(self, patched_offload_runtime):
hooks = []
for size in (4, 9):
current_block = nn.Linear(1, 1, bias=False)
next_block = nn.Module()
next_block.weight = nn.Parameter(torch.arange(size, dtype=torch.float32))
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
device=torch.device("cpu"),
dp_group=None,
dp_size=1,
rank=0,
pin_memory=False,
shared_buffers=[None, None],
rank_local_mmap=True,
)
hook.initialize_hook(current_block)
hooks.append(hook)
resident_block = nn.Module()
resident_block.weight = nn.Parameter(torch.arange(12, dtype=torch.float32))
resident_group = PinnedResidentLayerGroup(
[resident_block],
device=torch.device("cpu"),
copy_stream=DummyStream(),
pin_memory=False,
rank_local_mmap=True,
)
staging = DistributedLayerwiseOffloadBackend._allocate_shared_cpu_staging_buffers(
hooks,
resident_group,
)
assert len(staging) == 2
assert staging[0][torch.float32].numel() == 12
assert staging[1][torch.float32].numel() == 12
class TestPinnedResidentLayerGroup:
def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
@@ -386,23 +509,59 @@ class TestPinnedResidentLayerGroup:
group.offload()
def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self):
module = nn.Linear(2, 2, bias=False)
def test_load_preserves_noncontiguous_weight_stride(self, patched_offload_runtime):
block = nn.Module()
weight = torch.arange(6, dtype=torch.float32).reshape(2, 3).t()
expected = weight.detach().clone()
expected_stride = weight.stride()
block.weight = nn.Parameter(weight)
group = PinnedResidentLayerGroup(
[block],
device=torch.device("cpu"),
copy_stream=DummyStream(),
pin_memory=False,
)
def transform(tensor):
return tensor.flip(0)
group.load()
assert torch.equal(block.weight, expected)
assert block.weight.stride() == expected_stride
group.offload()
module.weight.mmap_weight_transform = transform
def test_rank_local_mmap_resident_layers_retain_sources(self, patched_offload_runtime):
blocks = [nn.Linear(2, 2, bias=False) for _ in range(3)]
expected = []
source_storage = []
for index, block in enumerate(blocks):
block.weight.data.fill_(index + 1)
expected.append(block.weight.detach().clone())
source_storage.append(block.weight.untyped_storage().data_ptr())
backend = object.__new__(DistributedLayerwiseOffloadBackend)
backend._remember_mmap_param_attrs(module)
module.to_empty(device="meta")
group = PinnedResidentLayerGroup(
blocks,
device=torch.device("cpu"),
copy_stream=DummyStream(),
pin_memory=False,
rank_local_mmap=True,
)
replacement = nn.Parameter(torch.arange(4).reshape(2, 2).float())
backend._attach_mmap_param_attrs("weight", replacement, module.weight)
for index, state in enumerate(group._states):
source = state["cpu_sources"][torch.float32][0]["tensor"]
assert source.untyped_storage().data_ptr() == source_storage[index]
assert state["cpu_shards"] == {}
assert len(group._cpu_staging_buffers) == 2
assert replacement.mmap_weight_transform is transform
assert replacement.mmap_weight_transform_pending is True
group.load()
for block, weight in zip(blocks, expected):
assert torch.equal(block.weight, weight)
blocks[0].weight.data.zero_()
group.offload()
group.load()
for block, weight in zip(blocks, expected):
assert torch.equal(block.weight, weight)
group.offload()
assert all(block.weight.numel() == 0 for block in blocks)
def test_shared_allgather_output_is_narrowed_to_current_block(
self, dist_group, patched_offload_runtime, monkeypatch
@@ -545,7 +704,6 @@ class TestMmapWeightLoading:
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
model_path=str(tmp_path),
),
torch.device("cpu"),
)
@@ -553,13 +711,81 @@ class TestMmapWeightLoading:
dits=[pipeline.transformer],
dit_names=["transformer"],
)
plan = HostWeightPlan(
backing_kind="checkpoint_mmap",
bindings={
name: TensorBinding(
checkpoint_key=name,
file_path=str(weight_file),
)
for name in weights
},
)
backend._load_weights_via_mmap(pipeline, modules)
backend._load_weights_via_mmap(pipeline, modules, plan)
assert pipeline.transformer.post_load_calls == 1
assert pipeline.transformer.time_embedder.weight.dtype == torch.float32
assert pipeline.transformer.blocks[0].weight.dtype == torch.bfloat16
def test_allgather_flag_with_group_size_one_uses_rank_local_mmap(
self,
tmp_path,
patched_offload_runtime,
):
class Transformer(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"]
def __init__(self):
super().__init__()
self.blocks = nn.ModuleList(
[
nn.Linear(2, 2, bias=False),
nn.Linear(2, 2, bias=False),
]
)
class Pipeline(nn.Module):
def __init__(self):
super().__init__()
self.transformer = Transformer()
@staticmethod
def _remap_ckpt_key(key):
return key
pipeline = Pipeline()
weights = {name: torch.ones_like(param) for name, param in pipeline.named_parameters()}
checkpoint_file = tmp_path / "model.safetensors"
save_file(weights, str(checkpoint_file))
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=(),
model_path=str(tmp_path),
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert result.plan is not None
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
dp_size=1,
dlo_use_allgather=True,
),
torch.device("cpu"),
host_weight_plan=result.plan,
)
backend.enable(pipeline)
assert backend._using_rank_local_mmap
assert backend.dp_group is None
assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group)
backend.disable()
class TestGetBlocksFromDit:
def test_get_blocks_from_dit_single_block_attr(self):
@@ -825,23 +1051,123 @@ class TestOffloadPlan:
class TestMmapValidation:
"""Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights."""
"""Tests for loader preflight and backend plan realization."""
def test_tp_rejected_when_tp_world_size_gt_1(self, tmp_path, patched_offload_runtime, monkeypatch):
"""DLO+AllGather should reject when TP world size > 1."""
import json
from types import SimpleNamespace
from safetensors.torch import save_file
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
# Mock TP world size = 2
monkeypatch.setattr(
"vllm.distributed.parallel_state.get_tensor_model_parallel_world_size",
lambda: 2,
def test_component_source_prefix_builds_model_namespace(self, tmp_path):
source_root = tmp_path / "partition"
checkpoint_dir = source_root / "transformer"
checkpoint_dir.mkdir(parents=True)
checkpoint_file = checkpoint_dir / "model.safetensors"
save_file(
{"blocks.0.weight": torch.arange(4, dtype=torch.float32).reshape(1, 4)},
str(checkpoint_file),
)
class Pipeline(nn.Module):
def __init__(self):
super().__init__()
self.transformer = nn.Module()
self.transformer.blocks = nn.ModuleList([nn.Linear(4, 1, bias=False)])
self.weights_sources = [
SimpleNamespace(
model_or_path=str(source_root),
subfolder="transformer",
revision=None,
prefix="transformer.",
)
]
pipeline = Pipeline()
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=tuple(pipeline.weights_sources),
model_path=None,
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert result.fallback_reason is None
assert result.plan is not None
assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding(
checkpoint_key="blocks.0.weight",
file_path=str(checkpoint_file),
)
assert result.plan.planned_source_prefixes == frozenset({"transformer."})
def test_non_dedicated_component_source_falls_back_before_discovery(self):
pipeline = nn.Module()
pipeline.transformer = nn.Linear(2, 2, bias=False)
sources = (
SimpleNamespace(
model_or_path="unused",
subfolder=None,
revision=None,
prefix="",
),
)
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=sources,
model_path=None,
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert result.plan is None
assert "dedicated component weight source" in result.fallback_reason
def test_preflight_falls_back_before_mutation_for_missing_model_tensors(self, tmp_path):
source_root = tmp_path / "partition"
checkpoint_dir = source_root / "transformer"
checkpoint_dir.mkdir(parents=True)
save_file(
{"blocks.0.weight": torch.ones((2, 2), dtype=torch.float32)},
str(checkpoint_dir / "model.safetensors"),
)
class Transformer(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"]
def __init__(self):
super().__init__()
self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False), nn.Linear(2, 2, bias=False)])
class Pipeline(nn.Module):
def __init__(self):
super().__init__()
self.transformer = Transformer()
self.weights_sources = [
SimpleNamespace(
model_or_path=str(source_root),
subfolder="transformer",
revision=None,
prefix="transformer.",
)
]
pipeline = Pipeline()
original_weight = pipeline.transformer.blocks[0].weight
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=tuple(pipeline.weights_sources),
model_path=None,
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert result.plan is None
assert "1 required DiT tensors" in result.fallback_reason
assert pipeline.transformer.blocks[0].weight is original_weight
assert not original_weight.is_meta
def test_tp_falls_back_before_checkpoint_discovery(self):
class SimpleModel(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"]
@@ -859,41 +1185,118 @@ class TestMmapValidation:
return key
pipeline = SimplePipeline()
weights = {name: torch.ones(p.shape, dtype=p.dtype) for name, p in pipeline.named_parameters() if not p.is_meta}
save_file(weights, str(tmp_path / "model.safetensors"))
(tmp_path / "model.safetensors.index.json").write_text(
json.dumps({"weight_map": {k: "model.safetensors" for k in weights}})
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=(),
model_path="unused",
tensor_parallel_size=2,
use_hsdp=False,
online_quantization=False,
)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
model_path=str(tmp_path),
),
torch.device("cpu"),
)
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
assert result.plan is None
assert result.fallback_reason == "TP=2 requires the ordinary loader"
with pytest.raises(ValueError, match="Tensor Parallel"):
backend._load_weights_via_mmap(pipeline, modules)
@pytest.mark.parametrize(
("use_hsdp", "online_quantization", "expected_reason"),
[
(True, False, "HSDP requires the ordinary loader"),
(False, True, "online quantization requires the ordinary loader"),
],
)
def test_runtime_layouts_that_need_loader_callbacks_fall_back(
self,
use_hsdp,
online_quantization,
expected_reason,
):
pipeline = nn.Module()
pipeline.transformer = nn.Linear(2, 2, bias=False)
def test_tp_allowed_when_tp_world_size_eq_1(self, tmp_path, patched_offload_runtime, monkeypatch):
"""DLO+AllGather should work when TP world size = 1, even if params
have custom weight_loader attributes (e.g. QKVParallelLinear at TP=1)."""
import json
from types import SimpleNamespace
from safetensors.torch import save_file
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
# Mock TP world size = 1 (default)
monkeypatch.setattr(
"vllm.distributed.parallel_state.get_tensor_model_parallel_world_size",
lambda: 1,
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=(),
model_path="unused",
tensor_parallel_size=1,
use_hsdp=use_hsdp,
online_quantization=online_quantization,
)
assert result.plan is None
assert result.fallback_reason == expected_reason
@pytest.mark.parametrize(
("checkpoint_tensor", "expected_reason"),
[
(torch.ones((3, 2), dtype=torch.float32), "shape mismatch"),
(torch.ones((2, 2), dtype=torch.bfloat16), "dtype mismatch"),
],
)
def test_source_metadata_mismatch_falls_back(self, tmp_path, checkpoint_tensor, expected_reason):
checkpoint_file = tmp_path / "model.safetensors"
save_file({"transformer.weight": checkpoint_tensor}, str(checkpoint_file))
class Pipeline(nn.Module):
def __init__(self):
super().__init__()
self.transformer = nn.Linear(2, 2, bias=False)
@staticmethod
def _remap_ckpt_key(key):
return key
pipeline = Pipeline()
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=(),
model_path=str(tmp_path),
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert result.plan is None
assert expected_reason in result.fallback_reason
def test_non_persistent_buffers_do_not_need_checkpoint_bindings(self, tmp_path):
class Transformer(nn.Module):
def __init__(self):
super().__init__()
self.weight = nn.Parameter(torch.ones(2, 2))
self.register_buffer("derived", torch.ones(2), persistent=False)
checkpoint_file = tmp_path / "model.safetensors"
save_file({"transformer.weight": torch.ones(2, 2)}, str(checkpoint_file))
class Pipeline(nn.Module):
def __init__(self):
super().__init__()
self.transformer = Transformer()
@staticmethod
def _remap_ckpt_key(key):
return key
pipeline = Pipeline()
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=(),
model_path=str(tmp_path),
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert result.plan is not None
assert set(result.plan.bindings) == {"transformer.weight"}
def test_unadapted_custom_weight_loader_falls_back_at_tp1(self, tmp_path):
def custom_weight_loader(param, weight):
param.data.copy_(weight)
@@ -923,27 +1326,67 @@ class TestMmapValidation:
json.dumps({"weight_map": {k: "model.safetensors" for k in weights}})
)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
model_path=str(tmp_path),
),
torch.device("cpu"),
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=(),
model_path=str(tmp_path),
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
# Should NOT raise — TP=1 is allowed even with custom weight_loader
backend._load_weights_via_mmap(pipeline, modules)
assert result.plan is None
assert "custom weight_loader" in result.fallback_reason
def test_cosmos3_adapter_preserves_existing_tp1_direct_mmap_contract(self, tmp_path):
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
Cosmos3OmniDiffusersPipeline,
)
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
Cosmos3VFMTransformer,
)
transformer = object.__new__(Cosmos3VFMTransformer)
nn.Module.__init__(transformer)
transformer.proj_in = nn.Linear(2, 2, bias=False)
transformer.proj_in.weight.weight_loader = lambda *_args: None
pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
nn.Module.__init__(pipeline)
pipeline.transformer = transformer
source_root = tmp_path / "partition"
checkpoint_dir = source_root / "transformer"
checkpoint_dir.mkdir(parents=True)
checkpoint_file = checkpoint_dir / "model.safetensors"
save_file({"proj_in.weight": torch.ones(2, 2)}, str(checkpoint_file))
sources = (
SimpleNamespace(
model_or_path=str(source_root),
subfolder="transformer",
revision=None,
prefix="transformer.",
),
)
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", transformer),),
sources=sources,
model_path=None,
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert result.plan is not None
assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding(
checkpoint_key="proj_in.weight",
file_path=str(checkpoint_file),
)
def test_validate_loaded_weights_called(self, tmp_path, patched_offload_runtime):
"""validate_loaded_weights should be called after mmap loading."""
import json
from types import SimpleNamespace
from safetensors.torch import save_file
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
class ModelWithValidate(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"]
@@ -985,19 +1428,64 @@ class TestMmapValidation:
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False, model_path=str(tmp_path)
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
),
torch.device("cpu"),
)
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=(),
model_path=str(tmp_path),
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
backend._load_weights_via_mmap(pipeline, modules)
assert result.plan is not None
backend._load_weights_via_mmap(pipeline, modules, result.plan)
assert pipeline.transformer.validate_called, "validate_loaded_weights should be called"
class TestConfigValidation:
"""Tests for configuration validation in OffloadConfig / execute_request."""
def test_loader_plan_cannot_be_silently_dropped_on_unsupported_platform(self, monkeypatch):
import vllm_omni.diffusion.offloader as offloader_module
od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
dlo_resident_layers=0,
pin_cpu_memory=False,
model="unused",
parallel_config=SimpleNamespace(
use_hsdp=False,
data_parallel_size=1,
sequence_parallel_size=1,
),
)
plan = HostWeightPlan(
backing_kind="checkpoint_mmap",
bindings={},
)
monkeypatch.setattr(
offloader_module.current_omni_platform,
"supports_cpu_offload",
lambda: False,
)
with pytest.raises(RuntimeError, match="skipped ordinary weight materialization"):
offloader_module.get_offload_backend(
od_config,
device=torch.device("cpu"),
host_weight_plan=plan,
)
def test_hsdp_with_allgather_rejected(self):
"""HSDP + DLO + AllGather should raise ValueError (double sharding)."""
from vllm_omni.diffusion.offloader.base import OffloadConfig
@@ -79,11 +79,11 @@ def dist_group():
@pytest.fixture
def patched_offload_runtime(mocker):
mocker.patch.object(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
mocker.patch.object(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
mocker.patch.object(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
mocker.patch.object(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
def patched_offload_runtime(monkeypatch):
monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
class TinyBlock(nn.Module):
@@ -128,6 +128,33 @@ class TestLayerwiseOffloadHook:
assert current_block.weight.to_local().shape == torch.Size([4])
assert not hook.is_materialized
def test_prefetch_preserves_transposed_weight_stride(self, patched_offload_runtime):
"""Online-FP8 Cutlass weights must retain their transposed layout."""
class StridedBlock(nn.Module):
def __init__(self, start: float):
super().__init__()
base = torch.arange(start, start + 12).reshape(3, 4)
self.weight = nn.Parameter(base.t(), requires_grad=False)
current_block = StridedBlock(1.0)
next_block = StridedBlock(20.0)
expected = next_block.weight.detach().clone()
expected_stride = next_block.weight.stride()
assert expected_stride == (1, 4)
hook = LayerwiseOffloadHook(
next_block=next_block,
device=torch.device("cpu"),
stream=DummyStream(),
pin_memory=False,
)
hook.initialize_hook(current_block)
hook.prefetch_layer(non_blocking=False)
assert next_block.weight.stride() == expected_stride
assert torch.equal(next_block.weight, expected)
class _DummyBlock(nn.Module):
def __init__(self):
@@ -16,10 +16,11 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
)
from vllm_omni.quantization.component_config import (
ComponentQuantizationConfig,
resolve_component_quant_config,
)
from vllm_omni.quantization.inc_config import OmniINCConfig
pytestmark = [pytest.mark.core_model]
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
# ---------------------------------------------------------------------------
@@ -347,6 +348,24 @@ class TestTalkerVisualRouting:
# ===================================================================
class TestResolveComponentQuantConfig:
def test_global_config_applies_to_each_quantization_aware_component(self):
config = _MockQuantConfig("fp8")
assert resolve_component_quant_config(config, "transformer") is config
assert resolve_component_quant_config(config, "text_encoder") is config
def test_component_map_is_the_only_scope_override(self):
transformer = _MockQuantConfig("fp8")
config = ComponentQuantizationConfig(
component_configs={"transformer": transformer},
default_config=None,
)
assert resolve_component_quant_config(config, "transformer") is transformer
assert resolve_component_quant_config(config, "text_encoder") is None
class TestComponentResolve:
def test_longest_prefix_match(self):
a = _MockQuantConfig("a")
+133
View File
@@ -3,6 +3,8 @@
"""Unit tests for WorkerProc async output logic."""
import queue
import threading
import time
from unittest.mock import MagicMock
import pytest
@@ -29,9 +31,59 @@ def _make_worker_proc(step_execution=False):
proc._running = True
proc._async_output_queue = queue.Queue() if not step_execution else None
proc._async_output_thread = None
proc._async_output_done = threading.Condition()
proc._async_output_pending = 0
proc._result_mq_lock = threading.Lock()
return proc
class _OverlapDetectingQueue:
"""Fails if two threads are inside enqueue() at the same time.
Mirrors vLLM's MessageQueue, whose acquire_write() reads current_idx,
writes that block, then advances it. Overlapping writers can target the
same ring block and silently drop a message.
"""
def __init__(self):
self._active = 0
self._guard = threading.Lock()
self.overlaps = 0
self.messages = []
def enqueue(self, msg, timeout=None):
with self._guard:
self._active += 1
if self._active > 1:
self.overlaps += 1
time.sleep(0.001) # widen the window a real ring-buffer write would have
with self._guard:
self.messages.append(msg)
self._active -= 1
class TestResultQueueSingleWriter:
"""Worker must serialize writes to the single-writer result queue."""
def test_concurrent_enqueues_do_not_overlap(self):
proc = _make_worker_proc(step_execution=False)
detector = _OverlapDetectingQueue()
proc.result_mq = detector
def writer(tag):
for i in range(20):
proc._enqueue_result(f"{tag}-{i}")
threads = [threading.Thread(target=writer, args=(tag,)) for tag in ("compute", "output")]
for t in threads:
t.start()
for t in threads:
t.join()
assert detector.overlaps == 0, f"{detector.overlaps} overlapping writes to result_mq"
assert len(detector.messages) == 40
class TestGenerateAsyncOutputId:
"""Test _generate_async_output_id uniqueness."""
@@ -192,6 +244,87 @@ class TestAsyncOutputLoopLogic:
assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls"
class TestDrainAsyncOutputs:
"""Async outputs must finish packing before device memory is released."""
def test_drain_returns_immediately_when_nothing_in_flight(self):
proc = _make_worker_proc(step_execution=False)
assert proc.drain_async_outputs(timeout=0.01) is True
def test_return_result_keeps_output_pending_until_loop_completes(self, mocker):
"""_return_result marks work pending; the bg loop clears it."""
proc = _make_worker_proc(step_execution=False)
mock_platform = mocker.MagicMock()
mock_platform.record_device_event.return_value = MagicMock()
mocker.patch(
"vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform",
mock_platform,
)
mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm")
mock_torch = mocker.MagicMock()
mock_torch.accelerator.current_accelerator.return_value.type = "cpu"
mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch)
proc._return_result(DiffusionOutput(output="data"), rpc_id="1")
# Packing has not run yet, so a sleep here would free tensors in flight.
assert proc._async_output_pending == 1
assert proc.drain_async_outputs(timeout=0.05) is False
thread = threading.Thread(target=proc._async_output_loop, daemon=True)
thread.start()
try:
assert proc.drain_async_outputs(timeout=5.0) is True
finally:
proc._async_output_queue.put(None)
thread.join(timeout=2.0)
def test_sleep_rpc_drains_before_releasing_memory(self, mocker):
proc = _make_worker_proc(step_execution=False)
proc.worker = MagicMock()
call_order: list[str] = []
mocker.patch.object(
proc,
"drain_async_outputs",
side_effect=lambda *a, **kw: call_order.append("drain") or True,
)
proc.worker.execute_method.side_effect = lambda *a, **kw: call_order.append("sleep")
proc._execute_rpc(
{
"method": "handle_sleep_task",
"args": (),
"kwargs": {},
"output_rank": 0,
"exec_all_ranks": False,
"collect_rank_status": False,
}
)
assert call_order == ["drain", "sleep"]
def test_non_sleep_rpc_does_not_drain(self, mocker):
proc = _make_worker_proc(step_execution=False)
proc.worker = MagicMock()
drain = mocker.patch.object(proc, "drain_async_outputs", return_value=True)
proc._execute_rpc(
{
"method": "execute_model",
"args": (),
"kwargs": {},
"output_rank": 0,
"exec_all_ranks": False,
"collect_rank_status": False,
}
)
drain.assert_not_called()
class TestWorkerProcShutdown:
def test_shutdown_stops_async_thread_and_releases_queues(self):
import threading
@@ -123,6 +123,7 @@ class _ConcurrencyTrackingExecutor:
def _make_request(tag: str):
sampling_params = dict(_SAMPLING_KEY_DEFAULTS)
sampling_params["num_inference_steps"] = 1
sampling_params["extra_args"] = {}
return SimpleNamespace(
request_id=tag,
prompt=f"prompt_{tag}",
+43
View File
@@ -305,6 +305,49 @@ def test_batch_runner_output_round_trips_nested_results_through_shm() -> None:
assert output["req-error"].result.error == "boom"
def test_per_request_views_of_a_batch_tensor_do_not_ship_whole_storage() -> None:
"""Regression: per-request slices used to pickle the whole batch storage.
Each view is under the packing threshold, but pickle writes the entire
shared storage per view, so a batch's message grew by the full batch
tensor once per request and overflowed the ring buffer.
"""
import pickle
num_requests = 16
# Each row stays under the threshold; the shared storage is well over it.
row_numel = _SHM_TENSOR_THRESHOLD // (2 * torch.empty((), dtype=torch.float32).element_size())
batch = torch.arange(num_requests * row_numel, dtype=torch.float32).reshape(num_requests, row_numel)
per_request = batch[0]
assert per_request.nelement() * per_request.element_size() <= _SHM_TENSOR_THRESHOLD
assert per_request.untyped_storage().nbytes() > _SHM_TENSOR_THRESHOLD
output = BatchRunnerOutput.from_list(
[
RunnerOutput(
request_id=f"req-{i}",
finished=True,
result=DiffusionOutput(trajectory_latents=batch[i]),
)
for i in range(num_requests)
]
)
unpacked_size = len(pickle.dumps(output, protocol=pickle.HIGHEST_PROTOCOL))
pack_diffusion_output_shm(output)
packed_size = len(pickle.dumps(output, protocol=pickle.HIGHEST_PROTOCOL))
assert output.runner_outputs[0].result.trajectory_latents["__tensor_shm__"] is True
# Whole storage per request before; handles only after.
storage_bytes = per_request.untyped_storage().nbytes()
assert unpacked_size > num_requests * storage_bytes / 2
assert packed_size < _SHM_TENSOR_THRESHOLD
unpack_diffusion_output_shm(output)
for i in range(num_requests):
torch.testing.assert_close(output[f"req-{i}"].result.trajectory_latents, batch[i])
def test_pack_value_keeps_tensor_at_threshold_inline() -> None:
tensor = torch.arange(
_SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(),
@@ -740,6 +740,9 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
del kwargs
return SimpleNamespace(transformer=torch.nn.Identity())
def take_host_weight_plan(self):
return None
class _DummyMemoryProfiler:
consumed_memory = 0
@@ -779,7 +782,11 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader)
monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
monkeypatch.setattr(
model_runner_module,
"get_offload_backend",
lambda od_config, device, host_weight_plan: None,
)
monkeypatch.setattr(
model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
)
+133 -1
View File
@@ -217,13 +217,21 @@ class TestGetRequestBatchSamplingParamsKey:
num_inference_steps: int = 2,
seed: int | None = 123,
generator: torch.Generator | None = None,
extra_args: dict | None = None,
condition_key: tuple | None = None,
) -> OmniDiffusionRequest:
sp = OmniDiffusionSamplingParams(
num_inference_steps=num_inference_steps,
seed=seed,
generator=generator,
extra_args=extra_args or {},
)
return OmniDiffusionRequest(
prompt="prompt",
sampling_params=sp,
request_id=f"req-{num_inference_steps}",
batch_compatibility_key=condition_key,
)
return OmniDiffusionRequest(prompt="prompt", sampling_params=sp, request_id=f"req-{num_inference_steps}")
def test_distinguishes_num_inference_steps(self) -> None:
scheduler = RequestScheduler()
@@ -240,6 +248,57 @@ class TestGetRequestBatchSamplingParamsKey:
self._make(seed=1, generator=gen_a)
) == scheduler._build_sampling_params_key(self._make(seed=2, generator=gen_b))
@pytest.mark.parametrize(
("first_extra_args", "second_extra_args"),
[
({"sample_solver": "unipc"}, {"sample_solver": "euler"}),
({"flow_shift": 3.0}, {"flow_shift": 5.0}),
],
)
def test_distinguishes_wan_scheduler_structure(
self,
first_extra_args: dict,
second_extra_args: dict,
) -> None:
scheduler = RequestScheduler()
assert scheduler._build_sampling_params_key(
self._make(extra_args=first_extra_args)
) != scheduler._build_sampling_params_key(self._make(extra_args=second_extra_args))
@pytest.mark.parametrize(
("first_extra_args", "second_extra_args"),
[
({"sample_solver": " Euler "}, {"sample_solver": "euler"}),
({"flow_shift": "5.0"}, {"flow_shift": 5.0}),
],
)
def test_normalizes_equivalent_wan_scheduler_structure(
self,
first_extra_args: dict,
second_extra_args: dict,
) -> None:
scheduler = RequestScheduler()
assert scheduler._build_sampling_params_key(
self._make(extra_args=first_extra_args)
) == scheduler._build_sampling_params_key(self._make(extra_args=second_extra_args))
def test_uses_none_for_unspecified_wan_scheduler_structure(self) -> None:
scheduler = RequestScheduler()
key = scheduler._build_sampling_params_key(self._make())
assert key.sample_solver is None
assert key.flow_shift is None
def test_distinguishes_pipeline_condition_structure(self) -> None:
scheduler = RequestScheduler()
assert scheduler._build_sampling_params_key(
self._make(condition_key=("wan22_s2v_condition", True))
) != scheduler._build_sampling_params_key(self._make(condition_key=("wan22_s2v_condition", False)))
class TestRequestScheduler:
def setup_method(self) -> None:
@@ -602,6 +661,38 @@ class TestRequestScheduler:
assert first.num_running_reqs == 1
assert first.num_waiting_reqs == 1
def test_batches_incompatible_pipeline_conditions_separately(self) -> None:
scheduler = RequestScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
request_a = OmniDiffusionRequest(
prompt="a",
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, seed=123),
request_id="a",
batch_compatibility_key=("wan22_s2v_condition", True),
)
request_b = OmniDiffusionRequest(
prompt="b",
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, seed=456),
request_id="b",
batch_compatibility_key=("wan22_s2v_condition", False),
)
scheduler.add_request(request_a)
scheduler.add_request(request_b)
first = scheduler.schedule()
assert _new_ids(first) == ["a"]
assert first.num_running_reqs == 1
assert first.num_waiting_reqs == 1
scheduler.update_from_output(first, _make_request_output("a"))
second = scheduler.schedule()
assert _new_ids(second) == ["b"]
assert second.num_running_reqs == 1
assert second.num_waiting_reqs == 0
def test_batches_different_request_local_seed_together(self) -> None:
scheduler = RequestScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
@@ -625,6 +716,47 @@ class TestRequestScheduler:
assert first.num_running_reqs == 2
assert first.num_waiting_reqs == 0
@pytest.mark.parametrize(
("first_extra_args", "second_extra_args"),
[
({"sample_solver": "unipc"}, {"sample_solver": "euler"}),
({"flow_shift": 3.0}, {"flow_shift": 5.0}),
({"sample_solver": "unipc"}, {}),
({"flow_shift": 3.0}, {}),
],
)
def test_batches_incompatible_wan_scheduler_structure_separately(
self,
first_extra_args: dict,
second_extra_args: dict,
) -> None:
scheduler = RequestScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
first_id = scheduler.add_request(
_make_step_request(
"a",
sampling_params=OmniDiffusionSamplingParams(
num_inference_steps=2,
extra_args=first_extra_args,
),
)
)
scheduler.add_request(
_make_step_request(
"b",
sampling_params=OmniDiffusionSamplingParams(
num_inference_steps=2,
extra_args=second_extra_args,
),
)
)
first = scheduler.schedule()
assert _new_ids(first) == [first_id]
assert first.num_running_reqs == 1
assert first.num_waiting_reqs == 1
def test_incompatible_waiting_head_blocks_later_compatible_request(self) -> None:
scheduler = RequestScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=3))
@@ -0,0 +1,171 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import threading
import time
from unittest.mock import Mock
import pytest
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
class _DelayedMultipartSocket:
def __init__(
self,
socket,
send_started: threading.Event,
release_send: threading.Event,
):
self._socket = socket
self._send_started = send_started
self._release_send = release_send
self._delay_next_send = True
def __getattr__(self, name):
return getattr(self._socket, name)
def send_multipart(self, *args, **kwargs):
if self._delay_next_send:
self._delay_next_send = False
self._send_started.set()
if not self._release_send.wait(timeout=2):
raise TimeoutError("Timed out waiting to release delayed overflow payload")
return self._socket.send_multipart(*args, **kwargs)
class _ObservedPollSocket:
def __init__(self, socket, poll_started: threading.Event):
self._socket = socket
self._poll_started = poll_started
def __getattr__(self, name):
return getattr(self._socket, name)
def poll(self, *args, **kwargs):
self._poll_started.set()
return self._socket.poll(*args, **kwargs)
class _RecordingWorker:
def __init__(self):
self.calls: list[str] = []
def execute_method(self, method, *args, **kwargs):
assert method == "record"
self.calls.append(args[0])
def test_recv_message_waits_indefinitely():
worker_proc = WorkerProc.__new__(WorkerProc)
worker_proc.mq = Mock()
worker_proc.mq.dequeue.return_value = {"type": "rpc"}
result = worker_proc.recv_message()
assert result == {"type": "rpc"}
worker_proc.mq.dequeue.assert_called_once_with(indefinite=True)
def test_busy_loop_preserves_order_when_overflow_payload_is_delayed(monkeypatch):
"""A consumed overflow marker must not be abandoned on a socket timeout."""
writer = MessageQueue(
n_reader=1,
n_local_reader=1,
local_reader_ranks=[0],
max_chunk_bytes=1024,
max_chunks=4,
)
reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0)
writer.wait_until_ready()
reader.wait_until_ready()
overflow_send_started = threading.Event()
release_overflow_send = threading.Event()
overflow_poll_started = threading.Event()
writer.local_socket = _DelayedMultipartSocket(
writer.local_socket,
overflow_send_started,
release_overflow_send,
)
reader.local_socket = _ObservedPollSocket(reader.local_socket, overflow_poll_started)
recording_worker = _RecordingWorker()
worker_proc = WorkerProc.__new__(WorkerProc)
worker_proc.gpu_id = 0
worker_proc.mq = reader
worker_proc.worker = recording_worker
worker_proc.result_mq = None
worker_proc.wake_event = None
worker_proc._running = True
# Make the pre-fix implementation fail quickly: it read this module-level
# value directly from _worker_busy_loop. The fixed implementation ignores it
# and delegates to recv_message(), which waits indefinitely.
monkeypatch.setattr(
diffusion_worker_module,
"_BROADCAST_DEQUEUE_TIMEOUT_S",
0.02,
raising=False,
)
overflow_message = {
"type": "rpc",
"method": "record",
"args": ("overflow-1",),
"padding": b"x" * 2048,
}
inline_message = {
"type": "rpc",
"method": "record",
"args": ("inline-2",),
}
shutdown_message = {"type": "shutdown"}
thread_errors: list[BaseException] = []
def send_messages():
try:
writer.enqueue(overflow_message)
writer.enqueue(inline_message)
writer.enqueue(shutdown_message)
except BaseException as exc:
thread_errors.append(exc)
def run_busy_loop():
try:
worker_proc._worker_busy_loop()
except BaseException as exc:
thread_errors.append(exc)
writer_thread = threading.Thread(target=send_messages, daemon=True)
worker_thread = threading.Thread(target=run_busy_loop, daemon=True)
try:
writer_thread.start()
worker_thread.start()
assert overflow_send_started.wait(timeout=2)
assert overflow_poll_started.wait(timeout=2)
# Keep the multipart unavailable beyond the pre-fix 20 ms timeout.
time.sleep(0.1)
release_overflow_send.set()
writer_thread.join(timeout=2)
worker_thread.join(timeout=2)
assert not writer_thread.is_alive()
assert not worker_thread.is_alive()
assert thread_errors == []
assert recording_worker.calls == ["overflow-1", "inline-2"]
finally:
release_overflow_send.set()
worker_proc._running = False
reader.shutdown()
writer.shutdown()
writer_thread.join(timeout=1)
worker_thread.join(timeout=1)
+114
View File
@@ -6,6 +6,7 @@ import concurrent.futures
import queue
import threading
import time
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
@@ -264,6 +265,119 @@ class TestResultPumpDispatch:
fut = executor.wait_output_ready(async_output_id)
assert fut.done()
assert fut.result(timeout=1.0) is output
assert async_output_id not in executor._output_futures
def test_output_ready_atomic_resolution_when_future_already_waiting(self):
"""When OUTPUT_READY arrives and a future is already waiting, resolve it directly."""
executor = _make_executor()
async_output_id = "abc123"
output = DiffusionOutput(output="waiting")
fut = executor.wait_output_ready(async_output_id)
assert not fut.done()
msg = AsyncDiffusionOutput(
kind=AsyncOutputKind.OUTPUT_READY,
async_output_id=async_output_id,
output=output,
)
_feed_one_msg_to_pump(executor, msg)
assert fut.done()
assert fut.result(timeout=1.0) is output
assert async_output_id not in executor._output_futures
assert async_output_id not in executor._completed_outputs
class _FakeBatchOutput:
"""Batch-level output exposing per-request results."""
def __init__(self, results):
self._results = results
def get_request_output(self, req_id):
result = self._results.get(req_id)
if result is None:
return None
return SimpleNamespace(result=result)
def _make_scheduler_output(req_ids):
return SimpleNamespace(
scheduled_new_reqs=[SimpleNamespace(request_id=rid, req=SimpleNamespace()) for rid in req_ids]
)
class TestBatchSplitDelivery:
"""execute_batch must resolve per-request futures in either arrival order."""
@staticmethod
def _run(executor, req_ids, batch_id, deliver_early):
# Keep execute_batch on the fused request-batch path (not DLO DP).
executor.od_config.parallel_config.data_parallel_size = 1
executor.od_config.enable_distributed_layerwise_offload = False
executor._ensure_open = lambda: None
outputs = {rid: DiffusionOutput(output=f"img-{rid}") for rid in req_ids}
ready = AsyncDiffusionOutput(
kind=AsyncOutputKind.OUTPUT_READY,
async_output_id=batch_id,
output=_FakeBatchOutput(outputs),
)
def fake_collective_rpc(*args, **kwargs):
if deliver_early:
# Worker's background D2H/SHM thread wins the race: OUTPUT_READY
# is pumped before execute_batch registers the split map.
_feed_one_msg_to_pump(executor, ready)
return AsyncDiffusionOutput(
kind=AsyncOutputKind.COMPUTE_DONE,
rpc_id="1",
async_output_id=batch_id,
)
executor.collective_rpc = fake_collective_rpc
batch = executor.execute_batch(_make_scheduler_output(req_ids))
if not deliver_early:
_feed_one_msg_to_pump(executor, ready)
return batch, outputs
def test_output_ready_after_split_map(self):
executor = _make_executor()
req_ids = ["r0", "r1", "r2"]
_, outputs = self._run(executor, req_ids, "batch-1", deliver_early=False)
for rid in req_ids:
fut = executor.wait_output_ready(f"batch-1/{rid}")
assert fut.done()
assert fut.result(timeout=1.0) is outputs[rid]
def test_output_ready_before_split_map(self):
"""Regression: the whole batch used to be lost, hanging every request."""
executor = _make_executor()
req_ids = ["r0", "r1", "r2"]
_, outputs = self._run(executor, req_ids, "batch-1", deliver_early=True)
for rid in req_ids:
fut = executor.wait_output_ready(f"batch-1/{rid}")
assert fut.done(), f"request {rid} never resolved"
assert fut.result(timeout=1.0) is outputs[rid]
# No stale batch-level state left behind.
assert executor._batch_split_map == {}
assert executor._completed_outputs == {}
def test_engine_waiter_before_output_is_resolved(self):
"""Consumers already blocked in wait_output_ready must be woken."""
executor = _make_executor()
req_ids = ["r0", "r1"]
batch_id = "batch-2"
futures = {rid: executor.wait_output_ready(f"{batch_id}/{rid}") for rid in req_ids}
_, outputs = self._run(executor, req_ids, batch_id, deliver_early=True)
for rid in req_ids:
assert futures[rid].done(), f"request {rid} never resolved"
assert futures[rid].result(timeout=1.0) is outputs[rid]
class TestShutdownCleansUpFutures:
@@ -0,0 +1,55 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Guard the public accessor that examples and docs use to read omni outputs.
``OmniRequestOutput`` used to hold a nested ``RequestOutput`` in a
``request_output`` field. Since #5146 it *inherits* ``RequestOutput``, so the
generation content lives on the object itself and the nested field is gone.
Call sites that kept the old accessor only fail once a real model runs, which
means a GPU nightly job is the first thing to notice. These checks pin both
halves of the contract on CPU instead: the field is really gone, and no doc,
example, test or tool still reaches for it.
"""
import re
from pathlib import Path
import pytest
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
ROOT_DIR = Path(__file__).parents[2]
SCANNED_DIRS = ("docs", "examples", "tests", "tools")
SCANNED_SUFFIXES = (".py", ".md")
# ``.request_outputs`` (plural, an OutputProcessor batch) is a different and
# still-valid attribute, so require a non-identifier character after the name.
NESTED_ACCESSOR = re.compile(r"\.request_output(?![\w])")
def _iter_scanned_files():
for directory in SCANNED_DIRS:
for path in sorted((ROOT_DIR / directory).rglob("*")):
if path.suffix in SCANNED_SUFFIXES and path.resolve() != Path(__file__).resolve():
yield path
def test_omni_request_output_has_no_nested_request_output():
assert not hasattr(OmniRequestOutput(request_id="req-accessor"), "request_output")
def test_no_source_reads_the_nested_request_output():
offenders = []
for path in _iter_scanned_files():
for lineno, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
if NESTED_ACCESSOR.search(line):
offenders.append(f"{path.relative_to(ROOT_DIR)}:{lineno}: {line.strip()}")
assert not offenders, (
"OmniRequestOutput no longer nests a RequestOutput; read the content off the "
"output itself (e.g. `outputs[0].images`) instead of `.request_output`:\n" + "\n".join(offenders)
)
@@ -0,0 +1,591 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Generate raw LTX video and audio outputs with the official or Omni runtime."""
from __future__ import annotations
import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any
import numpy as np
import torch
from PIL import Image
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", choices=("official", "omni"), required=True)
parser.add_argument("--request", type=Path, required=True)
parser.add_argument("--output-dir", type=Path, required=True)
parser.add_argument("--model")
parser.add_argument("--model-class-name")
parser.add_argument("--official-root", type=Path)
parser.add_argument("--checkpoint", type=Path)
parser.add_argument("--gemma-root", type=Path)
parser.add_argument("--transformer-path", type=Path)
parser.add_argument("--text-encoder-path", type=Path)
parser.add_argument("--video-vae-path", type=Path)
parser.add_argument("--audio-vae-path", type=Path)
parser.add_argument("--spatial-upsampler-path", type=Path)
parser.add_argument("--distilled-lora-path", type=Path)
parser.add_argument("--connector-model", type=Path)
parser.add_argument(
"--official-pipeline",
choices=("distilled_two_stage", "full_one_stage", "full_two_stage"),
default="distilled_two_stage",
)
parser.add_argument("--enable-layerwise-offload", action="store_true")
return parser.parse_args()
def _save_outputs(
output_dir: Path,
*,
video: torch.Tensor,
audio: torch.Tensor,
audio_sample_rate: int,
fps: float,
metadata: dict[str, Any],
) -> None:
output_dir.mkdir(parents=True, exist_ok=True)
video = video.detach().float().cpu().clamp(0.0, 1.0)
audio = audio.detach().float().cpu()
np.save(output_dir / "video.npy", video.numpy())
np.save(output_dir / "audio.npy", audio.numpy())
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
frames_u8 = video.mul(255.0).round().to(torch.uint8).numpy()
audio_array = audio.numpy()
while audio_array.ndim > 2 and audio_array.shape[0] == 1:
audio_array = audio_array[0]
(output_dir / "output.mp4").write_bytes(
mux_video_audio_bytes(
frames_u8,
audio_array,
fps=float(fps),
audio_sample_rate=int(audio_sample_rate),
video_codec_options={"preset": "ultrafast"},
)
)
frame_indices = sorted({0, video.shape[0] // 2, video.shape[0] - 1})
for index in frame_indices:
frame = video[index].mul(255.0).round().to(torch.uint8).numpy()
Image.fromarray(frame).save(output_dir / f"frame_{index:04d}.png")
metadata.update(
{
"video_shape": list(video.shape),
"audio_shape": list(audio.shape),
"audio_sample_rate": int(audio_sample_rate),
"frame_indices": frame_indices,
"mp4": "output.mp4",
}
)
(output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
def _reset_accelerator_peak_memory_stats() -> None:
if torch.accelerator.is_available():
torch.accelerator.reset_peak_memory_stats()
def _accelerator_peak_memory_stats() -> dict[str, float]:
if not torch.accelerator.is_available():
return {}
torch.accelerator.synchronize()
mib = 1024**2
return {
"peak_memory_allocated_mb": torch.accelerator.max_memory_allocated() / mib,
"peak_memory_reserved_mb": torch.accelerator.max_memory_reserved() / mib,
}
def _insert_official_paths(official_root: Path) -> None:
for relative_path in ("packages/ltx-core/src", "packages/ltx-pipelines/src"):
path = str((official_root / relative_path).resolve())
if path not in sys.path:
sys.path.insert(0, path)
def _configure_official_sdpa(pipeline: Any) -> None:
"""Pin official connector and denoiser attention to cuDNN, with MATH fallback."""
from ltx_core.loader.attention_ops import set_attention_module_op
from ltx_core.model.transformer.attention import PytorchAttention
from torch.nn.attention import SDPBackend
attention = PytorchAttention(priority=[SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH])
module_op = set_attention_module_op(
attention=attention,
masked_attention=attention,
)
transformer_owners = [
owner
for attribute in ("stage", "stage_1", "stage_2")
if (owner := getattr(pipeline, attribute, None)) is not None
]
owners_and_attributes = [
*((owner, "_transformer_builder") for owner in transformer_owners),
(pipeline.prompt_encoder, "_embeddings_processor_builder"),
]
for owner, attribute in owners_and_attributes:
builder = getattr(owner, attribute)
setattr(
owner,
attribute,
builder.with_module_ops((*builder.module_ops, module_op)),
)
def _use_diffusers_connector_weights(prompt_encoder: Any, model: Path) -> None:
"""Run official connector code with the checkpoint under test."""
from safetensors import safe_open
connector_root = model / "connectors"
shards = sorted(connector_root.glob("*.safetensors"))
if not shards:
raise ValueError(f"No connector safetensors found under {connector_root}")
original_build_embeddings_processor = prompt_encoder._build_embeddings_processor
def build_embeddings_processor():
processor = original_build_embeddings_processor()
parameters = dict(processor.named_parameters())
expected = {name for name in parameters if name.startswith(("video_connector.", "audio_connector."))}
loaded: set[str] = set()
with torch.no_grad():
for shard in shards:
with safe_open(str(shard), framework="pt", device="cpu") as weights:
for source_name in weights.keys():
if not source_name.startswith(("video_connector.", "audio_connector.")):
continue
target_name = (
source_name.replace(".transformer_blocks.", ".transformer_1d_blocks.")
.replace(".attn1.norm_q.", ".attn1.q_norm.")
.replace(".attn1.norm_k.", ".attn1.k_norm.")
)
parameter = parameters[target_name]
parameter.copy_(weights.get_tensor(source_name).to(parameter.device, parameter.dtype))
loaded.add(target_name)
if loaded != expected:
missing = sorted(expected - loaded)
unexpected = sorted(loaded - expected)
raise ValueError(f"Connector weight mismatch: missing={missing}, unexpected={unexpected}")
return processor
prompt_encoder._build_embeddings_processor = build_embeddings_processor
@torch.inference_mode()
def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
if args.official_root is None:
raise ValueError("Official backend requires --official-root")
_insert_official_paths(args.official_root)
# Keep Gemma on the same SDPA path in both subprocesses. DiT attention is
# still pinned independently to the explicit cuDNN backend below.
torch.backends.cuda.enable_cudnn_sdp(False)
_reset_accelerator_peak_memory_stats()
from ltx_pipelines.utils.args import ImageConditioningInput
from ltx_pipelines.utils.model_paths import ModelPaths
from ltx_pipelines.utils.types import OffloadMode
offload_mode = OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE
image_path = request.get("image")
images = (
[]
if image_path is None
else [
ImageConditioningInput(
path=str(image_path),
frame_idx=0,
strength=1.0,
crf=request.get("image_crf", 18),
)
]
)
if args.transformer_path is not None:
required = {
"--text-encoder-path": args.text_encoder_path,
"--video-vae-path": args.video_vae_path,
"--audio-vae-path": args.audio_vae_path,
}
if args.official_pipeline in {"distilled_two_stage", "full_two_stage"}:
required["--spatial-upsampler-path"] = args.spatial_upsampler_path
if args.official_pipeline == "full_two_stage":
required["--distilled-lora-path"] = args.distilled_lora_path
missing = [name for name, value in required.items() if value is None]
if missing:
raise ValueError(f"Official LTX-2.5 split backend is missing: {', '.join(missing)}")
model_paths = ModelPaths.from_split(
transformer_path=str(args.transformer_path),
text_encoder_path=str(args.text_encoder_path),
video_vae_path=str(args.video_vae_path),
audio_vae_path=str(args.audio_vae_path),
)
if args.official_pipeline == "distilled_two_stage":
from ltx_pipelines.distilled import DistilledPipeline
pipeline = DistilledPipeline(
model_paths=model_paths,
spatial_upsampler_path=str(args.spatial_upsampler_path),
loras=(),
offload_mode=offload_mode,
)
if args.connector_model is not None:
_use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model)
_configure_official_sdpa(pipeline)
video, audio, _, _ = pipeline(
prompt=request["prompt"],
seed=request["seed"],
height=request["height"],
width=request["width"],
num_frames=request["num_frames"],
frame_rate=request["fps"],
images=images,
)
pipeline_name = "DistilledPipeline"
elif args.official_pipeline == "full_two_stage":
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
pipeline = TI2VidTwoStagesPipeline(
model_paths=model_paths,
distilled_lora=[
LoraPathStrengthAndSDOps(
path=str(args.distilled_lora_path),
strength=1.0,
sd_ops=LTXV_LORA_COMFY_RENAMING_MAP,
)
],
spatial_upsampler_path=str(args.spatial_upsampler_path),
loras=[],
offload_mode=offload_mode,
)
if args.connector_model is not None:
_use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model)
_configure_official_sdpa(pipeline)
video, audio, _, _ = pipeline(
prompt=request["prompt"],
negative_prompt=request["negative_prompt"],
seed=request["seed"],
height=request["height"],
width=request["width"],
num_frames=request["num_frames"],
frame_rate=request["fps"],
num_inference_steps=request["num_inference_steps"],
video_guider_params=MultiModalGuiderParams(
cfg_scale=request["video_cfg_scale"],
stg_scale=request["video_stg_scale"],
rescale_scale=request["video_rescale_scale"],
modality_scale=request["video_modality_scale"],
skip_step=0,
stg_blocks=request["video_stg_blocks"],
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=request["audio_cfg_scale"],
stg_scale=request["audio_stg_scale"],
rescale_scale=request["audio_rescale_scale"],
modality_scale=request["audio_modality_scale"],
skip_step=0,
stg_blocks=request["audio_stg_blocks"],
),
images=images,
max_batch_size=4,
stage_1_sigmas=torch.tensor(request["stage_1_sigmas"], dtype=torch.float32),
stage_2_sigmas=torch.tensor(request["stage_2_sigmas"], dtype=torch.float32),
)
pipeline_name = "TI2VidTwoStagesPipeline"
else:
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
pipeline = TI2VidOneStagePipeline(
model_paths=model_paths,
loras=(),
offload_mode=offload_mode,
)
if args.connector_model is not None:
_use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model)
_configure_official_sdpa(pipeline)
request_sigmas = request.get("sigmas")
video, audio, _ = pipeline(
prompt=request["prompt"],
negative_prompt=request["negative_prompt"],
seed=request["seed"],
height=request["height"],
width=request["width"],
num_frames=request["num_frames"],
frame_rate=request["fps"],
num_inference_steps=request["num_inference_steps"],
video_guider_params=MultiModalGuiderParams(
cfg_scale=request["video_cfg_scale"],
stg_scale=request["video_stg_scale"],
rescale_scale=request["video_rescale_scale"],
modality_scale=request["video_modality_scale"],
skip_step=0,
stg_blocks=request["video_stg_blocks"],
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=request["audio_cfg_scale"],
stg_scale=request["audio_stg_scale"],
rescale_scale=request["audio_rescale_scale"],
modality_scale=request["audio_modality_scale"],
skip_step=0,
stg_blocks=request["audio_stg_blocks"],
),
images=images,
max_batch_size=4,
sigmas=None if request_sigmas is None else torch.tensor(request_sigmas, dtype=torch.float32),
)
pipeline_name = "TI2VidOneStagePipeline"
checkpoint_metadata = {
"transformer_path": str(args.transformer_path),
"video_vae_path": str(args.video_vae_path),
"connector_model": None if args.connector_model is None else str(args.connector_model),
"pipeline": pipeline_name,
}
else:
if args.checkpoint is None or args.gemma_root is None:
raise ValueError("Official monolith backend requires --checkpoint and --gemma-root")
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
pipeline = TI2VidOneStagePipeline(
model_paths=ModelPaths.from_monolith(str(args.checkpoint), str(args.gemma_root)),
loras=(),
offload_mode=offload_mode,
)
_configure_official_sdpa(pipeline)
video, audio, _ = pipeline(
prompt=request["prompt"],
negative_prompt=request["negative_prompt"],
seed=request["seed"],
height=request["height"],
width=request["width"],
num_frames=request["num_frames"],
frame_rate=request["fps"],
num_inference_steps=request["num_inference_steps"],
video_guider_params=MultiModalGuiderParams(
cfg_scale=request["video_cfg_scale"],
stg_scale=request["video_stg_scale"],
rescale_scale=request["video_rescale_scale"],
modality_scale=request["video_modality_scale"],
skip_step=0,
stg_blocks=request["video_stg_blocks"],
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=request["audio_cfg_scale"],
stg_scale=request["audio_stg_scale"],
rescale_scale=request["audio_rescale_scale"],
modality_scale=request["audio_modality_scale"],
skip_step=0,
stg_blocks=request["audio_stg_blocks"],
),
images=images,
max_batch_size=4,
)
checkpoint_metadata = {
"checkpoint": str(args.checkpoint),
"pipeline": "TI2VidOneStagePipeline",
}
video_tensor = torch.cat([chunk.detach().cpu() for chunk in video], dim=0)
_save_outputs(
args.output_dir,
video=_canonical_video(video_tensor),
audio=audio.waveform,
audio_sample_rate=audio.sampling_rate,
fps=float(request["fps"]),
metadata={
"backend": "official",
"attention_backend": "torch_sdpa_cudnn_bf16",
"official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"),
"seed": request["seed"],
"sigmas": request.get("sigmas"),
"stage_1_sigmas": request.get("stage_1_sigmas"),
"stage_2_sigmas": request.get("stage_2_sigmas"),
**_accelerator_peak_memory_stats(),
**checkpoint_metadata,
},
)
def _omni_worker_peak_memory_mb(output: Any) -> float:
result = output[0] if isinstance(output, list) and output else output
return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
def _unwrap_omni_output(output: Any) -> tuple[Any, Any, int]:
from vllm_omni.outputs import OmniRequestOutput
audio = None
audio_sample_rate = None
frames = output[0] if isinstance(output, list) and output else output
if isinstance(frames, OmniRequestOutput):
multimodal_output = frames.multimodal_output or {}
audio = multimodal_output.get("audio")
audio_sample_rate = multimodal_output.get("audio_sample_rate")
nested_output = getattr(frames, "request_output", None)
if frames.is_pipeline_output and isinstance(nested_output, OmniRequestOutput):
frames = nested_output
multimodal_output = frames.multimodal_output or {}
audio = multimodal_output.get("audio", audio)
audio_sample_rate = multimodal_output.get("audio_sample_rate", audio_sample_rate)
if isinstance(frames, OmniRequestOutput):
if not frames.images:
raise ValueError("No video frames found in OmniRequestOutput")
frames = frames.images
if isinstance(frames, list) and len(frames) == 1:
frames = frames[0]
if isinstance(frames, tuple) and len(frames) == 2:
frames, audio = frames
elif isinstance(frames, dict):
audio = frames.get("audio", audio)
audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate)
frames = frames.get("frames", frames.get("video"))
if frames is None or audio is None or audio_sample_rate is None:
raise ValueError("Omni output did not contain video, audio, and audio_sample_rate")
return frames, audio, int(audio_sample_rate)
def _canonical_video(video: Any) -> torch.Tensor:
if isinstance(video, list):
if len(video) == 1:
return _canonical_video(video[0])
tensors = [torch.as_tensor(np.asarray(item) if not isinstance(item, torch.Tensor) else item) for item in video]
if all(item.ndim == 3 for item in tensors):
video = torch.stack(tensors)
elif all(item.ndim == 4 for item in tensors):
video = torch.cat(tensors)
else:
raise ValueError(f"Cannot combine video items with dimensions {[item.ndim for item in tensors]}")
tensor = torch.as_tensor(np.asarray(video) if not isinstance(video, torch.Tensor) else video).detach().cpu()
if tensor.ndim == 5:
tensor = tensor[0]
if tensor.ndim != 4:
raise ValueError(f"Expected a 4D video tensor, got {tuple(tensor.shape)}")
if tensor.shape[-1] in (3, 4):
tensor = tensor[..., :3]
elif tensor.shape[1] in (3, 4):
tensor = tensor[:, :3].permute(0, 2, 3, 1)
elif tensor.shape[0] in (3, 4):
tensor = tensor[:3].permute(1, 2, 3, 0)
else:
raise ValueError(f"Cannot infer video channel dimension from {tuple(tensor.shape)}")
tensor = tensor.float()
if tensor.numel() and tensor.max() > 1:
tensor = tensor / 255.0
elif tensor.numel() and tensor.min() < 0:
tensor = tensor.clamp(-1.0, 1.0).add(1.0).mul(0.5)
return tensor.clamp(0.0, 1.0)
@torch.inference_mode()
def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
if args.model is None or args.model_class_name is None:
raise ValueError("Omni backend requires --model and --model-class-name")
# vLLM imports can otherwise change the process-wide Gemma dispatch.
torch.backends.cuda.enable_cudnn_sdp(False)
_reset_accelerator_peak_memory_stats()
attention_config = {"default": {"backend": "CUDNN_ATTN"}}
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_extras import get_extra_body_params
from vllm_omni.platforms import current_omni_platform
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(request["seed"])
omni = Omni(
model=args.model,
model_class_name=args.model_class_name,
enforce_eager=True,
enable_layerwise_offload=args.enable_layerwise_offload,
diffusion_attention_config=attention_config,
parallel_config=DiffusionParallelConfig(),
)
try:
# The explicit public class argument selects the worker pipeline. The
# parent Omni proxy can only expose the checkpoint's base class, which
# is insufficient to distinguish LTX execution modes.
model_class_name = args.model_class_name
sampling_params = OmniDiffusionSamplingParams(
height=request["height"],
width=request["width"],
generator=generator,
guidance_scale=None,
num_inference_steps=request["num_inference_steps"],
num_frames=request["num_frames"],
fps=request["fps"],
frame_rate=float(request["fps"]),
output_type="np",
)
extra_args = {
key: value for key, value in request.items() if key.startswith("video_") or key.startswith("audio_")
}
for key in ("sigmas", "stage_1_sigmas", "stage_2_sigmas", "image_crf"):
if key in request:
extra_args[key] = request[key]
apply_declared_extra_args(sampling_params, get_extra_body_params(model_class_name), extra_args)
prompt: dict[str, Any] = {"prompt": request["prompt"]}
if "negative_prompt" in request:
prompt["negative_prompt"] = request["negative_prompt"]
image_path = request.get("image")
if image_path is not None:
with Image.open(str(image_path)) as source_image:
image = source_image.convert("RGB")
image.load()
prompt["multi_modal_data"] = {"image": image}
output = omni.generate(prompt, sampling_params)
video, audio, audio_sample_rate = _unwrap_omni_output(output)
audio_tensor = torch.as_tensor(np.asarray(audio) if not isinstance(audio, torch.Tensor) else audio)
_save_outputs(
args.output_dir,
video=_canonical_video(video),
audio=audio_tensor,
audio_sample_rate=audio_sample_rate,
fps=float(request["fps"]),
metadata={
"backend": "omni",
"attention_backend": "torch_sdpa_cudnn_bf16",
"seed": request["seed"],
"sigmas": request.get("sigmas"),
"model": args.model,
"stage_1_sigmas": request.get("stage_1_sigmas"),
"stage_2_sigmas": request.get("stage_2_sigmas"),
"model_class_name": model_class_name,
"worker_peak_memory_mb": _omni_worker_peak_memory_mb(output),
**_accelerator_peak_memory_stats(),
},
)
finally:
omni.shutdown()
def main() -> None:
args = _parse_args()
request = json.loads(args.request.read_text())
if args.backend == "official":
_run_official(args, request)
else:
_run_omni(args, request)
if __name__ == "__main__":
main()
@@ -0,0 +1,753 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""E2E accuracy guards against pinned official Lightricks LTX references.
The LTX-2/2.3 comparison runs both runtimes through PyTorch SDPA and uses
``max_batch_size=4`` in the official reference to match Omni's fused guidance
batch. Video and audio guidance use the official non-HQ one-stage defaults;
only the generation shape and step count are reduced for CI runtime.
LTX-2.5 runs the official split-artifact pipelines with connector weights from
the same Diffusers checkpoint under test. The distilled case covers the fixed
two-stage schedule; the Full/SFT case covers raw dev weights and an explicit
shared one-stage schedule.
"""
from __future__ import annotations
import json
import math
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path
import numpy as np
import pytest
import torch
from huggingface_hub import hf_hub_download, snapshot_download
from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
from tests.e2e.accuracy.helpers import reset_artifact_dir
from tests.helpers import skip_if_gated_repo_inaccessible
from tests.helpers.mark import hardware_test
OFFICIAL_REPOSITORY = "https://github.com/Lightricks/LTX-2.git"
LTX25_OFFICIAL_REVISION = "7954dcb0d986bdc36ef272564a9789ade07fcc65"
LTX25_OMNI_MODEL_ID = "Lightricks/LTX-2.5-Diffusers"
LTX25_OMNI_MODEL_REVISION = "a6de4b5354f078db24d9cf4778c14846788aea3d"
LTX25_OFFICIAL_MODEL_ID = "Lightricks/LTX-2.5"
LTX25_OFFICIAL_MODEL_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1"
LTX25_OFFICIAL_COMMON_FILES = {
"text_encoder": "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors",
"video_vae": "vae/ltx-2.5-video-vae-conv-bf16.safetensors",
"audio_vae": "vae/ltx-2.5-audio-vae-bf16.safetensors",
}
LTX25_OFFICIAL_VARIANT_FILES = {
"distilled": {
"transformer": "diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors",
"spatial_upsampler": "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors",
},
"full": {
"transformer": "diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors",
},
"full_two_stage": {
"transformer": "diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors",
"spatial_upsampler": "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors",
"distilled_lora": "loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors",
},
}
# Version selected by the pinned official source's uv.lock. Keep it isolated
# from Omni's runtime and development dependencies.
OFFICIAL_OPENIMAGEIO_VERSION = "3.1.11.0"
LTX25_PROMPT = "A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside."
LTX25_I2V_IMAGE_REPO = "huggingface/documentation-images"
LTX25_I2V_IMAGE_FILENAME = "diffusers/svd/rocket.png"
LTX25_I2V_IMAGE_REVISION = "645d8364f0c7a101180b364811b5a11a362e4010"
NEGATIVE_PROMPT = (
"blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, "
"grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, "
"deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, "
"wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of field, "
"background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent lighting "
"direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny valley effect, "
"incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, mismatched lip sync, "
"silent or muted audio, distorted voice, robotic voice, echo, background noise, off-sync audio, incorrect "
"dialogue, added dialogue, repetitive speech, jittery movement, awkward pauses, incorrect timing, unnatural "
"transitions, inconsistent framing, tilted camera, flat lighting, inconsistent tone, cinematic oversaturation, "
"stylized filters, or AI artifacts."
)
# Release parity exercises the accelerated BF16 cuDNN path in both runtimes.
# Degenerate one-token attention falls back to PyTorch SDPA MATH.
ATTENTION_BACKEND = "torch_sdpa_cudnn_bf16"
VIDEO_SSIM_MEAN_THRESHOLD = 0.95
VIDEO_SSIM_MIN_THRESHOLD = 0.90
VIDEO_PSNR_MEAN_THRESHOLD = 30.0
AUDIO_RELATIVE_L2_THRESHOLD = 0.2
AUDIO_COSINE_THRESHOLD = 0.95
# Decoded-output gates for the pinned official/Omni BF16 cuDNN comparison.
# The metrics artifact preserves the exact observed values.
LTX25_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
LTX25_I2V_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
LTX25_VIDEO_SSIM_MIN_THRESHOLD = 0.99
LTX25_VIDEO_PSNR_MEAN_THRESHOLD = 40.0
LTX25_VIDEO_LPIPS_MEAN_THRESHOLD = 0.01
LTX25_AUDIO_RELATIVE_L2_THRESHOLD = 0.32
LTX25_AUDIO_COSINE_THRESHOLD = 0.95
# Full/SFT uses independently converted dev weights, but the release contract
# still requires decoded-video parity at the same 0.99 SSIM floor as the
# distilled paths when both runtimes use BF16 cuDNN attention.
LTX25_FULL_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
LTX25_FULL_VIDEO_SSIM_MIN_THRESHOLD = 0.99
LTX25_FULL_VIDEO_PSNR_MEAN_THRESHOLD = 40.0
LTX25_FULL_AUDIO_RELATIVE_L2_THRESHOLD = 0.3
LTX25_FULL_VIDEO_LPIPS_MEAN_THRESHOLD = 0.01
LTX25_FULL_I2V_VIDEO_LPIPS_MEAN_THRESHOLD = 0.02
LTX25_FULL_AUDIO_COSINE_THRESHOLD = AUDIO_COSINE_THRESHOLD
LTX25_STAGE_2_SIGMAS = [0.909375, 0.725, 0.421875, 0.0]
def test_ltx_reference_runner_unwraps_flattened_pipeline_output() -> None:
from vllm_omni.outputs import OmniRequestOutput
from .run_ltx25_reference import _omni_worker_peak_memory_mb, _unwrap_omni_output
frame = object()
audio = object()
output = OmniRequestOutput(
stage_id=0,
images=[frame],
_multimodal_output={"audio": audio, "audio_sample_rate": 48_000},
)
frames, actual_audio, sample_rate = _unwrap_omni_output(output)
assert frames is frame
assert actual_audio is audio
assert sample_rate == 48_000
output.peak_memory_mb = 12_345.0
assert _omni_worker_peak_memory_mb([output]) == 12_345.0
def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None:
start = time.perf_counter()
subprocess.run(command, env=env, timeout=timeout, check=True)
print(f"{' '.join(command[:3])} finished in {time.perf_counter() - start:.1f}s")
def _clone_pinned_source(root: Path, repository: str, revision: str) -> None:
root.mkdir(parents=True)
subprocess.run(["git", "init", "-q", str(root)], check=True)
subprocess.run(["git", "-C", str(root), "remote", "add", "origin", repository], check=True)
last_error: subprocess.CalledProcessError | None = None
for attempt in range(3):
try:
subprocess.run(
["git", "-C", str(root), "fetch", "--depth", "1", "origin", revision],
check=True,
)
last_error = None
break
except subprocess.CalledProcessError as error:
last_error = error
if attempt < 2:
time.sleep(5 * (attempt + 1))
if last_error is not None:
raise last_error
subprocess.run(["git", "-C", str(root), "checkout", "-q", "--detach", "FETCH_HEAD"], check=True)
def _git_revision(root: Path) -> str | None:
try:
return subprocess.check_output(
["git", "-C", str(root), "rev-parse", "HEAD"],
stderr=subprocess.DEVNULL,
text=True,
).strip()
except subprocess.CalledProcessError:
return None
def _ltx25_official_source(artifact_root: Path) -> tuple[Path, str]:
repository = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_REPOSITORY", OFFICIAL_REPOSITORY)
revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_REVISION", LTX25_OFFICIAL_REVISION)
configured_root = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_ROOT")
root = Path(configured_root) if configured_root else artifact_root / f"official-source-{revision[:12]}"
actual_revision = _git_revision(root) if root.exists() else None
if actual_revision != revision and configured_root:
raise AssertionError(f"Official LTX-2.5 source revision mismatch: {actual_revision} != {revision}")
if actual_revision != revision:
if root.exists():
shutil.rmtree(root)
_clone_pinned_source(root, repository, revision)
actual_revision = _git_revision(root)
assert actual_revision == revision, f"Official LTX-2.5 source revision mismatch: {actual_revision} != {revision}"
return root, revision
def _official_runner_prefix() -> list[str]:
"""Run the reference with its missing binary dependency isolated from CI."""
uv = shutil.which("uv")
assert uv is not None, "uv is required to run the pinned official LTX reference"
return [
uv,
"run",
"--no-project",
"--with",
f"openimageio=={OFFICIAL_OPENIMAGEIO_VERSION}",
"--python",
sys.executable,
"python",
]
def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") -> tuple[Path, str, str | None]:
if transformer_subfolder not in {"transformer", "transformer_full"}:
raise ValueError(f"Unsupported LTX-2.5 transformer subfolder: {transformer_subfolder!r}")
configured_model = os.environ.get("VLLM_TEST_LTX25_MODEL")
configured_revision = os.environ.get("VLLM_TEST_LTX25_MODEL_REVISION")
if configured_model and Path(configured_model).exists():
model = Path(configured_model).resolve()
snapshot_revision = model.name if model.parent.name == "snapshots" else None
assert (model / transformer_subfolder).is_dir(), f"LTX-2.5 component not found: {model / transformer_subfolder}"
return model, configured_model, configured_revision or snapshot_revision
model_id = configured_model or LTX25_OMNI_MODEL_ID
revision = configured_revision
if revision is None and model_id == LTX25_OMNI_MODEL_ID:
revision = LTX25_OMNI_MODEL_REVISION
model = Path(
snapshot_download(
repo_id=model_id,
revision=revision,
allow_patterns=[
"model_index.json",
"audio_vae/*",
"connectors/config.json",
"connectors/diffusion_pytorch_model.safetensors.index.json",
"connectors/diffusion_pytorch_model-*.safetensors",
"scheduler/*",
"text_encoder/config.json",
"text_encoder/generation_config.json",
"text_encoder/model*",
"tokenizer/*",
f"{transformer_subfolder}/*",
"vae/*",
"vocoder/*",
],
)
).resolve()
snapshot_revision = model.name if model.parent.name == "snapshots" else None
return model, model_id, snapshot_revision or revision
def _resolve_ltx25_official_artifacts(
*, checkpoint_variant: str = "distilled"
) -> tuple[dict[str, Path], str, str | None]:
try:
variant_files = LTX25_OFFICIAL_VARIANT_FILES[checkpoint_variant]
except KeyError as exc:
raise ValueError(f"Unsupported official LTX-2.5 checkpoint variant: {checkpoint_variant!r}") from exc
files = {**LTX25_OFFICIAL_COMMON_FILES, **variant_files}
configured_model = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL")
configured_revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL_REVISION")
if configured_model and Path(configured_model).exists():
root = Path(configured_model).resolve()
model_source = configured_model
snapshot_revision = root.name if root.parent.name == "snapshots" else None
resolved_revision = configured_revision or snapshot_revision
else:
model_source = configured_model or LTX25_OFFICIAL_MODEL_ID
revision = configured_revision
if revision is None and model_source == LTX25_OFFICIAL_MODEL_ID:
revision = LTX25_OFFICIAL_MODEL_REVISION
root = Path(
snapshot_download(
repo_id=model_source,
revision=revision,
allow_patterns=list(files.values()),
)
).resolve()
snapshot_revision = root.name if root.parent.name == "snapshots" else None
resolved_revision = snapshot_revision or revision
artifacts = {name: root / relative_path for name, relative_path in files.items()}
missing = [f"{name}: {path}" for name, path in artifacts.items() if not path.is_file()]
assert not missing, f"Official LTX-2.5 artifacts are missing: {', '.join(missing)}"
return artifacts, model_source, resolved_revision
def _resolve_gemma_root(model: Path) -> Path:
configured_root = os.environ.get("VLLM_TEST_LTX_GEMMA_ROOT")
if configured_root:
root = Path(configured_root)
assert root.is_dir(), f"Gemma root not found: {root}"
return root
return model
def _resolve_ltx25_image() -> Path:
configured_image = os.environ.get("VLLM_TEST_LTX25_I2V_IMAGE")
if configured_image:
image = Path(configured_image)
assert image.is_file(), f"LTX-2.5 I2V conditioning image not found: {image}"
return image
return Path(
hf_hub_download(
repo_id=LTX25_I2V_IMAGE_REPO,
repo_type="dataset",
filename=LTX25_I2V_IMAGE_FILENAME,
revision=LTX25_I2V_IMAGE_REVISION,
)
)
def _require_ltx25_model_access() -> None:
configured_omni_model = os.environ.get("VLLM_TEST_LTX25_MODEL")
if not (configured_omni_model and Path(configured_omni_model).exists()):
omni_model = configured_omni_model or LTX25_OMNI_MODEL_ID
omni_revision = os.environ.get("VLLM_TEST_LTX25_MODEL_REVISION")
if omni_revision is None and omni_model == LTX25_OMNI_MODEL_ID:
omni_revision = LTX25_OMNI_MODEL_REVISION
skip_if_gated_repo_inaccessible(
omni_model,
revision=omni_revision,
filename="model_index.json",
)
configured_official_model = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL")
if not (configured_official_model and Path(configured_official_model).exists()):
official_model = configured_official_model or LTX25_OFFICIAL_MODEL_ID
official_revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL_REVISION")
if official_revision is None and official_model == LTX25_OFFICIAL_MODEL_ID:
official_revision = LTX25_OFFICIAL_MODEL_REVISION
skip_if_gated_repo_inaccessible(
official_model,
revision=official_revision,
filename=LTX25_OFFICIAL_COMMON_FILES["text_encoder"],
)
def _ltx25_request(image: Path | None = None) -> dict[str, object]:
request: dict[str, object] = {
"prompt": LTX25_PROMPT,
# Keep the nightly parity test small enough to run the pinned official
# pipeline and Omni sequentially on one H100. The public recipe covers
# the release-resolution 1920x1088x121 configuration.
"width": 1024,
"height": 576,
"num_frames": 25,
"fps": 24,
"num_inference_steps": 8,
"seed": 42,
}
if image is not None:
request.update(image=str(image.resolve()), image_crf=18)
return request
def _ltx25_full_sigmas(num_inference_steps: int) -> list[float]:
"""Materialize the official LTX-2.5 one-stage schedule in FP32."""
sigmas = torch.linspace(1.0, 0.0, num_inference_steps + 1)
base_anchor = 1024
max_anchor = 4096
base_shift = 0.95
max_shift = 2.05
slope = (max_shift - base_shift) / (max_anchor - base_anchor)
sigma_shift = max_anchor * slope + (base_shift - slope * base_anchor)
exp_shift = math.exp(sigma_shift)
sigmas = torch.where(sigmas != 0, exp_shift / (exp_shift + (1 / sigmas - 1)), 0)
non_zero = sigmas != 0
one_minus_sigmas = 1.0 - sigmas[non_zero]
scale = one_minus_sigmas[-1] / (1.0 - 0.1)
sigmas[non_zero] = 1.0 - one_minus_sigmas / scale
return [float(sigma) for sigma in sigmas]
def _ltx25_full_request(image: Path | None = None) -> dict[str, object]:
num_inference_steps = 30
request: dict[str, object] = {
"prompt": LTX25_PROMPT,
"negative_prompt": NEGATIVE_PROMPT,
# Full guidance expands to a four-way batch. Keep the decoded output
# small enough for the pinned official and Omni runs to share one H100.
"width": 512,
"height": 384,
"num_frames": 25,
"fps": 24,
"num_inference_steps": num_inference_steps,
"seed": 42,
"video_cfg_scale": 3.0,
"audio_cfg_scale": 7.0,
"video_stg_scale": 1.0,
"audio_stg_scale": 1.0,
"video_modality_scale": 3.0,
"audio_modality_scale": 3.0,
"video_rescale_scale": 0.7,
"audio_rescale_scale": 0.7,
"video_stg_blocks": [28],
"audio_stg_blocks": [28],
"sigmas": _ltx25_full_sigmas(num_inference_steps),
}
if image is not None:
request.update(image=str(image.resolve()), image_crf=18)
return request
def test_ltx25_full_request_pins_official_schedule() -> None:
request = _ltx25_full_request()
sigmas = request["sigmas"]
assert isinstance(sigmas, list)
assert isinstance(request["num_inference_steps"], int)
assert len(sigmas) == request["num_inference_steps"] + 1
assert math.isclose(sigmas[0], 1.0, rel_tol=0.0, abs_tol=1e-6)
assert sigmas[-1] == 0.0
assert all(sigmas[index] > sigmas[index + 1] for index in range(len(sigmas) - 1))
def test_ltx25_i2v_requests_pin_official_crf() -> None:
image = Path("conditioning.png")
for request in (_ltx25_request(image), _ltx25_full_request(image)):
assert request["image"] == str(image.resolve())
assert request["image_crf"] == 18
def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
from benchmarks.diffusion.quantization_quality import compute_lpips_video
assert reference.shape == prediction.shape
assert reference.ndim == 4 and reference.shape[-1] == 3
ssim_scores: list[float] = []
psnr_scores: list[float] = []
for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0)
prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0)
ssim_scores.append(float(StructuralSimilarityIndexMeasure(data_range=1.0)(prediction_tensor, reference_tensor)))
psnr_scores.append(float(PeakSignalNoiseRatio(data_range=1.0)(prediction_tensor, reference_tensor)))
difference = np.abs(reference.astype(np.float64) - prediction.astype(np.float64))
return {
"ssim_mean": float(np.mean(ssim_scores)),
"ssim_min": float(np.min(ssim_scores)),
"psnr_mean_db": float(np.mean(psnr_scores)),
"lpips_alex_mean": compute_lpips_video(reference, prediction),
"max_abs": float(difference.max()),
"mean_abs": float(difference.mean()),
}
def _canonical_audio(audio: np.ndarray) -> np.ndarray:
while audio.ndim > 2 and audio.shape[0] == 1:
audio = audio[0]
return audio.astype(np.float64)
def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float | bool]:
reference = _canonical_audio(reference)
prediction = _canonical_audio(prediction)
assert reference.shape == prediction.shape
difference = reference - prediction
reference_norm = max(float(np.linalg.norm(reference)), 1e-12)
prediction_norm = max(float(np.linalg.norm(prediction)), 1e-12)
return {
"bitwise_equal": bool(np.array_equal(reference, prediction)),
"max_abs": float(np.abs(difference).max()),
"mean_abs": float(np.abs(difference).mean()),
"relative_l2": float(np.linalg.norm(difference) / reference_norm),
"cosine_similarity": float(np.vdot(reference.ravel(), prediction.ravel()) / (reference_norm * prediction_norm)),
}
@pytest.mark.full_model
@pytest.mark.benchmark
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@pytest.mark.parametrize("task", ["t2v", "i2v"])
def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path, task: str) -> None:
"""Compare Omni with the official runtime using the same checkpoint connector weights."""
_require_ltx25_model_access()
artifact_parent = accuracy_artifact_root / "ltx_official"
output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_distilled_{task}")
official_root, official_revision = _ltx25_official_source(artifact_parent)
official_artifacts, official_model_source, official_model_revision = _resolve_ltx25_official_artifacts()
omni_model, omni_model_source, omni_model_revision = _resolve_ltx25_omni_model()
image = _resolve_ltx25_image() if task == "i2v" else None
request_path = output_root / "request.json"
request_path.write_text(json.dumps(_ltx25_request(image), indent=2) + "\n")
runner = Path(__file__).with_name("run_ltx25_reference.py")
runner_args = [
str(runner),
"--request",
str(request_path),
"--enable-layerwise-offload",
]
env = os.environ.copy()
env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision
env["PYTHONUNBUFFERED"] = "1"
repository_root = Path(__file__).resolve().parents[4]
existing_pythonpath = env.get("PYTHONPATH")
env["PYTHONPATH"] = (
str(repository_root) if not existing_pythonpath else f"{repository_root}{os.pathsep}{existing_pythonpath}"
)
official_output = output_root / "official"
_run(
_official_runner_prefix()
+ runner_args
+ [
"--backend",
"official",
"--output-dir",
str(official_output),
"--official-root",
str(official_root),
"--transformer-path",
str(official_artifacts["transformer"]),
"--text-encoder-path",
str(official_artifacts["text_encoder"]),
"--video-vae-path",
str(official_artifacts["video_vae"]),
"--audio-vae-path",
str(official_artifacts["audio_vae"]),
"--spatial-upsampler-path",
str(official_artifacts["spatial_upsampler"]),
"--connector-model",
str(omni_model),
],
env=env,
)
omni_output = output_root / "omni"
_run(
[sys.executable]
+ runner_args
+ [
"--backend",
"omni",
"--output-dir",
str(omni_output),
"--model",
str(omni_model),
"--model-class-name",
"LTX2DistilledTwoStagePipeline",
],
env=env,
)
official_metadata = json.loads((official_output / "metadata.json").read_text())
omni_metadata = json.loads((omni_output / "metadata.json").read_text())
assert official_metadata["official_revision"] == official_revision
assert official_metadata["pipeline"] == "DistilledPipeline"
assert official_metadata["attention_backend"] == ATTENTION_BACKEND
assert official_metadata["connector_model"] == str(omni_model)
assert omni_metadata["model_class_name"] == "LTX2DistilledTwoStagePipeline"
assert omni_metadata["attention_backend"] == ATTENTION_BACKEND
assert official_metadata["audio_sample_rate"] == omni_metadata["audio_sample_rate"]
video_metrics = _video_metrics(
np.load(official_output / "video.npy"),
np.load(omni_output / "video.npy"),
)
audio_metrics = _audio_metrics(
np.load(official_output / "audio.npy"),
np.load(omni_output / "audio.npy"),
)
result = {
"case": "ltx2_5_distilled",
"task": task,
"attention_backend": ATTENTION_BACKEND,
"official_revision": official_revision,
"official_model": official_model_source,
"official_model_revision": official_model_revision,
"official_connector_model": omni_model_source,
"official_connector_revision": omni_model_revision,
"omni_model": omni_model_source,
"omni_model_revision": omni_model_revision,
"resolved_omni_model_path": str(omni_model),
"peak_memory_mb": {
"official_allocated": official_metadata["peak_memory_allocated_mb"],
"official_reserved": official_metadata["peak_memory_reserved_mb"],
"omni_worker_reported": omni_metadata["worker_peak_memory_mb"],
"omni_parent_allocated": omni_metadata["peak_memory_allocated_mb"],
"omni_parent_reserved": omni_metadata["peak_memory_reserved_mb"],
},
"video": video_metrics,
"audio": audio_metrics,
}
(output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result, indent=2))
ssim_mean_threshold = LTX25_VIDEO_SSIM_MEAN_THRESHOLD if task == "t2v" else LTX25_I2V_VIDEO_SSIM_MEAN_THRESHOLD
assert video_metrics["ssim_mean"] >= ssim_mean_threshold
assert video_metrics["ssim_min"] >= LTX25_VIDEO_SSIM_MIN_THRESHOLD
assert video_metrics["psnr_mean_db"] >= LTX25_VIDEO_PSNR_MEAN_THRESHOLD
assert video_metrics["lpips_alex_mean"] <= LTX25_VIDEO_LPIPS_MEAN_THRESHOLD
assert audio_metrics["cosine_similarity"] >= LTX25_AUDIO_COSINE_THRESHOLD
assert audio_metrics["relative_l2"] <= LTX25_AUDIO_RELATIVE_L2_THRESHOLD
@pytest.mark.full_model
@pytest.mark.benchmark
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@pytest.mark.parametrize("task", ["t2v", "i2v"])
@pytest.mark.parametrize("pipeline_mode", ["full_one_stage", "full_two_stage"])
def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pipeline_mode: str) -> None:
"""Compare raw official Full/SFT weights with Omni's converted Full pipeline."""
_require_ltx25_model_access()
artifact_parent = accuracy_artifact_root / "ltx_official"
output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_{pipeline_mode}_{task}")
model_class_name = {
"full_one_stage": "LTX2Pipeline",
"full_two_stage": "LTX2TwoStagePipeline",
}[pipeline_mode]
official_root, official_revision = _ltx25_official_source(artifact_parent)
official_artifacts, official_model_source, official_model_revision = _resolve_ltx25_official_artifacts(
checkpoint_variant="full" if pipeline_mode == "full_one_stage" else "full_two_stage"
)
omni_model, omni_model_source, omni_model_revision = _resolve_ltx25_omni_model(
transformer_subfolder="transformer_full"
)
image = _resolve_ltx25_image() if task == "i2v" else None
request = _ltx25_full_request(image)
if pipeline_mode == "full_two_stage":
request["stage_1_sigmas"] = request.pop("sigmas")
request["stage_2_sigmas"] = LTX25_STAGE_2_SIGMAS
request_path = output_root / "request.json"
request_path.write_text(json.dumps(request, indent=2) + "\n")
runner = Path(__file__).with_name("run_ltx25_reference.py")
runner_args = [
str(runner),
"--request",
str(request_path),
"--enable-layerwise-offload",
]
env = os.environ.copy()
env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision
env["PYTHONUNBUFFERED"] = "1"
repository_root = Path(__file__).resolve().parents[4]
existing_pythonpath = env.get("PYTHONPATH")
env["PYTHONPATH"] = (
str(repository_root) if not existing_pythonpath else f"{repository_root}{os.pathsep}{existing_pythonpath}"
)
official_output = output_root / "official"
official_sidecar_args: list[str] = []
if pipeline_mode == "full_two_stage":
official_sidecar_args = [
"--spatial-upsampler-path",
str(official_artifacts["spatial_upsampler"]),
"--distilled-lora-path",
str(official_artifacts["distilled_lora"]),
]
_run(
_official_runner_prefix()
+ runner_args
+ [
"--backend",
"official",
"--output-dir",
str(official_output),
"--official-root",
str(official_root),
"--official-pipeline",
pipeline_mode,
"--transformer-path",
str(official_artifacts["transformer"]),
"--text-encoder-path",
str(official_artifacts["text_encoder"]),
"--video-vae-path",
str(official_artifacts["video_vae"]),
"--audio-vae-path",
str(official_artifacts["audio_vae"]),
"--connector-model",
str(omni_model),
]
+ official_sidecar_args,
env=env,
)
omni_output = output_root / "omni"
_run(
[sys.executable]
+ runner_args
+ [
"--backend",
"omni",
"--output-dir",
str(omni_output),
"--model",
str(omni_model),
"--model-class-name",
model_class_name,
],
env=env,
)
official_metadata = json.loads((official_output / "metadata.json").read_text())
omni_metadata = json.loads((omni_output / "metadata.json").read_text())
assert official_metadata["official_revision"] == official_revision
assert (
official_metadata["pipeline"]
== {
"full_one_stage": "TI2VidOneStagePipeline",
"full_two_stage": "TI2VidTwoStagesPipeline",
}[pipeline_mode]
)
assert official_metadata["attention_backend"] == ATTENTION_BACKEND
assert official_metadata["connector_model"] == str(omni_model)
assert omni_metadata["model_class_name"] == model_class_name
assert omni_metadata["attention_backend"] == ATTENTION_BACKEND
assert official_metadata["seed"] == omni_metadata["seed"]
for schedule_name in ("sigmas", "stage_1_sigmas", "stage_2_sigmas"):
assert official_metadata[schedule_name] == omni_metadata[schedule_name]
assert official_metadata["audio_sample_rate"] == omni_metadata["audio_sample_rate"]
video_metrics = _video_metrics(
np.load(official_output / "video.npy"),
np.load(omni_output / "video.npy"),
)
audio_metrics = _audio_metrics(
np.load(official_output / "audio.npy"),
np.load(omni_output / "audio.npy"),
)
result = {
"case": f"ltx2_5_{pipeline_mode}",
"task": task,
"attention_backend": ATTENTION_BACKEND,
"official_revision": official_revision,
"official_model": official_model_source,
"official_model_revision": official_model_revision,
"official_connector_model": omni_model_source,
"official_connector_revision": omni_model_revision,
"omni_model": omni_model_source,
"omni_model_revision": omni_model_revision,
"resolved_omni_model_path": str(omni_model),
"peak_memory_mb": {
"official_allocated": official_metadata["peak_memory_allocated_mb"],
"official_reserved": official_metadata["peak_memory_reserved_mb"],
"omni_worker_reported": omni_metadata["worker_peak_memory_mb"],
"omni_parent_allocated": omni_metadata["peak_memory_allocated_mb"],
"omni_parent_reserved": omni_metadata["peak_memory_reserved_mb"],
},
"video": video_metrics,
"audio": audio_metrics,
}
(output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result, indent=2))
assert video_metrics["ssim_mean"] >= LTX25_FULL_VIDEO_SSIM_MEAN_THRESHOLD
assert video_metrics["ssim_min"] >= LTX25_FULL_VIDEO_SSIM_MIN_THRESHOLD
assert video_metrics["psnr_mean_db"] >= LTX25_FULL_VIDEO_PSNR_MEAN_THRESHOLD
lpips_threshold = (
LTX25_FULL_VIDEO_LPIPS_MEAN_THRESHOLD if task == "t2v" else LTX25_FULL_I2V_VIDEO_LPIPS_MEAN_THRESHOLD
)
assert audio_metrics["relative_l2"] <= LTX25_FULL_AUDIO_RELATIVE_L2_THRESHOLD
assert video_metrics["lpips_alex_mean"] <= lpips_threshold
assert audio_metrics["cosine_similarity"] >= LTX25_FULL_AUDIO_COSINE_THRESHOLD
+7 -7
View File
@@ -29,14 +29,14 @@ HEIGHT = 512
NUM_INFERENCE_STEPS = 20
TRUE_CFG_SCALE = 4.0
SEED = 42
SSIM_THRESHOLD = 0.97
# The FA3-hub-matched comparison clears 0.97 with margin. The SDPA-matched
# fallback (hub kernel unavailable for the image's torch) landed at 0.9687 in
# build 2954 — the attention math matches but the surrounding stacks (vLLM
# compiled path vs eager diffusers) still differ slightly. Keep the strict
# gate for the FA3 path and a slightly looser one for the fallback only.
SSIM_THRESHOLD = 0.94
# With --fa-deterministic, H100 nightlies stably land near SSIM 0.958 / PSNR 27.8
# vs Diffusers (see #5734 / #5963), so the primary gates are 0.94 / 27.0.
# The SDPA-matched fallback (hub kernel unavailable for the image's torch)
# landed at SSIM 0.9687 in build 2954 — keep a slightly tighter SSIM gate for
# that path only.
SSIM_THRESHOLD_SDPA_FALLBACK = 0.96
PSNR_THRESHOLD = 30.0
PSNR_THRESHOLD = 27.0
MODEL_2512_ID = "Qwen/Qwen-Image-2512"
MODEL_2512_ENV_VAR = "QWEN_IMAGE_2512_MODEL"
+1 -1
View File
@@ -46,7 +46,7 @@ REFERENCE_TEXT_IMG2TEXT = (
_OMNI_RUNNER_PARAM = (MODEL_NAME, STAGE_CONFIG)
pytestmark = [
pytest.mark.advanced_model,
pytest.mark.slow,
pytest.mark.diffusion,
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
]
@@ -72,7 +72,7 @@ def _send_inpaint_with_generator(
)
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_flux2_klein_inpaint_deterministic(omni_runner_handler: OmniRunnerHandler):
@@ -93,7 +93,7 @@ def test_flux2_klein_inpaint_deterministic(omni_runner_handler: OmniRunnerHandle
)
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_flux2_klein_inpaint_different_seeds_different_output(omni_runner_handler: OmniRunnerHandler):
@@ -0,0 +1,345 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""E2E offline inference tests for the Gepard-1.0 single-stage native-AR TTS.
Zero-shot only (the model's default learned voice). These tests need a GPU and
the NeMo NanoCodec, so the whole file runs in the weekly tier -- see the comment
above ``pytestmark``. The window arithmetic behind the streaming decode is
covered on CPU in ``tests/model_executor/models/test_gepard_window.py``.
"""
from __future__ import annotations
import functools
import tempfile
from pathlib import Path
from typing import NamedTuple
import pytest
import soundfile as sf
import torch
from transformers import AutoTokenizer
from tests.helpers.mark import hardware_test
from tests.helpers.media import convert_audio_file_to_text
from tests.helpers.runtime import OmniRunner # noqa: F401 (indirect fixture)
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni import Omni
from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
# The codec decoder is an optional dependency; skip rather than fail the engine
# inside the worker process when it is absent. The CI image does not carry it,
# so the weekly step installs it first -- see .buildkite/cuda/test-weekly.yml.
# Install recipe: examples/offline_inference/text_to_speech/README.md.
pytest.importorskip("nemo.collections.tts.models")
MODEL_NAME = "nineninesix/gepard-1.0"
STAGE_CONFIG = get_deploy_config_path("gepard.yaml")
SAMPLE_RATE = 22050
# NanoCodec upsamples one FSQ frame to exactly this many samples.
SAMPLES_PER_FRAME = 1024
@functools.cache
def _config() -> GepardConfig:
"""The checkpoint's own layout, not the class defaults.
The worker builds its config from the ``gepard_config.json`` sidecar, so a
prompt assembled from the defaults would silently use a different layout
the day the sidecar moves — and a mismatched layout degrades the speech
rather than failing, which is the hardest way for this to be found.
"""
return GepardConfig.from_checkpoint(MODEL_NAME)
# Prompt -> a word that appears in no other prompt, so a clip can be traced
# back to the request that asked for it.
#
# These are not interchangeable. Batching perturbs every request's sampling
# trajectory, and on some trajectories this checkpoint's stop head never fires;
# this set was screened over 10 batched runs with no runaway. Re-screen the
# whole set when changing one, not just the row changed.
_DISTINGUISHABLE_PROMPTS = {
"He drinks coffee every morning.": "coffee",
"Machine learning is interesting.": "learning",
"Please close the window before leaving.": "window",
"My favorite color is purple.": "purple",
}
# Above text_repetition.apply_below, so the layout is voiced once instead of
# repeated. Every other prompt here takes the repeated branch.
_LONG_TEXT = (
"The morning train was late again, so I walked the last two miles along the "
"river and reached the office just before nine."
)
_OMNI_RUNNER_PARAM = (
MODEL_NAME,
STAGE_CONFIG,
{"trust_remote_code": True},
)
# Only a common word is checked, which needs less ASR than the shared helper's
# ``small`` default; a miss escalates to it.
_ASR_MODEL = "base"
_ASR_ESCALATION_MODEL = "small"
# Transcription dominates this file's runtime -- Whisper runs on CPU on a
# single-GPU runner -- so only two of the four concurrent clips are read as
# speech. The pairwise waveform comparison covers all four.
_TRANSCRIBED_CLIPS = 2
# No level marker at module scope: each test carries the one matching how often
# its claim needs re-checking, though the weekly step collects the file whole.
# There is deliberately no core_model row -- under dummy weights the stop head is
# random, so every request would run to gepard.yaml's max_tokens.
pytestmark = [
pytest.mark.tts,
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
]
def _build_request(text: str) -> dict:
"""Assemble the [speaker slots, SOT, text, EOT, SOS] prompt."""
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
prompt_token_ids = build_gepard_prompt_ids(
tokenizer(text, add_special_tokens=False)["input_ids"],
config=_config(),
)
return {
"prompt_token_ids": prompt_token_ids,
"additional_information": {"text": [text]},
}
def _extract_audio(mm) -> torch.Tensor | None:
# The consolidation path renames "model_outputs" to "audio".
audio = mm.get("audio")
if audio is None:
audio = mm.get("model_outputs")
if isinstance(audio, list):
audio = audio[0] if len(audio) else None
return audio
class _Clip(NamedTuple):
wav: torch.Tensor
# Counts the frames the codec was handed using only what the engine
# returned, never the model's internals.
committed_frames: int
def _synthesize_all(omni: Omni, texts: list[str]) -> list[_Clip]:
"""Submit every text in one call; return one clip per request, in
submission order.
``Omni.generate`` yields in completion order and never sorts, so the pairing
is restored from the request id prefix it assigns (``f"{i}_{uuid4()}"``).
Without that, the only testable claim is that the batch as a whole looks
plausible, which a stable swap between two requests satisfies.
Nothing here authors a SamplingParams: a fresh one would replace the stage
defaults rather than merge over them, dropping the pipeline's stop token.
"""
outputs = omni.generate([_build_request(t) for t in texts])
assert len(outputs) == len(texts), f"expected {len(texts)} outputs, got {len(outputs)}"
by_index: dict[int, _Clip] = {}
for stage_outputs in outputs:
index = int(str(stage_outputs.request_id).split("_", 1)[0])
waveform = None
frames = 0
for out in stage_outputs.outputs:
# Explicit None check: `x or []` would ask an array for a truth value.
token_ids = out.token_ids if out.token_ids is not None else []
# STOP sits just past head0's range, so a token equal to it is
# never a committed frame.
frames += sum(1 for t in token_ids if int(t) != _config().stop_token)
mm = out.multimodal_output
if mm is None:
continue
audio = _extract_audio(mm)
if audio is None:
continue
sr_t = mm.get("sr")
if hasattr(sr_t, "item"):
assert int(sr_t.item()) == SAMPLE_RATE
waveform = audio.cpu().float()
assert waveform is not None, "no audio output produced"
by_index[index] = _Clip(wav=waveform, committed_frames=frames)
assert sorted(by_index) == list(range(len(texts))), (
f"request ids did not cover every submitted prompt exactly once: got {sorted(by_index)}"
)
return [by_index[i] for i in range(len(texts))]
def _synthesize(omni: Omni, text: str) -> _Clip:
return _synthesize_all(omni, [text])[0]
def _assert_clip_is_sane(clip: _Clip) -> None:
"""The checks every clip has to pass, whatever produced it."""
wav = clip.wav
assert wav.dim() == 1, f"expected mono 1-D, got shape {tuple(wav.shape)}"
assert wav.numel() > 0, "empty waveform"
assert torch.isfinite(wav).all(), "non-finite samples"
# Frame conservation against the engine's own token count, not just
# divisibility. This is what catches a payload carrying non-audio data or a
# dropped tail; both pass every other check here.
assert wav.numel() == clip.committed_frames * SAMPLES_PER_FRAME, (
f"{wav.numel()} samples for {clip.committed_frames} committed frames — "
f"expected exactly {clip.committed_frames * SAMPLES_PER_FRAME}"
)
# Anything far outside the waveform range is not codec output.
assert float(wav.abs().max()) <= 1.5, f"peak {float(wav.abs().max()):.2f} is out of range"
def _assert_stopped_on_its_own(clip: _Clip) -> None:
"""The request must stop on its own rather than run to ``max_tokens``.
Separate from the structural checks because it reads the stop head's
judgement, not the payload's shape: under ``load_format: dummy`` that head
is random and never fires, so every clip runs to the token cap.
"""
seconds = clip.wav.numel() / SAMPLE_RATE
assert 0.5 < seconds < 30.0, f"implausible duration {seconds:.2f}s — did STOP fire?"
def _transcribe(wav: torch.Tensor, tmp_dir: str, name: str, model_size: str = _ASR_MODEL) -> str:
path = str(Path(tmp_dir) / f"{name}.wav")
sf.write(path, wav.numpy(), SAMPLE_RATE)
return convert_audio_file_to_text(path, model_size=model_size, language="en").lower()
def _transcribe_for(wav: torch.Tensor, tmp_dir: str, name: str, accepts) -> str:
"""Transcribe with the fast model, retrying once with the stronger one.
``accepts`` decides whether the fast transcript is usable. Escalating only
on a miss keeps the average clip at ``_ASR_MODEL`` while a genuine model
artifact still fails the caller's assertion, since the stronger ASR
mistranscribes it too.
"""
transcript = _transcribe(wav, tmp_dir, name)
if accepts(transcript):
return transcript
return _transcribe(wav, tmp_dir, name, model_size=_ASR_ESCALATION_MODEL)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize(
("text", "keyword"),
[
# The canonical path: one request, one transcript, the cheapest check
# that reads the audio as speech.
pytest.param(
"Hello, this is Gepard speaking.",
"hello",
marks=pytest.mark.advanced_model,
id="canonical",
),
# No keyword: this row carries the structural and stopping checks over a
# second text and leaves the transcript to the row above.
pytest.param(
"The quick brown fox jumps over the lazy dog.",
None,
marks=pytest.mark.full_model,
id="second_text",
),
],
)
def test_gepard_offline_zero_shot(omni_runner, run_level: str, text: str, keyword: str | None) -> None:
"""Zero-shot synthesis produces a finite mono waveform at 22.05 kHz.
The transcript check is the one assertion here that reads the audio as
speech rather than as a tensor. It stands in for a parity comparison, which
needs a reference implementation this PR cannot run — noise, a repeated
loop or the wrong clip all pass every numeric check above it.
"""
clip = _synthesize(omni_runner, text)
_assert_clip_is_sane(clip)
if run_level not in {"advanced_model", "full_model"}:
return # dummy weights: the audio is structural only, so do not read it
_assert_stopped_on_its_own(clip)
if keyword is None:
return
with tempfile.TemporaryDirectory() as tmp_dir:
transcript = _transcribe_for(clip.wav, tmp_dir, "zero_shot", lambda t: keyword in t)
assert keyword in transcript, f"expected {keyword!r} in transcript, got {transcript!r}"
@pytest.mark.full_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_gepard_offline_long_text_skips_repetition(omni_runner, run_level: str) -> None:
"""The other branch of the prompt layout, end to end.
A text at or above ``apply_below`` tokens is voiced once; shorter ones are
repeated to carry enough mass against the speaker prefix. The layout
assertion keeps this row honest — if a tokenizer change pushed this text
back under the threshold it would silently retest the repeated path.
"""
prompt_ids = _build_request(_LONG_TEXT)["prompt_token_ids"]
assert prompt_ids.count(_config().start_of_text) == 1, "_LONG_TEXT is no longer above apply_below"
clip = _synthesize(omni_runner, _LONG_TEXT)
_assert_clip_is_sane(clip)
if run_level not in {"advanced_model", "full_model"}:
return
_assert_stopped_on_its_own(clip)
with tempfile.TemporaryDirectory() as tmp_dir:
transcript = _transcribe_for(clip.wav, tmp_dir, "long_text", lambda t: "river" in t)
assert "river" in transcript, f"expected 'river' in transcript, got {transcript!r}"
@pytest.mark.slow
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_gepard_offline_concurrent_requests_stay_isolated(omni_runner, run_level: str) -> None:
"""Four requests in one batch, each holding its own generation state.
The prompts are deliberately unlike each other: cross-talk between the
per-request frame history, the codec window or the output routing shows up
as a transcript matching the wrong prompt, which similar prompts would
hide. ``max_num_seqs`` is 4, so these really do run together.
"""
texts = list(_DISTINGUISHABLE_PROMPTS)
clips = _synthesize_all(omni_runner, texts)
for clip in clips:
_assert_clip_is_sane(clip)
if run_level not in {"advanced_model", "full_model"}:
return # dummy weights: the audio is structural only, so do not read it
for clip in clips:
_assert_stopped_on_its_own(clip)
# All four clips, not just the transcribed ones: delivering one request's
# audio for another leaves two byte-identical waveforms.
for i in range(len(clips)):
for j in range(i + 1, len(clips)):
assert not torch.equal(clips[i].wav, clips[j].wav), (
f"clips {i} and {j} are the same waveform — one request's audio was delivered for another"
)
# Each clip must say what its own prompt asked for: accepting any of the
# four would pass a stable permutation, which is what a mix-up produces.
keywords = set(_DISTINGUISHABLE_PROMPTS.values())
with tempfile.TemporaryDirectory() as tmp_dir:
for i, clip in enumerate(clips[:_TRANSCRIBED_CLIPS]):
expected = _DISTINGUISHABLE_PROMPTS[texts[i]]
# Escalate on the expected word being absent, not on "no keyword
# at all" -- a neighbour's word is the failure under test.
transcript = _transcribe_for(clip.wav, tmp_dir, f"concurrent_{i}", lambda t: expected in t)
present = sorted(k for k in keywords if k in transcript)
assert present == [expected], (
f"clip for {texts[i]!r} transcribed as {transcript!r}: expected {expected!r} and "
f"no other prompt's keyword, found {present} — requests bled into each other"
)
@@ -39,8 +39,7 @@ def _extract_audio(multimodal_output: dict) -> torch.Tensor:
return audio
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_voxcpm2_zero_shot_001(omni_runner: OmniRunner) -> None:
@@ -53,8 +52,7 @@ def test_voxcpm2_zero_shot_001(omni_runner: OmniRunner) -> None:
assert 0.5 < duration_s < 30.0, f"Audio duration out of range: {duration_s:.2f}s"
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_voxcpm2_voice_clone_002(omni_runner: OmniRunner) -> None:
@@ -95,8 +93,7 @@ def test_voxcpm2_voice_clone_002(omni_runner: OmniRunner) -> None:
assert 0.5 < duration_s < 30.0, f"Audio duration out of range: {duration_s:.2f}s"
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_voxcpm2_prefill_decode_mixed_batch_003(omni_runner: OmniRunner) -> None:
+1 -1
View File
@@ -11,7 +11,7 @@ MODEL = "Tongyi-MAI/Z-Image-Turbo"
PROMPT = "A high-detail studio photo of an orange tabby cat sitting on a laptop keyboard."
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 1, "rocm": 1, "xpu": 2})
@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
+2 -4
View File
@@ -75,8 +75,7 @@ def _build_img2img_messages(prompt: str, image_b64: str) -> list[dict]:
]
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@@ -98,8 +97,7 @@ def test_bagel_text2img_online(omni_server, openai_client) -> None:
openai_client.send_diffusion_request(request_config)
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+12 -15
View File
@@ -59,8 +59,7 @@ _REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know w
class TestHiggsAudioV3OnlineHappyPath:
"""Plain-text -> audio happy paths against the live HTTP server."""
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_plain_text_wav(self, omni_server, openai_client) -> None:
@@ -79,7 +78,7 @@ class TestHiggsAudioV3OnlineHappyPath:
}
)
@pytest.mark.core_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> None:
@@ -97,8 +96,7 @@ class TestHiggsAudioV3OnlineHappyPath:
)
@pytest.mark.skip(reason="issue#4411")
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_concurrent_plain_text(self, omni_server, openai_client) -> None:
@@ -115,7 +113,7 @@ class TestHiggsAudioV3OnlineHappyPath:
)
@pytest.mark.skip(reason="issue#4411")
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None:
@@ -147,7 +145,7 @@ class TestHiggsAudioV3OnlineHappyPath:
)
@pytest.mark.skip(reason="issue#4411")
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_concurrent_pcm_streaming(self, omni_server, openai_client) -> None:
@@ -189,7 +187,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
They do not assert audio quality (out of scope for CI).
"""
@pytest.mark.core_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_inline_emotion_and_expressive(self, omni_server, openai_client) -> None:
@@ -205,7 +203,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
}
)
@pytest.mark.core_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_inline_style_whispering(self, omni_server, openai_client) -> None:
@@ -221,7 +219,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
}
)
@pytest.mark.core_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_inline_prosody_speed_and_pitch(self, omni_server, openai_client) -> None:
@@ -239,7 +237,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
}
)
@pytest.mark.core_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_inline_pause_mid_text(self, omni_server, openai_client) -> None:
@@ -255,7 +253,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
}
)
@pytest.mark.core_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_inline_sfx_with_onomatopoeia(self, omni_server, openai_client) -> None:
@@ -279,8 +277,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
class TestHiggsAudioV3OnlineVoiceClone:
"""Voice clone via the two payload shapes the model serves."""
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_voice_clone_ref_audio_ref_text(self, omni_server, openai_client) -> None:
@@ -297,7 +294,7 @@ class TestHiggsAudioV3OnlineVoiceClone:
}
)
@pytest.mark.core_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_voice_clone_references_alias(self, omni_server, openai_client) -> None:
+89
View File
@@ -0,0 +1,89 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Online serving smokes for all four canonical LTX-2.5 pipelines."""
import os
import pytest
from tests.helpers import skip_if_gated_repo_inaccessible
from tests.helpers.mark import hardware_marks
from tests.helpers.media import generate_synthetic_image
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
DEFAULT_MODEL = "Lightricks/LTX-2.5-Diffusers"
DEFAULT_REVISION = "a6de4b5354f078db24d9cf4778c14846788aea3d"
MODEL = os.getenv("VLLM_TEST_LTX25_MODEL", DEFAULT_MODEL)
MODEL_REVISION = os.getenv("VLLM_TEST_LTX25_MODEL_REVISION", DEFAULT_REVISION if MODEL == DEFAULT_MODEL else "")
PROMPT = "A red fox walks through a snowy forest while the camera remains fixed."
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
@pytest.fixture(scope="module", autouse=True)
def require_ltx25_model_access() -> None:
if not os.path.isdir(MODEL):
skip_if_gated_repo_inaccessible(
MODEL,
revision=MODEL_REVISION or None,
filename="model_index.json",
)
def _server(model_class_name: str) -> OmniServerParams:
return OmniServerParams(
model=MODEL,
server_args=[
*(["--revision", MODEL_REVISION] if MODEL_REVISION else []),
"--model-class-name",
model_class_name,
"--enforce-eager",
"--enable-layerwise-offload",
"--diffusion-attention-backend",
"CUDNN_ATTN",
],
)
def _cases():
return [
pytest.param(_server("LTX2Pipeline"), 30, id="full_one_stage", marks=SINGLE_CARD_MARKS),
pytest.param(_server("LTX2TwoStagePipeline"), 30, id="full_two_stage", marks=SINGLE_CARD_MARKS),
pytest.param(_server("LTX2DistilledOneStagePipeline"), 8, id="distilled_one_stage", marks=SINGLE_CARD_MARKS),
pytest.param(_server("LTX2DistilledTwoStagePipeline"), 8, id="distilled_two_stage", marks=SINGLE_CARD_MARKS),
]
@pytest.mark.parametrize(("omni_server", "num_inference_steps"), _cases(), indirect=["omni_server"])
def test_ltx25_pipeline_entries(
omni_server: OmniServer,
num_inference_steps: int,
openai_client: OpenAIClientHandler,
subtests,
) -> None:
"""Generate T2V and first-frame I2V through the synchronous video API."""
for task in ("t2v", "i2v"):
with subtests.test(task=task):
request_config = {
"model": omni_server.model,
"form_data": {
"model": omni_server.model,
"prompt": PROMPT,
"height": 128,
"width": 128,
"num_frames": 9,
"fps": 24,
"num_inference_steps": num_inference_steps,
"seed": 42,
},
}
if task == "i2v":
request_config["image_reference"] = (
f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
)
openai_client.send_video_diffusion_request(request_config)
@@ -14,9 +14,10 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
import pytest
from tests.helpers.mark import hardware_test
from tests.helpers.media import load_test_audio_data_url
from tests.helpers.media import get_asset_path, load_test_audio_data_url
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.entrypoints.openai.serving_speech import _DEFAULT_VOICE_NAME
try:
from transformers import HiggsAudioV2TokenizerModel # noqa: F401
@@ -140,3 +141,121 @@ class TestOmniVoiceVoiceCloning:
"min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
}
openai_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
@pytest.mark.parametrize("voice", [_DEFAULT_VOICE_NAME, {"id": _DEFAULT_VOICE_NAME}])
def test_speech_with_voice_param_accepted(omni_server, openai_client, voice) -> None:
"""The default voice should be accepted as both a plain string and a
VoiceID dict, even when the model has no uploaded speakers."""
request_config = {
"model": omni_server.model,
"input": get_prompt("text"),
"voice": voice,
"response_format": "wav",
"timeout": 180.0,
"min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
}
openai_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
def test_unknown_voice_rejected(omni_server, openai_client) -> None:
"""An unrecognized voice name should be rejected with a 400."""
request_config = {
"model": omni_server.model,
"input": get_prompt("text"),
"voice": "nonexistent_voice",
"status_code": 400,
"err_message": "Invalid voice",
}
openai_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
def test_voice_request_before_and_after_clone_registration(omni_server, openai_client) -> None:
"""The default voice should always be listed and usable, even after
a different cloned voice is registered."""
speech_config = {
"model": omni_server.model,
"input": get_prompt("text"),
"voice": _DEFAULT_VOICE_NAME,
}
# 1) "default" should appear in the voices list before any uploads
voices = openai_client.send_audio_voices_list_http_request()[0]
assert _DEFAULT_VOICE_NAME in voices.json_body["voices"]
# 2) Speech with the default voice should succeed
openai_client.send_audio_speech_request(speech_config)
# 3) Register a cloned voice under a different name
register_config = {
"data": {"name": "bar", "consent": "test-consent"},
"files": {
"audio_sample": (
"clone_2.wav",
get_asset_path("qwen3_tts/clone_2.wav").read_bytes(),
"audio/wav",
),
},
}
openai_client.send_audio_voices_create_http_request(register_config)
# 4) "default" should still be listed alongside the new voice
voices = openai_client.send_audio_voices_list_http_request()[0]
assert _DEFAULT_VOICE_NAME in voices.json_body["voices"]
assert "bar" in voices.json_body["voices"]
# 5) Speech with the default voice should still succeed
openai_client.send_audio_speech_request(speech_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
def test_registered_default_voice_overrides_placeholder(omni_server, openai_client) -> None:
"""When a real voice is uploaded under the default name, it should be
used as a cloned voice; after deletion, the placeholder behavior resumes."""
speech_config = {
"model": omni_server.model,
"input": get_prompt("text"),
"voice": _DEFAULT_VOICE_NAME,
}
# 1) Placeholder default works before any registration
openai_client.send_audio_speech_request(speech_config)
# 2) Register a real voice under the default name
register_config = {
"data": {"name": _DEFAULT_VOICE_NAME, "consent": "test-consent"},
"files": {
"audio_sample": (
"clone_2.wav",
get_asset_path("qwen3_tts/clone_2.wav").read_bytes(),
"audio/wav",
),
},
}
openai_client.send_audio_voices_create_http_request(register_config)
# 3) After creating the voice with the default name, it's visible as an uploaded_voice
voices = openai_client.send_audio_voices_list_http_request()[0]
assert _DEFAULT_VOICE_NAME in voices.json_body["voices"]
uploaded_names = [v["name"] for v in voices.json_body["uploaded_voices"]]
assert _DEFAULT_VOICE_NAME in uploaded_names
openai_client.send_audio_speech_request(speech_config)
# 4) Delete the uploaded default voice
openai_client.send_audio_voices_delete_http_request(
{"name": _DEFAULT_VOICE_NAME},
)
# 5) Ensure the placeholder behavior is consistent, but it's not an uploaded voice now
voices = openai_client.send_audio_voices_list_http_request()[0]
assert _DEFAULT_VOICE_NAME in voices.json_body["voices"]
uploaded_names = [v["name"] for v in voices.json_body["uploaded_voices"]]
assert _DEFAULT_VOICE_NAME not in uploaded_names
openai_client.send_audio_speech_request(speech_config)
@@ -6,7 +6,7 @@ Online serving tests for the Qwen-Image-Edit family (image-to-image via chat com
- ``test_single_image_to_image_001``: ``Qwen/Qwen-Image-Edit`` — one reference image, fixed 512×512.
- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2511`` — two reference images, fixed 512×512.
- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2511`` only, ``advanced_model`` — mixed input
- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2511`` only, ``slow`` — mixed input
resolutions; ``extra_body`` uses per-output ``width``/``height`` lists; the test client sends one
scalar-size request per list index in parallel and merges images (see ``OpenAIClientHandler.send_diffusion_request``).
@@ -52,7 +52,7 @@ def _get_diffusion_feature_cases(model: str):
]
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize(
"omni_server",
@@ -82,8 +82,7 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
openai_client.send_diffusion_request(request_config)
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize(
"omni_server",
@@ -114,7 +113,7 @@ def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenA
openai_client.send_diffusion_request(request_config)
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize(
"omni_server",
@@ -44,8 +44,7 @@ def _get_diffusion_feature_cases(model: str):
]
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+2 -3
View File
@@ -46,8 +46,7 @@ tts_server_params = [
]
@pytest.mark.advanced_model
@pytest.mark.core_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
@@ -72,7 +71,7 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
openai_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT)
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+2 -2
View File
@@ -32,7 +32,7 @@ def _get_diffusion_feature_cases():
]
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize(
"omni_server",
@@ -60,7 +60,7 @@ def test_basic_001(
openai_client.send_diffusion_request(request_config)
@pytest.mark.advanced_model
@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize(
"omni_server",
@@ -161,6 +161,17 @@ def _make_diffusion_plan(
)
def _make_stage_runtime() -> StageRuntime:
return StageRuntime(
stage_configs=[],
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
"""Regression test for forward references in make_async_mp_client."""
import vllm_omni.engine.stage_engine_core_client as client_mod
@@ -233,6 +244,48 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients():
assert cleanup_clients == [shared, extra]
def test_initialize_local_diffusion_replica_scopes_runtime_env(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
from vllm_omni.platforms import current_omni_platform
runtime = _make_stage_runtime()
plan = _make_diffusion_plan(0, stage_id=0).replicas[0]
runtime_env_var = "VLLM_OMNI_TEST_STAGE_RUNTIME_ENV"
device_env_var = current_omni_platform.device_control_env_var
runtime._init_visible_devices_baseline = "0,1"
plan.metadata.runtime_cfg = {
"devices": "0",
"env": {runtime_env_var: "stage-value"},
}
monkeypatch.delenv(runtime_env_var, raising=False)
monkeypatch.setenv(device_env_var, "0,1")
captured: dict[str, str | None] = {}
monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None)
def _capture_launch_diffusion_stage_replica(**_kwargs):
captured["runtime_env"] = os.environ.get(runtime_env_var)
captured["device_env"] = os.environ.get(device_env_var)
raise RuntimeError("stop after capturing launch environment")
monkeypatch.setattr(
runtime_mod,
"launch_diffusion_stage_replica",
_capture_launch_diffusion_stage_replica,
)
with pytest.raises(RuntimeError, match="stop after capturing launch environment"):
runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1)
assert captured == {
"runtime_env": "stage-value",
"device_env": "0",
}
assert runtime_env_var not in os.environ
assert os.environ[device_env_var] == "0,1"
def test_initialize_local_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
from vllm_omni.engine.stage_engine_startup import StageReplicaResources
@@ -311,6 +364,136 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline
}
def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
runtime = _make_stage_runtime()
plan = _make_llm_plan(0, stage_id=0, vllm_config=object()).replicas[0]
plan.engine_args_dict = {}
runtime_env_var = "VLLM_OMNI_TEST_STAGE_RUNTIME_ENV"
plan.metadata.runtime_cfg = {
"devices": "0",
"env": {runtime_env_var: "stage-value"},
}
monkeypatch.setenv(runtime_env_var, "parent-value")
captured: dict[str, str | None] = {}
@contextlib.contextmanager
def _capture_launch_stage_replica(*, stage_visible_devices, **_kwargs):
captured["runtime_env"] = os.environ.get(runtime_env_var)
captured["stage_visible_devices"] = stage_visible_devices
yield None
monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: [])
monkeypatch.setattr(runtime_mod, "launch_stage_replica", _capture_launch_stage_replica)
with pytest.raises(RuntimeError, match="launcher returned no resources"):
runtime._initialize_local_llm_replica(plan, stage_init_timeout=1)
assert captured == {
"runtime_env": "stage-value",
"stage_visible_devices": "0",
}
assert os.environ[runtime_env_var] == "parent-value"
def test_initialize_diffusion_stage_applies_client_batch_size_to_engine(monkeypatch):
import vllm_omni.diffusion.stage_diffusion_client as client_mod
import vllm_omni.engine.stage_init_utils as init_mod
od_config = types.SimpleNamespace(max_num_seqs=1)
captured: dict[str, object] = {}
monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config)
def _capture_client(model, config, metadata, stage_init_timeout, batch_size, use_inline):
captured.update(
model=model,
config=config,
metadata=metadata,
stage_init_timeout=stage_init_timeout,
batch_size=batch_size,
use_inline=use_inline,
)
return object()
monkeypatch.setattr(client_mod, "create_diffusion_client", _capture_client)
metadata = _make_diffusion_metadata(0)
init_mod.initialize_diffusion_stage(
0,
"dummy-model",
types.SimpleNamespace(),
metadata,
stage_init_timeout=12,
batch_size=4,
use_inline=True,
)
assert od_config.max_num_seqs == 4
assert captured == {
"model": "dummy-model",
"config": od_config,
"metadata": metadata,
"stage_init_timeout": 12,
"batch_size": 4,
"use_inline": True,
}
def test_launch_diffusion_stage_replica_applies_batch_size_to_config(monkeypatch):
import vllm_omni.diffusion.stage_diffusion_client as client_mod
import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
import vllm_omni.engine.stage_engine_startup as startup_mod
od_config = types.SimpleNamespace(max_num_seqs=1, parallel_config=types.SimpleNamespace(world_size=1))
monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)
monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: [])
monkeypatch.setattr(
startup_mod,
"register_stage_with_omni_master",
lambda **kwargs: types.SimpleNamespace(
handshake_address="tcp://127.0.0.1:26001",
input_address="tcp://127.0.0.1:26002",
output_address="tcp://127.0.0.1:26003",
),
)
omni_master_server = types.SimpleNamespace(
address="127.0.0.1",
port=25000,
release_route_port_reservations=lambda *args, **kwargs: None,
)
proc_manager = types.SimpleNamespace(
addresses=types.SimpleNamespace(
inputs=["tcp://127.0.0.1:26002"],
outputs=["tcp://127.0.0.1:26003"],
)
)
monkeypatch.setattr(proc_mod, "StageDiffusionProcManager", lambda **kwargs: proc_manager)
sentinel_client = object()
monkeypatch.setattr(
client_mod.StageDiffusionClient,
"from_addresses",
lambda metadata, **kwargs: sentinel_client,
)
result, resources = startup_mod.launch_diffusion_stage_replica(
model="dummy-model",
stage_config=types.SimpleNamespace(),
metadata=types.SimpleNamespace(stage_id=0),
stage_init_timeout=12,
batch_size=4,
use_inline=False,
omni_master_server=omni_master_server,
)
assert result is sentinel_client
assert od_config.max_num_seqs == 4
assert resources.manager is proc_manager
def test_stage_runtime_initializes_stage_pools(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
+53
View File
@@ -270,3 +270,56 @@ class TestMsgpackRoundTrip:
# Content merged from legacy nested key
assert result.prompt == "legacy prompt"
assert result.prompt_token_ids == [1, 2]
class TestRequestOutputFieldParity:
"""``OmniRequestOutput`` must expose everything ``RequestOutput`` sets.
``RequestOutput`` is a plain class, so the dataclass-generated ``__init__``
replaces its ``__init__`` and the inherited attributes only exist because
they are redeclared as fields. Any vLLM field that is not redeclared simply
does not exist on an omni output, and vLLM's own serving code reads those
attributes unconditionally — ``/v1/completions`` returned HTTP 500 with
``AttributeError: 'OmniRequestOutput' object has no attribute
'ec_transfer_params'`` for exactly that reason. This test fails on the next
vLLM bump that adds a field, instead of leaving it for an endpoint to hit.
"""
def test_carries_every_attribute_request_output_sets(self) -> None:
reference = _make_text_request_output()
omni = OmniRequestOutput(request_id="req-parity")
missing = sorted(name for name in vars(reference) if not name.startswith("_") and not hasattr(omni, name))
assert not missing, (
f"OmniRequestOutput is missing RequestOutput attributes {missing}; "
"redeclare them as dataclass fields in vllm_omni/outputs/__init__.py"
)
def test_transfer_param_fields_default_to_none(self) -> None:
omni = OmniRequestOutput(request_id="req-defaults")
assert omni.kv_transfer_params is None
assert omni.ec_transfer_params is None
assert omni.num_cached_tokens is None
assert omni.num_cache_creation_tokens is None
def test_transfer_param_fields_are_copied_from_the_stage_output(self) -> None:
"""Declaring the fields is not enough — they must also be copied.
``from_stage_output`` only carries the names listed in
``_REQUEST_OUTPUT_CONTENT_ATTRS``. With the fields declared but absent
from that list the endpoint stops raising ``AttributeError`` while
silently dropping external-connector metadata and cache-creation
accounting, which is harder to notice than the crash it replaced.
"""
source = _make_text_request_output()
source.num_cache_creation_tokens = 7
source.ec_transfer_params = {"connector": "shm", "handle": "abc"}
source.num_cached_tokens = 3
source.kv_transfer_params = {"role": "producer"}
omni = OmniRequestOutput.from_stage_output(source, stage_id=1)
assert omni.num_cache_creation_tokens == 7
assert omni.ec_transfer_params == {"connector": "shm", "handle": "abc"}
assert omni.num_cached_tokens == 3
assert omni.kv_transfer_params == {"role": "producer"}
+6 -3
View File
@@ -1138,7 +1138,8 @@ class TestSingleStageReplicaInitialization:
runtime._init_visible_devices_baseline = "0"
mocker.patch.object(runtime_mod, "inject_kv_stage_info")
mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value="diffusion-config")
od_config = SimpleNamespace(max_num_seqs=None, parallel_config=SimpleNamespace(world_size=1))
mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config)
mock_register = mocker.patch(
"vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
return_value=StageRegistrationResponse(
@@ -1174,6 +1175,7 @@ class TestSingleStageReplicaInitialization:
os.environ[device_env_var] = prev_device_env
assert result is sentinel_client
assert od_config.max_num_seqs == 4
mock_register.assert_called_once_with(
omni_master_address="127.0.0.1",
omni_master_port=25000,
@@ -1183,7 +1185,7 @@ class TestSingleStageReplicaInitialization:
)
mock_manager.assert_called_once_with(
model="fake-model",
od_config="diffusion-config",
od_config=od_config,
stage_init_timeout=60,
handshake_address="tcp://127.0.0.1:26001",
addresses=mocker.ANY,
@@ -1230,7 +1232,8 @@ class TestSingleStageReplicaInitialization:
runtime._init_visible_devices_baseline = "0"
mocker.patch.object(runtime_mod, "inject_kv_stage_info")
mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value="diffusion-config")
od_config = SimpleNamespace(max_num_seqs=None, parallel_config=SimpleNamespace(world_size=1))
mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config)
mocker.patch(
"vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
return_value=StageRegistrationResponse(
@@ -32,7 +32,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
VideoParams,
VideoResponse,
)
from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage
from vllm_omni.entrypoints.openai.storage import LocalStorageManager
from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry
from vllm_omni.errors import GuardrailViolationError
@@ -461,6 +461,51 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client,
assert input_image.size == (96, 64)
def test_i2v_resize_policy_can_defer_to_pipeline(monkeypatch):
engine = FakeAsyncOmni()
engine.get_diffusion_od_config = lambda: SimpleNamespace(
model="org/model",
model_class_name="ExamplePipeline",
revision="pinned-revision",
)
captured = {}
def fake_policy(model_class_name, *, model, revision=None):
captured.update(
model_class_name=model_class_name,
model=model,
revision=revision,
)
return True
monkeypatch.setattr(
"vllm_omni.entrypoints.openai.serving_video.should_preserve_reference_image_size",
fake_policy,
)
handler = OmniOpenAIServingVideo.for_diffusion(
diffusion_engine=engine,
model_name="fallback/model",
)
image = Image.new("RGB", (48, 32))
asyncio.run(
handler._run_and_extract(
VideoGenerationRequest(prompt="A bear playing with yarn.", width=96, height=64),
"pipeline-owned-resize",
reference_image=ReferenceImage(image),
)
)
input_image = engine.captured_prompt["multi_modal_data"]["image"]
assert isinstance(input_image, Image.Image)
assert input_image.size == (48, 32)
assert captured == {
"model_class_name": "ExamplePipeline",
"model": "org/model",
"revision": "pinned-revision",
}
def test_i2v_extra_params_dimensions_preserve_input_image_geometry(test_client, mocker: MockerFixture):
image_bytes = _make_test_image_bytes((48, 48))
mocker.patch(
@@ -10,8 +10,10 @@ ModelScope path.
import sys
import types
from pathlib import Path
import pytest
from pytest_mock import MockerFixture
from vllm import envs
from vllm_omni.entrypoints import omni_base
@@ -21,6 +23,64 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
MODEL_ID = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
@pytest.mark.parametrize(
"model_uri",
[
"s3://bucket/model",
"gs://bucket/model",
"az://bucket/model",
],
)
def test_omni_snapshot_download_preserves_object_storage_uri(
model_uri: str,
mocker: MockerFixture,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False)
hf_download = mocker.patch.object(omni_base, "download_weights_from_hf_specific")
assert omni_base.omni_snapshot_download(model_uri) == model_uri
hf_download.assert_not_called()
def test_omni_snapshot_download_preserves_existing_local_path(
tmp_path: Path,
mocker: MockerFixture,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False)
absolute_model_path = tmp_path / "absolute-model"
absolute_model_path.mkdir()
relative_model_path = "relative-model"
(tmp_path / relative_model_path).mkdir()
monkeypatch.chdir(tmp_path)
hf_download = mocker.patch.object(omni_base, "download_weights_from_hf_specific")
assert omni_base.omni_snapshot_download(str(absolute_model_path)) == str(absolute_model_path)
assert omni_base.omni_snapshot_download(relative_model_path) == relative_model_path
hf_download.assert_not_called()
def test_omni_snapshot_download_uses_hf_for_relative_repo_id(
mocker: MockerFixture,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False)
model_id = "org/model"
# Stub the modular-index probe so the test never reaches huggingface.co.
hf_file_probe = mocker.patch.object(omni_base, "file_or_path_exists", return_value=False)
hf_download = mocker.patch.object(omni_base, "download_weights_from_hf_specific")
assert omni_base.omni_snapshot_download(model_id) == model_id
hf_file_probe.assert_called_once_with(model_id, "modular_model_index.json", revision=None)
hf_download.assert_called_once_with(
model_name_or_path=model_id,
cache_dir=None,
allow_patterns=["*"],
require_all=True,
)
@pytest.fixture
def download_backend(monkeypatch: pytest.MonkeyPatch):
"""Run ``omni_snapshot_download`` and report which backend it selected.
@@ -49,7 +49,9 @@ def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
def test_dreamzero_enrich_config_preserves_explicit_model_class_name(monkeypatch):
monkeypatch.setattr(
"vllm.transformers_utils.config.get_hf_file_to_dict",
lambda path, _model: None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]},
lambda path, _model, **_kwargs: (
None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]}
),
)
od_config = OmniDiffusionConfig(
+43
View File
@@ -22,6 +22,7 @@ from vllm_omni.engine.stage_init_utils import (
from vllm_omni.entrypoints.utils import (
_convert_dataclasses_to_dict,
_filter_dict_like_object,
_try_resolve_omni_model_type,
coerce_param_message_types,
filter_dataclass_kwargs,
filter_stages,
@@ -371,6 +372,48 @@ class TestResolveModelConfigPath:
assert sender["extra"]["connector_get_max_wait_first_chunk"] == 3000
assert sender["extra"]["connector_get_max_wait"] == 300
def test_object_storage_uri_resolves_via_materialized_configs(
self,
mocker: MockerFixture,
tmp_path,
):
"""An object-storage URI must not reach HF helpers; resolution reads the
locally materialized config files instead."""
hf_config = SimpleNamespace(
model_type="qwen3_omni_moe",
architectures=["Qwen3OmniMoeForConditionalGeneration"],
)
uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct"
materialize = mocker.patch(
"vllm_omni.entrypoints.utils._materialize_object_storage_configs",
return_value=str(tmp_path),
)
get_config = mocker.patch(
"vllm_omni.entrypoints.utils.get_config",
return_value=hf_config,
)
result = resolve_model_config_path(uri)
materialize.assert_called_once_with(uri)
get_config.assert_called_once()
assert get_config.call_args[0][0] == str(tmp_path)
assert result is not None
assert Path(result).as_posix().endswith("qwen3_omni_moe.yaml")
class TestTryResolveOmniModelType:
"""Name matching must only scan the last path component."""
def test_uri_bucket_name_cannot_hijack_pipeline(self):
# "qwen3-tts-models" normalizes to "qwen3ttsmodels", which contains the
# registered key "qwen3tts" — only the basename may participate, so a
# misleadingly named bucket selects nothing.
assert _try_resolve_omni_model_type("s3://qwen3-tts-models/plain-checkpoint") is None
def test_basename_still_matches_pipeline(self):
assert _try_resolve_omni_model_type("gs://any-bucket/Fun-CosyVoice3-0.5B-2512") == "cosyvoice3"
class TestLoadAndResolveStageConfigs:
def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture):
+9 -4
View File
@@ -11,11 +11,16 @@ affect in-process Omni (``OmniRunner`` / offline e2e) vs subprocess-based
import pytest
def skip_if_gated_repo_inaccessible(repo_id: str) -> None:
def skip_if_gated_repo_inaccessible(
repo_id: str,
*,
revision: str | None = None,
filename: str = "config.json",
) -> None:
"""Skip the test if a gated HuggingFace repo is not accessible.
Tries to download the model's config.json via ``hf_hub_download``,
which performs an actual file-access check (unlike ``HfApi().model_info()``
Tries to download ``filename`` via ``hf_hub_download``, which performs an
actual file-access check (unlike ``HfApi().model_info()``
that only checks metadata). If the token has metadata access but not
file-download access, ``hf_hub_download`` will raise ``GatedRepoError``
and we skip cleanly.
@@ -26,7 +31,7 @@ def skip_if_gated_repo_inaccessible(repo_id: str) -> None:
except Exception:
return
try:
hf_hub_download(repo_id=repo_id, filename="config.json")
hf_hub_download(repo_id=repo_id, filename=filename, revision=revision)
except GatedRepoError as exc:
pytest.skip(
f"Skipping: gated HF repo {repo_id!r} inaccessible to the current "
@@ -0,0 +1,363 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""CPU coverage for Gepard's prompt layout, streaming-decode window and config.
The e2e test needs a GPU, NeMo and the checkpoint, so it only runs nightly.
These pin the parts that are pure logic and break silently: which frames each
decode covers, which samples survive the lookback trim, the prompt layout, and
that the packaged deploy config sits where the example looks for it.
The codec is replaced by a shift-invariant stand-in whose output depends only
on the codes, which makes "chunked emissions concatenate to one whole-history
decode" an exact assertion — the contract lookback exists to provide.
"""
from __future__ import annotations
from pathlib import Path
import pytest
import torch
import torch.nn as nn
import vllm_omni
from vllm_omni.config.stage_config import load_deploy_config
from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
from vllm_omni.model_executor.models.gepard.gepard_talker import (
CHUNK_FRAMES,
FIRST_CHUNK_FRAMES,
GepardTalkerForConditionalGeneration,
_GepardState,
)
from vllm_omni.model_executor.models.gepard.nanocodec import (
LOOKBACK_FRAMES,
NanoCodec,
apply_end_of_speech_tail,
)
from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
SAMPLES_PER_FRAME = 1024
NUM_HEADS = 32
SAMPLE_RATE = 22050
# Text-token count at or above which the layout is left unrepeated, read from
# the checkpoint defaults so the boundary cases follow the config.
APPLY_BELOW = GepardConfig().text_repetition_apply_below
class _ShiftInvariantCodec:
"""Frame tagged ``v`` -> samples ``v*SPF .. v*SPF+SPF-1``, wherever it sits."""
def __init__(self) -> None:
self.windows: list[int] = [] # frame count of every decode call
def decode_from_codes(self, codes: torch.Tensor, codes_len: torch.Tensor):
assert codes.shape[0] == 1, "streaming decode is per-request"
assert codes.shape[1] == NUM_HEADS, f"expected {NUM_HEADS} code channels"
tags = codes[0, 0, :].to(torch.int64) # (T,)
self.windows.append(int(tags.shape[0]))
offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64)
audio = (tags.unsqueeze(1) * SAMPLES_PER_FRAME + offsets).reshape(1, -1)
return audio.to(torch.float32), codes_len * SAMPLES_PER_FRAME
def _make_codec() -> tuple[NanoCodec, _ShiftInvariantCodec]:
codec = NanoCodec(codec_id="test/fake-codec", sample_rate=SAMPLE_RATE)
fake = _ShiftInvariantCodec()
codec._codec = fake # bypasses load(); NeMo is never imported
assert codec.is_loaded
return codec, fake
def _make_talker(codec: NanoCodec) -> GepardTalkerForConditionalGeneration:
"""A real talker with only the audio-path attributes populated.
``__init__`` needs a VllmConfig and builds the backbone; the methods under
test touch only these four attributes.
"""
talker = GepardTalkerForConditionalGeneration.__new__(GepardTalkerForConditionalGeneration)
nn.Module.__init__(talker) # sets up _parameters/_buffers/_modules only
talker._codec = codec
talker._audio_queue = []
talker._active_states = {}
talker._deferred_cleanup_ids = set()
return talker
def _frame(tag: int) -> torch.Tensor:
"""One committed frame: 32 per-dimension codes, all carrying the tag."""
return torch.full((NUM_HEADS,), tag, dtype=torch.long)
def _expected_samples(num_frames: int) -> torch.Tensor:
offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64)
tags = torch.arange(num_frames, dtype=torch.int64).unsqueeze(1)
return (tags * SAMPLES_PER_FRAME + offsets).reshape(-1).to(torch.float32)
def _drain(talker) -> torch.Tensor:
"""Concatenate and clear the queued deltas, as make_omni_output does."""
if not talker._audio_queue:
return torch.empty(0)
out = torch.cat([a.reshape(-1) for _, a in talker._audio_queue])
talker._audio_queue.clear()
return out
def _generate(talker, state, num_frames: int, *, stop_at_end: bool) -> None:
"""Drive the commit loop the way forward() does, one frame per step."""
for tag in range(num_frames):
state.frames.append(_frame(tag))
state.frame_count += 1
state.past_first_step = True
talker._emit_audio(state, is_final=stop_at_end and tag == num_frames - 1)
# Counts straddling both cadence thresholds, where an off-by-one hides.
@pytest.mark.parametrize(
"num_frames",
[
1,
2,
FIRST_CHUNK_FRAMES - 1,
FIRST_CHUNK_FRAMES,
FIRST_CHUNK_FRAMES + 1,
CHUNK_FRAMES - 1,
CHUNK_FRAMES,
CHUNK_FRAMES + 1,
FIRST_CHUNK_FRAMES + CHUNK_FRAMES,
79,
200,
],
)
def test_chunked_emissions_equal_a_whole_history_decode(num_frames: int) -> None:
codec, _fake = _make_codec()
talker = _make_talker(codec)
state = _GepardState(request_id="req-0")
_generate(talker, state, num_frames, stop_at_end=True)
torch.testing.assert_close(_drain(talker), _expected_samples(num_frames))
assert state.emitted_frames == num_frames
def test_every_frame_is_emitted_exactly_once() -> None:
"""The trim must not drop or duplicate samples at a chunk boundary."""
codec, _fake = _make_codec()
talker = _make_talker(codec)
state = _GepardState(request_id="req-0")
num_frames = FIRST_CHUNK_FRAMES + 2 * CHUNK_FRAMES + 3
_generate(talker, state, num_frames, stop_at_end=True)
audio = _drain(talker)
assert audio.numel() == num_frames * SAMPLES_PER_FRAME
# Tags are globally unique, so a duplicated or dropped frame shows up here.
assert len(torch.unique(audio)) == audio.numel()
def test_decode_windows_carry_lookback_context() -> None:
"""Every decode after the first must re-decode LOOKBACK_FRAMES of context."""
codec, fake = _make_codec()
talker = _make_talker(codec)
state = _GepardState(request_id="req-0")
num_frames = FIRST_CHUNK_FRAMES + 2 * CHUNK_FRAMES
_generate(talker, state, num_frames, stop_at_end=True)
assert len(fake.windows) >= 3, "expected a first chunk plus steady-state chunks"
assert fake.windows[0] == FIRST_CHUNK_FRAMES, "the first chunk has no left context"
for window in fake.windows[1:]:
assert window == CHUNK_FRAMES + LOOKBACK_FRAMES
def test_codec_runs_once_per_chunk_not_once_per_frame() -> None:
codec, fake = _make_codec()
talker = _make_talker(codec)
state = _GepardState(request_id="req-0")
num_frames = FIRST_CHUNK_FRAMES + 3 * CHUNK_FRAMES
_generate(talker, state, num_frames, stop_at_end=True)
assert len(fake.windows) == 4
assert sum(fake.windows) < num_frames * 2 # far below one decode per frame
def test_final_emit_drains_a_partial_chunk() -> None:
"""A request can end mid-chunk (max_tokens, STOP); the frames committed
since the last cadence boundary must still reach the queue.
Arithmetic only: this drives ``_emit_audio`` directly, so it says nothing
about *whether* the engine reaches it on the request's last step. That
wiring lives in ``test_gepard_wiring.py`` — a bug there leaves this green.
"""
codec, _fake = _make_codec()
talker = _make_talker(codec)
state = _GepardState(request_id="req-0")
num_frames = FIRST_CHUNK_FRAMES + CHUNK_FRAMES + 5
_generate(talker, state, num_frames, stop_at_end=False)
assert state.emitted_frames < num_frames, "a partial chunk should still be pending"
talker._emit_audio(state, is_final=True)
torch.testing.assert_close(_drain(talker), _expected_samples(num_frames))
assert state.emitted_frames == num_frames
def test_emitting_with_nothing_pending_is_a_no_op() -> None:
codec, fake = _make_codec()
talker = _make_talker(codec)
state = _GepardState(request_id="req-0")
_generate(talker, state, 4, stop_at_end=True)
calls_after_flush = len(fake.windows)
talker._audio_queue.clear()
talker._emit_audio(state, is_final=True)
assert len(fake.windows) == calls_after_flush, "codec ran with no new frames"
assert talker._audio_queue == []
def test_concurrent_requests_keep_separate_windows() -> None:
codec, _fake = _make_codec()
talker = _make_talker(codec)
states = [_GepardState(request_id=f"req-{i}") for i in range(2)]
for state in states:
_generate(talker, state, FIRST_CHUNK_FRAMES, stop_at_end=True)
by_req: dict[str, list[torch.Tensor]] = {}
for req_id, audio in talker._audio_queue:
by_req.setdefault(req_id, []).append(audio.reshape(-1))
assert set(by_req) == {"req-0", "req-1"}
for chunks in by_req.values():
torch.testing.assert_close(torch.cat(chunks), _expected_samples(FIRST_CHUNK_FRAMES))
def test_end_of_speech_tail_fades_and_pads() -> None:
audio = torch.ones(SAMPLE_RATE) # 1 s of full-scale signal
out = apply_end_of_speech_tail(audio, SAMPLE_RATE, fade_ms=10.0, silence_ms=20.0)
fade = int(SAMPLE_RATE * 10.0 / 1000.0)
pad = int(SAMPLE_RATE * 20.0 / 1000.0)
assert out.numel() == SAMPLE_RATE + pad
assert torch.all(out[-pad:] == 0.0), "trailing pad must be silent"
assert out[SAMPLE_RATE - fade - 1] == 1.0, "the fade must not reach back past fade_ms"
# Monotonically ramped to zero — the pad alone would leave the click.
ramp = out[SAMPLE_RATE - fade : SAMPLE_RATE]
assert torch.all(ramp[:-1] >= ramp[1:])
assert ramp[-1] == 0.0
def test_end_of_speech_tail_is_a_no_op_by_default() -> None:
audio = torch.ones(128)
out = apply_end_of_speech_tail(audio, SAMPLE_RATE, fade_ms=0.0, silence_ms=0.0)
assert out is audio
def _repeat_count(ids: list[int], cfg: GepardConfig) -> int:
"""How many ``[SOT ... EOT]`` copies the layout carries."""
return ids.count(cfg.start_of_text)
def test_gepard_prompt_layout() -> None:
"""The prompt carries the speaker slots and exactly one SOS, at the end."""
cfg = GepardConfig()
ids = build_gepard_prompt_ids([1, 2, 3], config=cfg)
slots = [cfg.speaker_token_base + i for i in range(cfg.num_speaker_prefix)]
assert ids[: cfg.num_speaker_prefix] == slots
assert ids.count(cfg.start_of_speech) == 1, "only the canonical copy may carry SOS"
assert ids[-1] == cfg.start_of_speech, "the first frame is sampled from the SOS position"
def test_short_text_repeats_to_reach_the_target_token_mass() -> None:
"""A short text is repeated until its region holds ~target_text_tokens."""
cfg = GepardConfig()
text = [1, 2, 3, 4] # target 16 / 4 -> 4 copies
ids = build_gepard_prompt_ids(text, config=cfg)
assert _repeat_count(ids, cfg) == 4
assert ids.count(cfg.start_of_speech) == 1, "only the last copy opens the speech region"
assert ids[-1] == cfg.start_of_speech
@pytest.mark.parametrize(
("num_text_tokens", "repeated"),
[
(APPLY_BELOW - 1, True),
(APPLY_BELOW, False),
(APPLY_BELOW + 1, False),
],
)
def test_apply_below_is_the_repetition_boundary(num_text_tokens: int, repeated: bool) -> None:
"""Both sides of the threshold: at apply_below the text is left alone.
The layout has to match training here or WER collapses, and the boundary
is where an off-by-one would be invisible in a single happy-path clip.
"""
cfg = GepardConfig()
ids = build_gepard_prompt_ids(list(range(1, num_text_tokens + 1)), config=cfg)
assert (_repeat_count(ids, cfg) > 1) is repeated
def test_max_repeats_caps_a_one_token_text() -> None:
"""target/1 would ask for 16 copies; the cap keeps the prefill bounded."""
cfg = GepardConfig()
ids = build_gepard_prompt_ids([1], config=cfg)
assert _repeat_count(ids, cfg) == cfg.text_repetition_max_repeats
def test_repetition_thresholds_are_read_from_the_config() -> None:
"""The numbers come from the checkpoint sidecar, not from literals.
A checkpoint whose text_repetition block differs must change the layout;
duplicating the values in the builder would silently ignore it.
"""
cfg = GepardConfig(text_repetition={"enabled": True, "target_text_tokens": 9, "apply_below": 8, "max_repeats": 2})
assert _repeat_count(build_gepard_prompt_ids([1, 2, 3], config=cfg), cfg) == 2 # ceil(9/3)=3, capped at 2
assert _repeat_count(build_gepard_prompt_ids(list(range(8)), config=cfg), cfg) == 1 # at apply_below
def test_repetition_can_be_disabled() -> None:
cfg = GepardConfig(text_repetition={"enabled": False})
ids = build_gepard_prompt_ids([1], config=cfg)
assert _repeat_count(ids, cfg) == 1
assert ids[-1] == cfg.start_of_speech
def test_empty_text_is_rejected() -> None:
"""An empty text would still assemble a valid layout and be voiced.
``[slots, SOT, EOT, SOS]`` parses, prefills and generates: the failure is
audible, not raised, which is the worst shape for a caller to debug.
"""
with pytest.raises(ValueError, match="at least one text token"):
build_gepard_prompt_ids([], config=GepardConfig())
def test_packaged_deploy_config_is_where_the_example_looks_for_it() -> None:
"""The example resolves its default off ``vllm_omni.__file__``; tests reach
the same YAML through ``tests.helpers.stage_config``, so nothing covered the
packaged location until this.
``pipeline: gepard`` is the load-bearing key: the checkpoint self-identifies
as qwen3_5_text, so without the pin the architectures fallback routes it to
the diffusion registry.
This asserts the location, not the example's constant — importing the
example would pull ``Omni`` and NeMo into a CPU test. Running the quick
start verbatim stays a release check.
"""
path = Path(vllm_omni.__file__).resolve().parent / "deploy" / "gepard.yaml"
assert path.is_file(), f"packaged deploy config missing at {path}"
deploy = load_deploy_config(path)
assert deploy.pipeline == "gepard"
assert deploy.stages, "deploy config declares no stages"
@@ -0,0 +1,731 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""CPU coverage for the seam between Gepard's talker and the AR runner.
``test_gepard_window.py`` pins the decode arithmetic *below* this layer, and
the e2e test exercises the layer *above* it but only along the path where the
stop head fires. Neither covers the wiring itself: whether anything calls the
final emit on a request's last step, and whether the payload it produces
survives the runner's sparse-audio routing.
These drive the real ``preprocess`` -> ``forward`` -> ``make_omni_output``
sequence through a stand-in for ``GPUARModelRunner.execute_model``, and use the
runner's own ``_resolve_sparse_mm_routing`` to decide what actually reaches the
caller. Only the LM backbone and the frame sampler are stubbed out; every
method under test is the real one, so a test here goes red when the engine
path breaks — which a test that calls the flush by hand cannot do.
"""
from __future__ import annotations
import logging
import pytest
import torch
import torch.nn as nn
from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
from vllm_omni.model_executor.models.gepard.gepard_talker import (
CHUNK_FRAMES,
FIRST_CHUNK_FRAMES,
GepardTalkerForConditionalGeneration,
_ReqInfo,
)
from vllm_omni.model_executor.models.gepard.nanocodec import NanoCodec
from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
SAMPLES_PER_FRAME = 1024
NUM_HEADS = 32
SAMPLE_RATE = 22050
LOGGER_NAME = "vllm_omni.model_executor.models.gepard.gepard_talker"
# Each request tags its frames from a distinct base, so a sample that reaches
# the wrong request is visible in the value, not just the length.
_TAG_STRIDE = 100_000
class _TagCodec:
"""Frame carrying tag ``v`` -> samples ``v*SPF .. v*SPF+SPF-1``, wherever it sits.
Shift-invariant, so "what the caller received" can be compared against one
whole-history decode exactly. The tag rides channel 1: channel 0 is head0,
which the talker keeps inside its real (small) vocab.
"""
def __init__(self) -> None:
self.decode_calls = 0
def decode_from_codes(self, codes: torch.Tensor, codes_len: torch.Tensor):
assert codes.shape[0] == 1, "streaming decode is per-request"
assert codes.shape[1] == NUM_HEADS, f"expected {NUM_HEADS} code channels"
self.decode_calls += 1
tags = codes[0, 1, :].to(torch.int64) # (T,)
offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64)
audio = (tags.unsqueeze(1) * SAMPLES_PER_FRAME + offsets).reshape(1, -1)
return audio.to(torch.float32), codes_len * SAMPLES_PER_FRAME
class _StubBackbone(nn.Module):
"""Stands in for the Qwen3.5 backbone; only shapes matter to this seam."""
def __init__(self, hidden: int) -> None:
super().__init__()
self.hidden = hidden
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return torch.zeros(int(input_ids.shape[0]), self.hidden)
def forward(
self,
input_ids: torch.Tensor | None = None,
positions: torch.Tensor | None = None,
intermediate_tensors=None,
inputs_embeds: torch.Tensor | None = None,
) -> torch.Tensor:
src = inputs_embeds if inputs_embeds is not None else input_ids
return torch.zeros(int(src.shape[0]), self.hidden)
def _expected_samples(tag_base: int, num_frames: int) -> torch.Tensor:
offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64)
tags = (torch.arange(num_frames, dtype=torch.int64) + tag_base).unsqueeze(1)
return (tags * SAMPLES_PER_FRAME + offsets).reshape(-1).to(torch.float32)
class _MiniEngine:
"""Drives the talker the way ``GPUARModelRunner.execute_model`` does.
Reproduces the three orderings the dropped-tail bug lived in:
* ``on_requests_finished`` fires at the START of the step AFTER the one
that finished the request (it reads ``scheduler_output.finished_req_ids``);
* a finished request is not scheduled again, so it is absent from
``req_ids_output_copy`` and from that step's routing;
* with nothing left to schedule there is no step at all, so no further
``forward`` runs;
* a stop-token finish gets one more scheduled step, a budget finish does
not -- under async scheduling only the latter is guarded against it.
"""
def __init__(self, talker: GepardTalkerForConditionalGeneration, *, max_model_len: int) -> None:
self.talker = talker
self.max_model_len = max_model_len
self.reqs: dict[str, dict] = {}
self.order: list[str] = []
self._finished_last_step: list[str] = []
self.delivered: dict[str, list[torch.Tensor]] = {}
self.committed_frames: dict[str, int] = {}
self.steps = 0
self._tags_handed_out = 0
def add_request(
self,
req_id: str,
*,
max_tokens: int,
stop_after: int | None = None,
text_len: int = 3,
seed: int | None = None,
) -> None:
"""Submit a request. An id already seen is a resubmit of that id.
Tags come from a counter, so a resubmitted id gets frames the previous
request could not have produced.
"""
prompt_ids = build_gepard_prompt_ids(list(range(10, 10 + text_len)), config=self.talker.config)
self._tags_handed_out += 1
self.reqs[req_id] = {
"prompt_ids": torch.tensor(prompt_ids, dtype=torch.long),
"prompt_len": len(prompt_ids),
"max_tokens": max_tokens,
"stop_after": stop_after,
"stop_fired": False,
"seed": seed,
"computed": 0,
"generated": 0,
"finished": False,
"ghost_step_pending": False,
"next_token": None,
"tag_base": _TAG_STRIDE * self._tags_handed_out,
}
if req_id not in self.order:
self.order.append(req_id)
self.delivered[req_id] = []
self.committed_frames[req_id] = 0
def abort(self, req_id: str) -> None:
"""Kill a request the way an aborted client does: no STOP, no budget."""
self.reqs[req_id]["finished"] = True
self._finished_last_step.append(req_id)
def _live(self) -> list[str]:
return [r for r in self.order if not self.reqs[r]["finished"]]
def step(self) -> bool:
"""Run one engine step. False when there was nothing to schedule."""
live = self._live()
if not live:
return False
talker = self.talker
self.steps += 1
# --- runner: notify the model of requests finished since last step ---
if self._finished_last_step:
talker.on_requests_finished(set(self._finished_last_step))
self._finished_last_step = []
# --- runner: per-request preprocess, in input_batch.req_ids order ---
embeds_parts, ids_parts = [], []
for req_id in live:
r = self.reqs[req_id]
is_prefill = r["computed"] < r["prompt_len"]
if is_prefill:
ids = r["prompt_ids"]
else:
ids = torch.tensor([r["next_token"] or 0], dtype=torch.long)
span = int(ids.shape[0])
out_ids, out_embeds, _ = talker.preprocess(
input_ids=ids,
input_embeds=None,
request_id=req_id,
_omni_prompt_len=r["prompt_len"],
_omni_num_computed_tokens=r["computed"],
_omni_is_prefill=is_prefill,
_omni_max_tokens=r["max_tokens"],
_omni_seed=r["seed"],
)
ids_parts.append(out_ids)
embeds_parts.append(out_embeds)
r["span"] = span
inputs_embeds = torch.cat(embeds_parts, dim=0)
input_ids = torch.cat(ids_parts, dim=0)
positions = torch.arange(int(input_ids.shape[0]), dtype=torch.long)
# --- runner: forward, then build the multimodal payload ---
hidden = talker.forward(input_ids, positions, inputs_embeds=inputs_embeds)
results = list(talker._results_queue)
omni_out = talker.make_omni_output(hidden)
mm = omni_out.multimodal_outputs
# compute_logits drains this; it does not touch the audio path.
talker._results_queue.clear()
for req_id in live:
state = talker._active_states.get(req_id)
if state is not None:
self.committed_frames[req_id] = state.frame_count
# --- runner: sparse-audio routing decides what reaches the caller ---
downstream, sparse_index, is_sparse = GPUARModelRunner._resolve_sparse_mm_routing(
engine_output_type="audio",
req_ids_output_copy=list(live),
downstream_req_ids=list(live),
multimodal_outputs=mm,
)
assert is_sparse, "gepard always marks meta.sparse_audio"
for rid in downstream:
self.delivered[rid].append(mm["model_outputs"][sparse_index[rid]].reshape(-1))
# --- runner/scheduler bookkeeping: advance and detect finishes ---
stopped = {rid for rid, _head0, do_stop in results if do_stop}
for req_id in live:
r = self.reqs[req_id]
r["computed"] += r["span"]
r["generated"] = r["computed"] - r["prompt_len"] + 1
r["next_token"] = 0
length_cap = min(r["max_tokens"], self.max_model_len - r["prompt_len"])
if r["ghost_step_pending"]:
# This step WAS the extra one. The engine finished the request
# on the stop token last step and discards whatever came back
# here, whether or not the model chose to stop again.
r["finished"] = True
self._finished_last_step.append(req_id)
elif req_id in stopped:
# Ends on the stop token: one more step is already scheduled.
r["ghost_step_pending"] = True
elif r["generated"] >= length_cap:
r["finished"] = True
self._finished_last_step.append(req_id)
return True
def run(self, max_steps: int = 4000) -> None:
for _ in range(max_steps):
if not self.step():
return
raise AssertionError("engine did not drain")
def audio_for(self, req_id: str) -> torch.Tensor:
chunks = self.delivered[req_id]
return torch.cat(chunks) if chunks else torch.empty(0)
def _make_talker(*, max_model_len: int = 4096) -> tuple[GepardTalkerForConditionalGeneration, _TagCodec]:
"""A real talker with only the audio-path collaborators populated.
``__init__`` needs a VllmConfig and builds the Qwen3.5 backbone; the seam
under test needs none of that, so the backbone and the frame sampler are
stubbed and everything else is the class's own code.
"""
cfg = GepardConfig()
# Fixed, not cfg.get_text_config().hidden_size: nothing in this seam reads
# the real width, and a small one keeps the test off the backbone config.
hidden = 64
codec = NanoCodec(codec_id="test/fake-codec", sample_rate=SAMPLE_RATE)
fake = _TagCodec()
codec._codec = fake # bypasses load(); NeMo is never imported
assert codec.is_loaded
talker = GepardTalkerForConditionalGeneration.__new__(GepardTalkerForConditionalGeneration)
nn.Module.__init__(talker)
talker.config = cfg
talker.model = _StubBackbone(hidden)
# requires_grad=False: preprocess writes the prefix into inputs_embeds
# in place, and production runs that under inference_mode.
talker.null_prefix = nn.Parameter(torch.zeros(cfg.num_speaker_prefix, hidden), requires_grad=False)
talker._codec = codec
talker._sample_rate = SAMPLE_RATE
talker._max_model_len = max_model_len
talker._active_states = {}
talker._pending_requests = []
talker._results_queue = []
talker._audio_queue = []
talker._deferred_cleanup_ids = set()
return talker, fake
def _install_scripted_sampling(talker: GepardTalkerForConditionalGeneration, engine: _MiniEngine) -> None:
"""Replace the 32-head sampler with a scripted one.
Frames carry a per-request, per-step tag so the caller's audio is
self-describing. STOP fires once, on the request's configured frame: on the
extra step the engine then runs, re-firing is likely but not guaranteed,
and the model may not assume it.
"""
head0_vocab = talker.config.head0_vocab_size
# Signature mirrors the real ``_sample_frame`` minus ``self``; the scripted
# sampler ignores ``generators`` because seeding is decided in
# ``_build_with_real_sampler``, which keeps the model's own sampler.
def _sample_frame(rows: torch.Tensor, generators: list[torch.Generator | None] | None = None):
n = int(rows.shape[0])
head0 = torch.zeros(n, dtype=torch.long)
heads = torch.zeros(n, NUM_HEADS - 1, dtype=torch.long)
stop = torch.zeros(n, dtype=torch.bool)
for i, (req_id, _span, _samples_frame, *_rest) in enumerate(talker._pending_requests):
state = talker._active_states.get(req_id)
frame_index = state.frame_count if state is not None else 0
tag = engine.reqs[req_id]["tag_base"] + frame_index
head0[i] = tag % head0_vocab
heads[i] = tag
r = engine.reqs[req_id]
if r["stop_after"] is not None and frame_index >= r["stop_after"] and not r["stop_fired"]:
r["stop_fired"] = True
stop[i] = True
return head0, heads, stop
talker._sample_frame = _sample_frame
talker._audio_frame_embed = lambda h0, h: torch.zeros(int(h0.shape[0]), talker.model.hidden)
def _build(*, max_model_len: int = 4096) -> tuple[GepardTalkerForConditionalGeneration, _TagCodec, _MiniEngine]:
talker, codec = _make_talker(max_model_len=max_model_len)
engine = _MiniEngine(talker, max_model_len=max_model_len)
_install_scripted_sampling(talker, engine)
return talker, codec, engine
def _build_with_real_sampler(*, max_model_len: int = 4096) -> tuple[GepardTalkerForConditionalGeneration, _MiniEngine]:
"""Same seam, but ``_sample_frame`` is the model's own 32-head sampler.
The stub backbone hands every row the same hidden state, so the sampled
codes — and therefore the audio the codec tags them into — vary with
nothing but the Gumbel draws. That makes this the one place where "did the
caller's seed reach the audio" is a decidable question.
"""
talker, _codec = _make_talker(max_model_len=max_model_len)
engine = _MiniEngine(talker, max_model_len=max_model_len)
cfg = talker.config
talker.vocab_sizes = list(cfg.audio_head_levels)
talker.num_heads = cfg.num_audio_heads
talker.head0_vocab = cfg.head0_vocab_size
talker.stop_token = cfg.stop_token
talker.stop_threshold = cfg.stop_threshold
talker.temperature = cfg.temperature
talker._greedy = False
gather_idx, mask = GepardTalkerForConditionalGeneration._build_gather_mask(talker.vocab_sizes)
talker._cb_gather_idx = gather_idx
talker._cb_mask = mask
# Fixed weights: two engine runs must differ only by their RNG.
torch.manual_seed(0)
talker.fused_codebook_head = nn.Linear(talker.model.hidden, sum(talker.vocab_sizes))
talker.stop_head = nn.Linear(talker.model.hidden, 1)
# Never stop: these requests end on their token budget, so the frame count
# is fixed and only the sampled codes can differ.
torch.nn.init.constant_(talker.stop_head.bias, -20.0)
talker._audio_frame_embed = lambda h0, h: torch.zeros(int(h0.shape[0]), talker.model.hidden)
return talker, engine
# --------------------------------------------------------------------------
# The tail must reach the caller, whichever way the request ends.
# --------------------------------------------------------------------------
def test_last_in_flight_request_delivers_every_committed_frame() -> None:
"""The only request runs out of budget: there is no later step to flush in.
This is the case a deferred flush cannot serve at all — nothing runs after
the final forward.
"""
_talker, _codec, engine = _build()
engine.add_request("req-0", max_tokens=FIRST_CHUNK_FRAMES + CHUNK_FRAMES + 5)
engine.run()
frames = engine.committed_frames["req-0"]
assert frames > 0
torch.testing.assert_close(
engine.audio_for("req-0"),
_expected_samples(engine.reqs["req-0"]["tag_base"], frames),
)
@pytest.mark.parametrize(
"max_tokens",
[
1,
FIRST_CHUNK_FRAMES,
FIRST_CHUNK_FRAMES + 1,
30,
FIRST_CHUNK_FRAMES + CHUNK_FRAMES,
100,
150,
],
)
def test_budget_truncated_request_delivers_every_committed_frame(max_tokens: int) -> None:
"""Lowering max_tokens must shorten the clip, never silently cut it back
to the last cadence boundary."""
_talker, _codec, engine = _build()
engine.add_request("req-0", max_tokens=max_tokens)
engine.run()
frames = engine.committed_frames["req-0"]
assert frames == max_tokens, "every sampled step commits one frame when STOP never fires"
torch.testing.assert_close(
engine.audio_for("req-0"),
_expected_samples(engine.reqs["req-0"]["tag_base"], frames),
)
def test_stop_head_path_still_delivers_every_committed_frame() -> None:
"""Regression guard: the path that already worked must keep working."""
_talker, _codec, engine = _build()
engine.add_request("req-0", max_tokens=4000, stop_after=FIRST_CHUNK_FRAMES + 7)
engine.run()
frames = engine.committed_frames["req-0"]
assert frames == FIRST_CHUNK_FRAMES + 7, "the STOP frame itself is not committed"
torch.testing.assert_close(
engine.audio_for("req-0"),
_expected_samples(engine.reqs["req-0"]["tag_base"], frames),
)
def test_the_step_after_a_stop_token_commits_nothing(caplog: pytest.LogCaptureFixture) -> None:
"""Async scheduling runs one more step on a request that already stopped.
Committing its frame would strand it — nothing ships after the stop step —
and would make the truncation warning fire on a request that delivered
everything it produced, which is the one thing that warning must not do.
"""
talker, _codec, engine = _build()
stop_after = FIRST_CHUNK_FRAMES + 6
engine.add_request("req-0", max_tokens=4000, stop_after=stop_after)
target = logging.getLogger(LOGGER_NAME)
target.addHandler(caplog.handler)
prev = target.level
target.setLevel(logging.WARNING)
try:
engine.run()
# The cleanup that would warn runs on the step after the request
# leaves, and here there is none. Drive it the way the runner would.
talker.on_requests_finished({"req-0"})
talker._flush_deferred_cleanup()
finally:
target.removeHandler(caplog.handler)
target.setLevel(prev)
assert engine.reqs["req-0"]["ghost_step_pending"], "the harness did not run the extra step"
frames = engine.committed_frames["req-0"]
assert frames == stop_after, f"the extra step committed a frame: {frames} > {stop_after}"
torch.testing.assert_close(
engine.audio_for("req-0"),
_expected_samples(engine.reqs["req-0"]["tag_base"], frames),
)
assert not [m for m in caplog.messages if "undelivered frame" in m], (
f"a fully delivered request was reported as truncated; saw {caplog.messages}"
)
def test_a_resubmitted_request_id_keeps_its_new_state(caplog: pytest.LogCaptureFixture) -> None:
"""An aborted id can come back before its cleanup has run.
The runner treats the two as distinct requests, so the pending free must
not take the second one's state: without that, the state is popped at the
end of the very forward that created it and every later decode falls back
to the zero-embedding path — the request keeps producing frames, so only
the audio is wrong.
"""
_talker, _codec, engine = _build()
engine.add_request("req-0", max_tokens=4000)
for _ in range(FIRST_CHUNK_FRAMES + 2):
engine.step()
engine.abort("req-0")
# Same id, resubmitted before the model has been told about the abort.
engine.add_request("req-0", max_tokens=FIRST_CHUNK_FRAMES + CHUNK_FRAMES + 5)
target = logging.getLogger(LOGGER_NAME)
target.addHandler(caplog.handler)
prev = target.level
target.setLevel(logging.WARNING)
try:
engine.run()
finally:
target.removeHandler(caplog.handler)
target.setLevel(prev)
frames = engine.committed_frames["req-0"]
assert frames == engine.reqs["req-0"]["max_tokens"], "the second request lost frames"
torch.testing.assert_close(
engine.audio_for("req-0"),
_expected_samples(engine.reqs["req-0"]["tag_base"], frames),
msg="the second request did not receive its own audio",
)
# The abort still gets reported: the cleanup that would have said so is the
# one this cancelled.
assert any("resubmitted" in m for m in caplog.messages), (
f"the aborted request's truncation went unreported; saw {caplog.messages}"
)
def test_a_preempted_request_is_reported_rather_than_resumed(caplog: pytest.LogCaptureFixture) -> None:
"""Recompute preemption re-prefills a request that was mid-generation.
This model cannot survive that, and the point of the test is that it says
so. The recomputed ids are the prompt plus the head0 codes sampled so far,
which go through the text embedding table, and each frame's other 31 codes
are never in the token stream at all — there is nothing to stitch the two
halves back together with, so generation restarts from frame 0.
The harness re-prefills with the bare prompt, which understates the damage.
What is under test is the report and the reset, not the corruption, which
needs a real backbone to observe.
"""
_talker, _codec, engine = _build()
engine.add_request("req-0", max_tokens=4000)
for _ in range(FIRST_CHUNK_FRAMES + 2):
engine.step()
assert engine.committed_frames["req-0"] > 0
# Preemption the way the scheduler does it: computed tokens back to 0, so
# the next step re-prefills, and no id reaches on_requests_finished.
engine.reqs["req-0"]["computed"] = 0
target = logging.getLogger(LOGGER_NAME)
target.addHandler(caplog.handler)
prev = target.level
target.setLevel(logging.WARNING)
try:
engine.step()
finally:
target.removeHandler(caplog.handler)
target.setLevel(prev)
assert any("preempted" in m for m in caplog.messages), (
f"a preempted request restarted its audio silently; saw {caplog.messages}"
)
assert not any("resubmitted" in m for m in caplog.messages), (
"a preemption was reported as an id resubmit; the two need different words"
)
assert engine.committed_frames["req-0"] == 1, "the re-prefill must start the frame history over"
def test_concurrent_requests_each_receive_their_own_full_audio() -> None:
"""One request finishes while another stays live.
The earlier finisher is the case where a deferred flush *does* run but the
payload is filtered out again, because its id has left the output batch.
"""
_talker, _codec, engine = _build()
engine.add_request("req-short", max_tokens=FIRST_CHUNK_FRAMES + 3)
engine.add_request("req-long", max_tokens=FIRST_CHUNK_FRAMES + CHUNK_FRAMES + 11)
engine.run()
for req_id in ("req-short", "req-long"):
frames = engine.committed_frames[req_id]
assert frames == engine.reqs[req_id]["max_tokens"]
torch.testing.assert_close(
engine.audio_for(req_id),
_expected_samples(engine.reqs[req_id]["tag_base"], frames),
msg=f"{req_id} lost or gained samples",
)
def test_codec_still_runs_once_per_chunk_not_once_per_frame() -> None:
"""Emitting on the final step must not degrade into per-frame decoding."""
_talker, codec, engine = _build()
num_frames = FIRST_CHUNK_FRAMES + 3 * CHUNK_FRAMES
engine.add_request("req-0", max_tokens=num_frames)
engine.run()
assert codec.decode_calls == 4, f"expected one decode per chunk, got {codec.decode_calls}"
# --------------------------------------------------------------------------
# Deferred cleanup is cleanup: it must not try to produce output.
# --------------------------------------------------------------------------
def test_deferred_cleanup_never_queues_audio() -> None:
"""Anything queued from the cleanup hook is unroutable by construction.
By the time it runs the id has left ``req_ids_output_copy``, so a payload
under it is dropped; queueing there only risks leaking one request's audio
into the next step's payload for a different request.
"""
talker, _codec, engine = _build()
engine.add_request("req-0", max_tokens=FIRST_CHUNK_FRAMES + 4)
engine.add_request("req-1", max_tokens=FIRST_CHUNK_FRAMES + CHUNK_FRAMES)
while engine.step():
if engine.reqs["req-0"]["finished"] and not engine.reqs["req-1"]["finished"]:
break
talker._audio_queue.clear()
talker.on_requests_finished({"req-0"})
talker._flush_deferred_cleanup()
assert talker._audio_queue == [], "cleanup emitted audio that routing cannot deliver"
assert "req-0" not in talker._active_states, "cleanup must still free the state"
def test_cleanup_warns_when_frames_are_left_undelivered(caplog: pytest.LogCaptureFixture) -> None:
"""An abort is the one case with no last step to emit on. Say so."""
_talker, _codec, engine = _build()
engine.add_request("req-0", max_tokens=4000)
engine.add_request("req-1", max_tokens=4000)
for _ in range(FIRST_CHUNK_FRAMES + 3):
engine.step()
engine.abort("req-0")
target = logging.getLogger(LOGGER_NAME)
target.addHandler(caplog.handler)
prev = target.level
target.setLevel(logging.WARNING)
try:
engine.step() # cleanup for req-0 runs at the start of this step
finally:
target.removeHandler(caplog.handler)
target.setLevel(prev)
assert any("undelivered frame" in m for m in caplog.messages), (
f"aborted request dropped audio without a warning; saw {caplog.messages}"
)
# --------------------------------------------------------------------------
# A request's seed must reach the in-model sampler.
# --------------------------------------------------------------------------
_SEEDED_FRAMES = FIRST_CHUNK_FRAMES + 5
def _run_seeded(seed: int | None, *, neighbour: bool = False, neighbour_seed: int | None = None) -> torch.Tensor:
"""Audio delivered to ``req-a``, optionally sharing the batch with another."""
_talker, engine = _build_with_real_sampler()
engine.add_request("req-a", max_tokens=_SEEDED_FRAMES, seed=seed)
if neighbour:
engine.add_request("req-b", max_tokens=_SEEDED_FRAMES, seed=neighbour_seed)
engine.run()
return engine.audio_for("req-a")
def test_a_seeded_request_is_reproducible_within_one_process() -> None:
"""The bug this pins: the 32 heads are sampled in ``forward``, so a seed
that only reaches vLLM's sampler leaves the audio at the mercy of the
global RNG — two identical requests in one process then diverge."""
torch.testing.assert_close(_run_seeded(1234), _run_seeded(1234))
def test_different_seeds_produce_different_audio() -> None:
"""Guards the opposite failure: a seed that is threaded but never used."""
assert not torch.equal(_run_seeded(1234), _run_seeded(4321))
def test_a_seeded_request_is_unaffected_by_its_batch() -> None:
"""Each seeded request draws its own noise, so the noise it sees cannot
depend on who else was scheduled alongside it.
The stub backbone is what makes that checkable. A real one is a bf16
batched matmul whose rows are not reproducible across batch shapes, so
end-to-end batch invariance is not attainable there; this pins the half
the model actually controls.
"""
torch.testing.assert_close(_run_seeded(1234), _run_seeded(1234, neighbour=True, neighbour_seed=999))
def test_an_unseeded_neighbour_does_not_disturb_a_seeded_request() -> None:
"""The mixed case: one row draws from the global RNG, the other must not."""
torch.testing.assert_close(_run_seeded(1234), _run_seeded(1234, neighbour=True))
# --------------------------------------------------------------------------
# Budget arithmetic, straight through the helper.
# --------------------------------------------------------------------------
def _info(**kwargs) -> _ReqInfo:
base = {
"request_id": "req-0",
"is_prefill": False,
"prompt_len": 20,
"num_computed_tokens": 20,
"max_tokens": 10,
"seed": None,
}
base.update(kwargs)
return _ReqInfo(**base)
def test_prefill_step_is_the_first_output_token() -> None:
talker, _codec, _engine = _build()
first = _info(is_prefill=True, num_computed_tokens=0, max_tokens=1)
assert talker._is_last_output_token(first, span=20, samples_frame=True) is True
more = _info(is_prefill=True, num_computed_tokens=0, max_tokens=2)
assert talker._is_last_output_token(more, span=20, samples_frame=True) is False
def test_partial_prefill_chunk_does_not_consume_budget() -> None:
"""vLLM discards that row's token, so it must not count against max_tokens."""
talker, _codec, _engine = _build()
info = _info(is_prefill=True, num_computed_tokens=0, prompt_len=40, max_tokens=1)
assert talker._is_last_output_token(info, span=10, samples_frame=False) is False
def test_budget_is_capped_by_max_model_len() -> None:
"""A max_tokens larger than the context window never gets reached."""
talker, _codec, _engine = _build(max_model_len=30)
# prompt 20 + 10 generated fills the window, whatever max_tokens says.
at_cap = _info(num_computed_tokens=28, max_tokens=1000)
assert talker._is_last_output_token(at_cap, span=1, samples_frame=True) is True
below_cap = _info(num_computed_tokens=27, max_tokens=1000)
assert talker._is_last_output_token(below_cap, span=1, samples_frame=True) is False
def test_missing_budget_is_not_treated_as_exhausted() -> None:
"""A runner that does not thread max_tokens must not trigger a fade."""
talker, _codec, _engine = _build(max_model_len=0)
assert talker._is_last_output_token(_info(max_tokens=None), span=1, samples_frame=True) is False
+52
View File
@@ -20,6 +20,7 @@ from vllm_omni.model_extras import (
get_output_tensor_range,
get_x_to_text_model_family,
should_init_extra_args_for_non_diffusion_stages,
should_preserve_reference_image_size,
)
@@ -230,18 +231,62 @@ def test_ltx_extra_registry_declares_official_guidance_params() -> None:
"audio_rescale_scale",
"video_stg_blocks",
"audio_stg_blocks",
"sigmas",
"stage_1_sigmas",
"stage_2_sigmas",
"image_crf",
}
)
for pipeline_name in (
"LTX2Pipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledOneStagePipeline",
"LTX2DistilledTwoStagePipeline",
"LTX2DistilledPipeline",
):
assert get_extra_body_params(pipeline_name) == expected
assert get_extra_output_params(pipeline_name) == frozenset()
@pytest.mark.parametrize(
("model_version", "expected"),
[("2", False), ("2.3", False), ("2.5", True)],
)
def test_ltx_reference_image_size_policy(tmp_path, model_version: str, expected: bool) -> None:
(tmp_path / "model_index.json").write_text(
'{"_class_name": "LTX2Pipeline", "model_version": "' + model_version + '"}'
)
assert (
should_preserve_reference_image_size(
None,
model=str(tmp_path),
)
is expected
)
def test_reference_image_size_policy_threads_revision(monkeypatch: pytest.MonkeyPatch) -> None:
captured = {}
def fake_policy(*, model, revision=None):
captured.update(model=model, revision=revision)
return True
monkeypatch.setattr(
"vllm_omni.diffusion.models.ltx2.ltx2_components.preserves_reference_image_size",
fake_policy,
)
assert should_preserve_reference_image_size(
"LTX2Pipeline",
model="org/model",
revision="pinned-revision",
)
assert captured == {"model": "org/model", "revision": "pinned-revision"}
@pytest.mark.core_model
@pytest.mark.cpu
def test_ming_flash_omni_extra_registry_declares_request_and_response_params() -> None:
@@ -424,6 +469,13 @@ def test_vace_extra_registry_has_no_pipeline_params() -> None:
def test_unknown_pipeline_has_empty_extra_registry() -> None:
assert get_extra_body_params("UnknownPipeline") == frozenset()
assert get_extra_output_params("UnknownPipeline") == frozenset()
assert (
should_preserve_reference_image_size(
"UnknownPipeline",
model="org/model",
)
is False
)
assert should_init_extra_args_for_non_diffusion_stages("UnknownPipeline") is False
@@ -28,6 +28,8 @@ EXCLUDED_MODELS = [
"WanPipeline",
"WanVACEPipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledOneStagePipeline",
"LTX2DistilledTwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
+2 -3
View File
@@ -159,11 +159,10 @@ def _run_once(model: str, yaml_path: str, *, async_mode: bool) -> list[dict]:
wav: bytes = b""
# Batch size 1: submit a single prompt so scheduling is deterministic.
for step in omni.generate([_text_prompt(question)], _sampling_params(), py_generator=False):
out = step.request_output
if step.final_output_type == "text":
text = out.outputs[0].text # exact text, no strip -> truly bit-identical
text = step.outputs[0].text # exact text, no strip -> truly bit-identical
elif step.final_output_type == "audio":
wav = _audio_bytes(out.outputs[0].multimodal_output["audio"])
wav = _audio_bytes(step.outputs[0].multimodal_output["audio"])
if text is None:
raise RuntimeError(f"prompt_id={prompt_id} async={async_mode}: produced no text output")
if not wav:
+4
View File
@@ -19,6 +19,10 @@ ALLOWED_FILES = {
"tests/helpers/process.py",
"vllm_omni/diffusion/distributed/group_coordinator.py",
"tests/diffusion/attention/test_attention_sp.py",
# Measures the wire size of diffusion outputs. The transport pickles them,
# and only pickle reproduces the per-view storage duplication under test
# (torch.save deduplicates shared storages). No untrusted data is loaded.
"tests/diffusion/test_diffusion_ipc.py",
}
PICKLE_RE = re.compile(
+43 -5
View File
@@ -15,6 +15,7 @@ from transformers import PretrainedConfig
from vllm.logger import init_logger
from vllm.transformers_utils.config import get_config
from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
from vllm.transformers_utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
from vllm_omni.config.endpoint_policy import EndpointRestriction
from vllm_omni.config.omni_config import VllmOmniConfig
@@ -45,6 +46,39 @@ logger = init_logger(__name__)
_DEFAULT_PARALLEL_DEGREE = 1
@functools.cache
def _materialize_object_storage_configs(model: str) -> str:
"""Materialize an object-storage model URI's config files locally.
vLLM's Run:AI streamer keeps ``s3://``/``gs://``/``az://`` URIs opaque until
each stage builds its ``ModelConfig``; parent-process resolution (HF config
lookup, pipeline/pipeline-key matching) would instead hand the URI to
``huggingface_hub`` helpers, which reject it with ``HFValidationError``.
Pull the lightweight files once into vLLM's deterministic
``model_streamer/<hash>`` directory so config reads work here, and so the
stage processes' own pull lands in that same directory.
Returns the input unchanged for non object-storage paths.
"""
if not is_runai_obj_uri(model):
return model
object_storage_model = ObjectStorageModel(url=model)
object_storage_model.pull_files(model, allow_pattern=["*.model", "*.py", "*.json"])
logger.info("Materialized object-storage configs for %s at %s", model, object_storage_model.dir)
return object_storage_model.dir
def _name_match_candidate(model: str) -> str:
"""Last path component of a model reference, used for name-based matching.
Object-storage URIs and HF repo ids carry non-model segments (bucket name,
organization) that must not participate in substring matching; e.g. a
bucket named ``qwen3-tts-models`` holding a ``Qwen3-Omni`` checkpoint must
not resolve to the ``qwen3_tts`` pipeline.
"""
return model.rstrip("/").rsplit("/", 1)[-1]
def with_trust_remote_code_override(
overrides: Mapping[str, Any],
trust_remote_code: bool | None,
@@ -117,7 +151,7 @@ class StageConfigFactory:
"""
hf_config = None
try:
return get_config(model, trust_remote_code=trust_remote_code)
return get_config(_materialize_object_storage_configs(model), trust_remote_code=trust_remote_code)
except Exception as e:
logger.debug(f"`get_config` failed with exception {e}; inferred HF config is None")
return hf_config
@@ -163,10 +197,12 @@ class StageConfigFactory:
if hf_config is not None:
return hf_config.model_type
config_source = _materialize_object_storage_configs(model)
# Fallback: read config.json directly for custom model types that
# are not registered with transformers (e.g. qwen3_tts).
try:
config_dict = get_hf_file_to_dict("config.json", model, revision=None)
config_dict = get_hf_file_to_dict("config.json", config_source, revision=None)
if config_dict:
if "model_type" in config_dict:
return config_dict["model_type"]
@@ -183,7 +219,7 @@ class StageConfigFactory:
# model_index.json with _class_name that maps to a pipeline key via
# PipelineConfig.diffusers_class_name.
try:
model_index = get_hf_file_to_dict("model_index.json", model, revision=None)
model_index = get_hf_file_to_dict("model_index.json", config_source, revision=None)
if model_index and "_class_name" in model_index:
class_name = model_index["_class_name"]
for obj in OMNI_PIPELINES.values():
@@ -203,8 +239,10 @@ class StageConfigFactory:
# Final fallback: some models (e.g. CosyVoice3) ship an empty
# config.json and rely on naming conventions. Match the model path
# basename against registered pipeline keys — longest match wins
# so "cosyvoice3" (length 10) beats "cosyvoice" (length 9).
model_lower = model.lower().replace("-", "").replace("_", "")
# so "cosyvoice3" (length 10) beats "cosyvoice" (length 9). Only
# the basename is scanned so URI segments such as the bucket name
# cannot select an unrelated pipeline.
model_lower = _name_match_candidate(model).lower().replace("-", "").replace("_", "")
best: str | None = None
best_len = 0
for registered_key in OMNI_PIPELINES.keys():
+2
View File
@@ -51,6 +51,7 @@ from vllm_omni.model_executor.models.dots_tts.pipeline import DOTS_TTS_PIPELINE
from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE
from vllm_omni.model_executor.models.dynin_omni.pipeline import DYNIN_OMNI_PIPELINE
from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE
from vllm_omni.model_executor.models.gepard.pipeline import GEPARD_PIPELINE
from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE
from vllm_omni.model_executor.models.glm_tts.pipeline import GLM_TTS_PIPELINE
from vllm_omni.model_executor.models.gr00t.pipeline import GR00T_N1D7_PIPELINE
@@ -138,6 +139,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
"lance": LANCE_PIPELINE,
"dreamzero": DREAMZERO_PIPELINE,
"Gr00tN1d7": GR00T_N1D7_PIPELINE,
"gepard": GEPARD_PIPELINE,
"glm_image": GLM_IMAGE_PIPELINE,
"hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE,
"hunyuan_image3_ar": HUNYUAN_IMAGE3_AR_PIPELINE,
+40
View File
@@ -0,0 +1,40 @@
# Gepard-1.0: single-stage native-AR FSQ/NanoCodec TTS (~0.5B).
# Target 1x L4 24GB. Zero-shot only; enforce_eager until the perf PR.
#
# The pipeline pin is required, not cosmetic: the checkpoint's config.json
# self-identifies as qwen3_5_text, so model_type inference cannot reach the
# gepard registry entry, and the architectures fallback would match
# Qwen3_5ForCausalLM against the diffusion pipeline.
pipeline: gepard
async_chunk: false
dtype: bfloat16
trust_remote_code: true
stages:
- stage_id: 0
max_num_seqs: 4
gpu_memory_utilization: 0.3 # AR LM + NanoCodec fit in ~2 GiB
enforce_eager: true
enable_prefix_caching: false
async_scheduling: true
max_num_batched_tokens: 4096
max_model_len: 4096
devices: "0"
# Override architectures so vLLM resolves our talker rather than the native
# Qwen3.5 text model. Do NOT override model_type: the backbone's own
# model_type check needs it to stay qwen3_5_text.
hf_overrides:
architectures: ["GepardTalkerForConditionalGeneration"]
default_sampling_params:
# All 32 heads are sampled in-model, so vLLM's sampler is pinned
# deterministic and just argmaxes what forward chose.
temperature: 0.0
top_p: 1.0
top_k: -1
# One output token is one frame, so this is the frame budget: 1000 frames
# is ~46 s at 21.5 fps. A backstop, not a target -- normal utterances are
# 40-420 frames -- but a load-bearing one, because this checkpoint's stop
# head does not fire on every batched trajectory.
max_tokens: 1000
seed: 42
detokenize: false
Loaded 100 of 165 files, more files were not shown because too many files have changed in this diff. Show more