Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7ec78360bb | ||
|
|
5569abf0a2 | ||
|
|
10853fb4f3 | ||
|
|
a52c19ea8b | ||
|
|
8b61499662 | ||
|
|
7872663bf0 |
No files matched your search
@@ -76,9 +76,6 @@ 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:
|
||||
@@ -95,9 +92,6 @@ 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
|
||||
@@ -115,15 +109,27 @@ 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
|
||||
@@ -135,11 +141,6 @@ 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:
|
||||
@@ -161,35 +162,70 @@ 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 '
|
||||
@@ -198,12 +234,23 @@ 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 "
|
||||
@@ -211,3 +258,21 @@ 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
|
||||
@@ -459,12 +459,6 @@ 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:
|
||||
@@ -493,13 +487,6 @@ 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
|
||||
|
||||
@@ -169,10 +169,6 @@ 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:
|
||||
@@ -185,11 +181,6 @@ 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:
|
||||
- |
|
||||
@@ -201,11 +192,6 @@ 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:
|
||||
- |
|
||||
@@ -224,24 +210,16 @@ 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 "
|
||||
@@ -261,13 +239,25 @@ 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:
|
||||
@@ -282,9 +272,6 @@ 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:
|
||||
@@ -297,11 +284,52 @@ 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 "
|
||||
@@ -314,9 +342,6 @@ 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 "
|
||||
@@ -329,8 +354,6 @@ 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 "
|
||||
@@ -338,3 +361,15 @@ 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
|
||||
@@ -108,22 +108,6 @@ 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:
|
||||
|
||||
@@ -254,7 +254,6 @@ outputs/
|
||||
results/
|
||||
generated/
|
||||
output_*/
|
||||
/tests/e2e/accuracy/artifacts/
|
||||
|
||||
# Configuration overrides
|
||||
configs/local.yaml
|
||||
|
||||
@@ -329,13 +329,11 @@ On PR builds, `upload_pipeline.py` logs `skip '…' (no changes under …)` for
|
||||
|
||||
### Per-model 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.
|
||||
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.
|
||||
|
||||
**Naming convention**: `coverage-<model_id>-<mode>-<step_id>.xml.gz`, where
|
||||
`<model_id>` is the model's directory name under
|
||||
@@ -412,8 +410,7 @@ 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` covers the filter behavior with
|
||||
a synthetic job (it does not pin real merge labels).
|
||||
nightly. `tests/buildkite/test_upload_pipeline.py` asserts this for the pilots.
|
||||
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
|
||||
|
||||
@@ -117,22 +117,7 @@ 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. 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.
|
||||
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.
|
||||
|
||||
## Model Coverage
|
||||
|
||||
@@ -146,7 +131,6 @@ 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
|
||||
complete rank-local block sources and avoid an additional collective.
|
||||
the standard loader's rank-local weights and avoid an additional collective.
|
||||
|
||||
For user-facing commands, see the
|
||||
[distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
|
||||
@@ -14,20 +14,10 @@ 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 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.
|
||||
`--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.
|
||||
|
||||
The compatibility matrix below describes the current implementation. The
|
||||
unit-level guards are covered, but not every parallelism combination has a
|
||||
@@ -51,36 +41,6 @@ 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
|
||||
@@ -101,15 +61,6 @@ 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.
|
||||
@@ -117,20 +68,9 @@ 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 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.
|
||||
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.
|
||||
|
||||
This mode means:
|
||||
|
||||
@@ -140,18 +80,21 @@ 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 share one checkpoint-backed copy per node when direct
|
||||
mmap is selected; the ordinary-loader fallback keeps one private runtime
|
||||
copy per rank.
|
||||
- Pure DP deployments keep a full host-side model copy per rank, subject to
|
||||
shared OS page-cache behavior.
|
||||
- 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. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. |
|
||||
| **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. |
|
||||
| **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** | 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. |
|
||||
| **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. |
|
||||
| **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
|
||||
@@ -179,15 +122,11 @@ AllGather DP multi-concurrency requires:
|
||||
The no-AllGather path does not impose these DLO-specific synchronized-wave
|
||||
requirements.
|
||||
|
||||
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).
|
||||
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.
|
||||
|
||||
## Validation coverage
|
||||
|
||||
@@ -195,75 +134,15 @@ Current source-level validation includes:
|
||||
|
||||
- HSDP + DLO + AllGather rejection;
|
||||
- HSDP + DLO without AllGather acceptance at configuration level;
|
||||
- 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;
|
||||
- TP rejection in the DLO+AllGather mmap path;
|
||||
- resident-layer requests requiring no-AllGather;
|
||||
- DP request-wave validation for denoising-step compatibility;
|
||||
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
|
||||
tests.
|
||||
|
||||
### 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.
|
||||
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.
|
||||
|
||||
## Recommendations
|
||||
|
||||
@@ -271,9 +150,7 @@ change the Phase A direct-mmap TP1 support boundary.
|
||||
scaling path.
|
||||
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
|
||||
not the goal.
|
||||
- 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.
|
||||
- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
|
||||
the expectation of higher host-memory use and lower validation confidence.
|
||||
- Prefer **HSDP alone** for production HSDP deployments until the combined
|
||||
HSDP + DLO no-AllGather path has broader end-to-end coverage.
|
||||
@@ -76,12 +76,6 @@ 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/
|
||||
@@ -160,11 +154,9 @@ 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. 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.
|
||||
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.
|
||||
|
||||
Weight quantization and quantized KV cache are separate paths. The former is
|
||||
resolved through `quantization_config`; runtime attention/KV-cache quantization
|
||||
|
||||
@@ -101,7 +101,7 @@ outputs = omni.generate(
|
||||
)
|
||||
|
||||
# Access custom trajectory data
|
||||
output = outputs[0]
|
||||
output = outputs[0].request_output
|
||||
print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
|
||||
print(f"Trajectory latents shape: {output.latents.shape}")
|
||||
```
|
||||
|
||||
@@ -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].images
|
||||
images = outputs[0].request_output.images
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
@@ -75,7 +75,8 @@ if __name__ == "__main__":
|
||||
]
|
||||
omni_outputs = omni.generate(prompts)
|
||||
for i_prompt, prompt_output in enumerate(omni_outputs):
|
||||
this_images = prompt_output.images
|
||||
this_request_output = prompt_output.request_output
|
||||
this_images = this_request_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")
|
||||
|
||||
@@ -43,12 +43,7 @@ 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 | ✅︎ | ✅︎ | | | — |
|
||||
| `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) |
|
||||
| `LTX2DistilledPipeline` | LTX-2 / LTX-2.3 full-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + `Lightricks/LTX-2.3` upsampler | ✅︎ | ✅︎ | | | — |
|
||||
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
|
||||
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
|
||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
@@ -103,7 +98,6 @@ 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,9 +3,7 @@
|
||||
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 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.
|
||||
each rank streams its standard-loader rank-local weights independently.
|
||||
|
||||
See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
|
||||
for the implementation contract and compatibility matrix.
|
||||
@@ -61,48 +59,21 @@ 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` | Stream complete rank-local blocks without a DLO weight collective | `false` |
|
||||
| `--dlo-no-use-allgather` | Keep standard-loader rank-local weights | `false` |
|
||||
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
|
||||
|
||||
## Host-weight loading
|
||||
## mmap weight loading
|
||||
|
||||
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:
|
||||
The DLO plus AllGather path:
|
||||
|
||||
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. 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.
|
||||
4. calls model-specific `post_load_weights()` conversions; and
|
||||
5. restores the saved non-persistent buffers.
|
||||
|
||||
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.
|
||||
This avoids one full checkpoint RSS copy per rank and does not require
|
||||
model-specific loading code.
|
||||
|
||||
## Declarative topology
|
||||
|
||||
@@ -132,22 +103,16 @@ must enter each collective.
|
||||
|
||||
## Limitations
|
||||
|
||||
- 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.
|
||||
- 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.
|
||||
- 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.
|
||||
@@ -161,7 +161,6 @@ 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** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
@@ -175,7 +174,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, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni
|
||||
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, Helios, HunyuanVideo-1.5, DreamID-Omni
|
||||
|
||||
### AudioGen
|
||||
|
||||
|
||||
@@ -43,7 +43,7 @@ if __name__ == "__main__":
|
||||
),
|
||||
],
|
||||
)
|
||||
outputs[0].images[0].save("output.png")
|
||||
outputs[0].request_output.images[0].save("output.png")
|
||||
```
|
||||
|
||||
## Image-to-Image (Image Editing)
|
||||
@@ -72,7 +72,7 @@ if __name__ == "__main__":
|
||||
),
|
||||
],
|
||||
)
|
||||
outputs[0].images[0].save("output.png")
|
||||
outputs[0].request_output.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].images
|
||||
images = outputs[0].request_output.images
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
@@ -179,7 +179,7 @@ if __name__ == "__main__":
|
||||
]
|
||||
outputs = omni.generate(prompts)
|
||||
for i, output in enumerate(outputs):
|
||||
output.images[0].save(f"{i}.jpg")
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
!!! info
|
||||
@@ -207,7 +207,7 @@ if __name__ == "__main__":
|
||||
}
|
||||
])
|
||||
for i, output in enumerate(outputs):
|
||||
output.images[0].save(f"{i}.jpg")
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
|
||||
|
||||
@@ -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` (`FL2VA` / `Ref2VA`) | Yes | Not validated | `quantization="fp8"` quantizes eligible DiT and text-encoder linears; mixed-precision input/output heads stay FP32 | None | ✅︎ |
|
||||
| MiniMax H3 | `MiniMaxAI/MiniMax-H3` | Yes | Not validated | DiT linears except patch, timestep, and final projections; not compatible with layerwise offload | None | |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
@@ -138,39 +138,6 @@ 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].images
|
||||
frames = outputs[0].request_output.images
|
||||
export_to_video(frames, "quick_test_i2v.mp4", fps=16)
|
||||
```
|
||||
|
||||
|
||||
@@ -68,7 +68,6 @@ 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
|
||||
@@ -455,13 +454,7 @@ def main():
|
||||
|
||||
media_inputs: dict[str, Any] = {}
|
||||
if image is not None:
|
||||
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)
|
||||
)
|
||||
media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
if last_image is not None:
|
||||
media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
if mask_image is not None:
|
||||
|
||||
@@ -335,7 +335,7 @@ def main(args):
|
||||
print(f"use_tts (tts_bos): {use_tts}")
|
||||
|
||||
for stage_outputs in omni_generator:
|
||||
output = stage_outputs
|
||||
output = stage_outputs.request_output
|
||||
if stage_outputs.final_output_type == "text":
|
||||
request_id = output.request_id
|
||||
text_output = output.outputs[0].text
|
||||
|
||||
+25
-60
@@ -1,32 +1,21 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""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.
|
||||
"""Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior.
|
||||
|
||||
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_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode dlo-no-allgather --dp-size 2
|
||||
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
|
||||
CUDA_VISIBLE_DEVICES=0,1 \
|
||||
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode request --tp-size 2
|
||||
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode request
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -51,33 +40,14 @@ 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", "dlo-no-allgather", "request"),
|
||||
choices=("dlo-dp2", "request"),
|
||||
required=True,
|
||||
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)",
|
||||
help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
|
||||
)
|
||||
parser.add_argument("--steps", type=int, default=2)
|
||||
parser.add_argument("--duration", type=float, default=5.0)
|
||||
@@ -90,11 +60,10 @@ 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": args.dp_size if is_dlo else args.tp_size,
|
||||
"num_gpus": 2,
|
||||
"usp": 1,
|
||||
"ring": 1,
|
||||
"vae_parallel_mode": "tile",
|
||||
@@ -105,22 +74,22 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
|
||||
"stage_init_timeout": args.init_timeout,
|
||||
"init_timeout": args.init_timeout,
|
||||
}
|
||||
if is_dlo:
|
||||
if args.mode == "dlo-dp2":
|
||||
common.update(
|
||||
tensor_parallel_size=1,
|
||||
data_parallel_size=args.dp_size,
|
||||
data_parallel_size=2,
|
||||
text_encoder_tp_size=1,
|
||||
vae_patch_parallel_size=1,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=args.mode == "dlo",
|
||||
dlo_use_allgather=True,
|
||||
dlo_resident_layers=0,
|
||||
)
|
||||
else:
|
||||
common.update(
|
||||
tensor_parallel_size=args.tp_size,
|
||||
tensor_parallel_size=2,
|
||||
data_parallel_size=1,
|
||||
text_encoder_tp_size=args.tp_size,
|
||||
vae_patch_parallel_size=args.tp_size,
|
||||
text_encoder_tp_size=2,
|
||||
vae_patch_parallel_size=2,
|
||||
enable_distributed_layerwise_offload=False,
|
||||
)
|
||||
return common
|
||||
@@ -203,8 +172,7 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
|
||||
"engine_kwargs": engine_kwargs(args),
|
||||
}
|
||||
try:
|
||||
uses_dp_wave = args.mode == "dlo" and args.dp_size > 1
|
||||
if uses_dp_wave:
|
||||
if args.mode == "dlo-dp2":
|
||||
started = time.perf_counter()
|
||||
asymmetric = await asyncio.gather(
|
||||
generate_one(
|
||||
@@ -214,15 +182,12 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
|
||||
prompt="",
|
||||
seed=1001,
|
||||
),
|
||||
*(
|
||||
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)
|
||||
generate_one(
|
||||
engine,
|
||||
args,
|
||||
request_id="invalid-peer",
|
||||
prompt=DEFAULT_PROMPTS[0],
|
||||
seed=1002,
|
||||
),
|
||||
return_exceptions=True,
|
||||
)
|
||||
@@ -230,18 +195,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"]) != args.dp_size:
|
||||
raise RuntimeError("Every request in the asymmetric DP wave must fail before dispatch")
|
||||
if len(summary["asymmetric_errors"]) != 2:
|
||||
raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch")
|
||||
|
||||
started = time.perf_counter()
|
||||
request_count = args.dp_size if uses_dp_wave else 1
|
||||
request_count = 2 if args.mode == "dlo-dp2" else 1
|
||||
outputs = await asyncio.gather(
|
||||
*(
|
||||
generate_one(
|
||||
engine,
|
||||
args,
|
||||
request_id=f"recovery-{index}",
|
||||
prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
|
||||
prompt=DEFAULT_PROMPTS[index],
|
||||
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].images
|
||||
images = outputs[0].request_output.images
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
@@ -265,7 +265,7 @@ if __name__ == "__main__":
|
||||
]
|
||||
outputs = omni.generate(prompts)
|
||||
for i, output in enumerate(outputs):
|
||||
output.images[0].save(f"{i}.jpg")
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
!!! info
|
||||
@@ -301,7 +301,7 @@ if __name__ == "__main__":
|
||||
}
|
||||
])
|
||||
for i, output in enumerate(outputs):
|
||||
output.images[0].save(f"{i}.jpg")
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
|
||||
|
||||
@@ -16,7 +16,6 @@ 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 |
|
||||
@@ -396,58 +395,6 @@ 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`.
|
||||
|
||||
@@ -1,139 +0,0 @@
|
||||
# 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,8 +570,6 @@ 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
|
||||
|
||||
|
||||
@@ -1,175 +0,0 @@
|
||||
# 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)
|
||||
+9
-14
@@ -8,12 +8,10 @@
|
||||
|---|---|---|
|
||||
| `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` |
|
||||
| `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` |
|
||||
| `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` |
|
||||
| `LTX2DistilledPipeline` | LTX-2 full-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
|
||||
| `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` |
|
||||
| `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers`<br>`Lightricks/LTX-2.3` |
|
||||
| `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` |
|
||||
| `LTX2DistilledPipeline` | LTX-2.3 full-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |
|
||||
|
||||
Repositories in the table are download units. A full pipeline repository
|
||||
contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler,
|
||||
@@ -28,12 +26,9 @@ 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. `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`.
|
||||
then refines with the distilled LoRA. `LTX2DistilledPipeline` uses a fully
|
||||
merged distilled Transformer in both stages. Both entries support T2V and I2V;
|
||||
select their class explicitly.
|
||||
|
||||
## API Migration
|
||||
|
||||
@@ -57,8 +52,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` | `LTX2DistilledTwoStagePipeline` |
|
||||
| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` with `image=` |
|
||||
| `LTX2TwoStagesPipeline` | `LTX2DistilledPipeline` |
|
||||
| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledPipeline` with `image=` |
|
||||
|
||||
Passing any second positional argument now raises `TypeError`. These changes
|
||||
affect direct Python callers and explicit `--model-class-name` overrides;
|
||||
@@ -117,10 +112,10 @@ Start a two-stage checkpoint with its explicit class:
|
||||
|
||||
```bash
|
||||
vllm serve rootonchair/LTX-2-19b-distilled --omni \
|
||||
--model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
|
||||
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
|
||||
# LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent
|
||||
vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \
|
||||
--model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
|
||||
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
|
||||
# Ordinary LTX-2.3 two-stage
|
||||
vllm serve diffusers/LTX-2.3-Diffusers --omni \
|
||||
--model-class-name LTX2TwoStagePipeline \
|
||||
|
||||
@@ -1,195 +0,0 @@
|
||||
# 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)
|
||||
@@ -357,19 +357,12 @@ No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while
|
||||
|
||||
### Online FP8 quantization
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
Add this option to an existing H3 server command:
|
||||
|
||||
@@ -377,28 +370,6 @@ 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
|
||||
@@ -415,10 +386,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 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.
|
||||
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.
|
||||
|
||||
## AMD ROCm (gfx942 / gfx950)
|
||||
|
||||
@@ -779,13 +750,6 @@ 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
|
||||
@@ -870,8 +834,7 @@ 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.
|
||||
- Runtime-created FP8 weights do not support the sharded DLO AllGather path;
|
||||
use `--dlo-no-use-allgather` when combining online FP8 with DLO.
|
||||
- Online FP8 is not compatible with layerwise offload.
|
||||
- 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
|
||||
|
||||
@@ -55,7 +55,6 @@ 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 |
|
||||
|
||||
@@ -1,15 +1,15 @@
|
||||
# Common dependencies for all platforms
|
||||
# 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
|
||||
# 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
|
||||
# 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.10.1, < 5.15
|
||||
transformers >= 5.5.3
|
||||
av>=14.0.0
|
||||
omegaconf>=2.3.0
|
||||
diffusers==0.38.0
|
||||
|
||||
@@ -39,6 +39,8 @@ 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([]),
|
||||
@@ -140,7 +142,12 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None:
|
||||
]
|
||||
|
||||
|
||||
def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]:
|
||||
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"))
|
||||
rendered = _render_test_pipeline(doc, changed_files=changed_files)
|
||||
labels: set[str] = set()
|
||||
|
||||
@@ -156,29 +163,6 @@ def _surviving_labels(doc: dict, 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",
|
||||
[
|
||||
@@ -186,17 +170,18 @@ _COVERAGE_SHARED_INPUTS_DOC = {
|
||||
"pyproject.toml",
|
||||
],
|
||||
)
|
||||
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_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_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
|
||||
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
|
||||
@@ -14,8 +14,7 @@ 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 import config_factory as config_factory_module
|
||||
from vllm_omni.config.config_factory import StageConfigFactory, _materialize_object_storage_configs
|
||||
from vllm_omni.config.config_factory import StageConfigFactory
|
||||
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
|
||||
@@ -56,7 +55,6 @@ 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)
|
||||
@@ -2767,89 +2765,3 @@ 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"
|
||||
@@ -424,10 +424,10 @@
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
82.6344
|
||||
82.1886
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1277
|
||||
0.1241
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -453,13 +453,13 @@
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
170.3737
|
||||
160.9168
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.2076
|
||||
0.1907
|
||||
],
|
||||
"audio_throughput": [
|
||||
37.0834
|
||||
39.2332
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -485,13 +485,13 @@
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
101.0293
|
||||
100.1265
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1498
|
||||
0.1437
|
||||
],
|
||||
"audio_throughput": [
|
||||
10.1186
|
||||
11.7605
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,7 +5,6 @@ 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
|
||||
@@ -13,48 +12,6 @@ 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 = {}
|
||||
|
||||
|
||||
+1
-7
@@ -36,6 +36,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
SEPARATE_CFG_TRANSFORMERS = [
|
||||
DreamIdOmniModel,
|
||||
LTX2VideoTransformer3DModel,
|
||||
HeliosTransformer3DModel,
|
||||
LongCatImageTransformer2DModel,
|
||||
Cosmos3VFMTransformer,
|
||||
@@ -123,13 +124,6 @@ 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,10 +53,14 @@ 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 isinstance(first_output, OmniRequestOutput):
|
||||
raise ValueError(f"Expected an OmniRequestOutput, got {type(first_output).__name__}")
|
||||
if not hasattr(first_output, "request_output") or not first_output.request_output:
|
||||
raise ValueError("No request_output found in OmniRequestOutput")
|
||||
|
||||
images = first_output.images
|
||||
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
|
||||
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,7 +72,8 @@ def _to_float_array(data: Any) -> np.ndarray:
|
||||
|
||||
|
||||
def _extract_output_array(output: Any) -> np.ndarray:
|
||||
payload = output.images[0]
|
||||
req_out = output.request_output
|
||||
payload = req_out.images[0]
|
||||
return _to_float_array(payload)
|
||||
|
||||
|
||||
|
||||
@@ -16,11 +16,6 @@ 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
|
||||
|
||||
@@ -115,85 +110,6 @@ 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__()
|
||||
@@ -263,137 +179,6 @@ 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
|
||||
|
||||
@@ -1,72 +0,0 @@
|
||||
# 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 = ComponentQuantizationConfig(
|
||||
component_configs={"text_encoder": text_encoder_config},
|
||||
default_config=None,
|
||||
)
|
||||
component_config = Mock()
|
||||
component_config.resolve.side_effect = {
|
||||
"text_encoder": text_encoder_config,
|
||||
"transformer": None,
|
||||
}.get
|
||||
|
||||
assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config
|
||||
assert _resolve_component_quant_config(component_config, "transformer") is None
|
||||
|
||||
@@ -1,881 +0,0 @@
|
||||
# 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,7 +3,6 @@
|
||||
|
||||
"""Unit tests for LTX image-to-video input and conditioning behavior."""
|
||||
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
@@ -24,7 +23,7 @@ def _make_ltx23_request_pipe(cls):
|
||||
|
||||
|
||||
class TestLTXImageToVideoForwardStages:
|
||||
def test_legacy_i2v_pil_preprocessing_preserves_aspect_ratio_and_center_crops(self):
|
||||
def test_i2v_pil_preprocessing_matches_official_center_crop(self):
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _preprocess_i2v_pil_images
|
||||
|
||||
pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
|
||||
@@ -36,136 +35,6 @@ 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
|
||||
@@ -229,7 +98,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_and_keeps_audio_per_token(self):
|
||||
def test_denoise_timestep_kwargs_masks_video_only(self):
|
||||
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
||||
|
||||
pipe = object.__new__(LTX2Pipeline)
|
||||
@@ -320,18 +189,13 @@ 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]]]))
|
||||
|
||||
@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):
|
||||
def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
|
||||
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
|
||||
@@ -373,20 +237,14 @@ 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 == [sampled_shape]
|
||||
assert sampled_shapes == [(1, 3, 2)]
|
||||
|
||||
@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):
|
||||
def test_ltx23_i2v_image_noise_is_sampled_after_packing(self, monkeypatch):
|
||||
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
|
||||
@@ -408,7 +266,6 @@ 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),
|
||||
@@ -422,54 +279,10 @@ class TestLTXImageToVideoConditioning:
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
assert sampled_shapes == [sampled_shape]
|
||||
assert sampled_shapes == [(1, 3, 2)]
|
||||
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
|
||||
|
||||
@@ -572,99 +385,6 @@ 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,68 +202,3 @@ 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,15 +28,6 @@ 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,7 +20,6 @@ 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,
|
||||
@@ -48,7 +47,6 @@ 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,
|
||||
@@ -66,14 +64,12 @@ 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,
|
||||
)
|
||||
|
||||
@@ -129,21 +125,15 @@ 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)
|
||||
@@ -264,7 +254,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="not qualified for this LTX recipe"):
|
||||
with pytest.raises(ValueError, match="only one-stage LTX recipes"):
|
||||
LTX2TwoStagePipeline(
|
||||
od_config=SimpleNamespace(
|
||||
model=str(tmp_path),
|
||||
@@ -719,48 +709,15 @@ 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)
|
||||
|
||||
|
||||
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():
|
||||
@pytest.mark.parametrize("has_learned_registers", [False, True])
|
||||
def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_registers):
|
||||
calls = []
|
||||
|
||||
class Backend:
|
||||
@@ -775,16 +732,12 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel():
|
||||
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=UpcastingNorm(),
|
||||
norm_k=UpcastingNorm(),
|
||||
norm_q=torch.nn.Identity(),
|
||||
norm_k=torch.nn.Identity(),
|
||||
to_q=torch.nn.Identity(),
|
||||
to_k=torch.nn.Identity(),
|
||||
to_v=torch.nn.Identity(),
|
||||
@@ -793,29 +746,24 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel():
|
||||
rope_type="split",
|
||||
omni_attention=OmniAttention(),
|
||||
)
|
||||
hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16)
|
||||
hidden_states = torch.randn(2, 3, 8)
|
||||
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()(
|
||||
output = ltx2_components._LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers)(
|
||||
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)
|
||||
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)
|
||||
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)
|
||||
|
||||
|
||||
def test_ltx2_distilled_pipeline_shares_recipe_runtime():
|
||||
@@ -1759,7 +1707,6 @@ class TestLTXForwardStages:
|
||||
assert seen["request_inputs"].output_type == "latent"
|
||||
assert seen["kwargs"] == {
|
||||
"request_sigmas": request_sigmas,
|
||||
"request_phase_sigmas": None,
|
||||
"image": None,
|
||||
}
|
||||
|
||||
@@ -1877,16 +1824,6 @@ 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",
|
||||
@@ -1911,8 +1848,6 @@ class TestRegistryIntegration:
|
||||
expected = [
|
||||
"LTX2Pipeline",
|
||||
"LTX2TwoStagePipeline",
|
||||
"LTX2DistilledOneStagePipeline",
|
||||
"LTX2DistilledTwoStagePipeline",
|
||||
"LTX2DistilledPipeline",
|
||||
"LTX2T2VDMD2Pipeline",
|
||||
"LTX2I2VDMD2Pipeline",
|
||||
@@ -1926,8 +1861,6 @@ class TestRegistryIntegration:
|
||||
[
|
||||
"LTX2Pipeline",
|
||||
"LTX2TwoStagePipeline",
|
||||
"LTX2DistilledOneStagePipeline",
|
||||
"LTX2DistilledTwoStagePipeline",
|
||||
"LTX2DistilledPipeline",
|
||||
"LTX2T2VDMD2Pipeline",
|
||||
"LTX2I2VDMD2Pipeline",
|
||||
|
||||
@@ -7,54 +7,18 @@ 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 (
|
||||
LTX2FeedForward,
|
||||
LTX2AudioVideoAttnProcessor,
|
||||
LTX2VideoTransformer3DModel,
|
||||
LTX2VideoTransformerBlock,
|
||||
_LTX2ParallelAttention,
|
||||
_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)
|
||||
|
||||
@@ -72,53 +36,11 @@ def test_ltx_rms_norm_affine_weight_remains_parameter():
|
||||
assert "weight" in norm.state_dict()
|
||||
|
||||
|
||||
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():
|
||||
def test_ltx_transformer_has_separate_cfg_cache_dit_config():
|
||||
adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
|
||||
|
||||
assert isinstance(adapter_config, CacheDiTAdapterConfig)
|
||||
assert not adapter_config.has_separate_cfg
|
||||
assert adapter_config.has_separate_cfg
|
||||
|
||||
|
||||
def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
|
||||
@@ -154,9 +76,8 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
|
||||
|
||||
|
||||
def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):
|
||||
processor_cls = ltx2_transformer.LTX2AudioVideoAttnProcessor
|
||||
monkeypatch.setattr(processor_cls, "_is_sp_enabled", staticmethod(lambda: True))
|
||||
processor = processor_cls()
|
||||
monkeypatch.setattr(LTX2AudioVideoAttnProcessor, "_is_sp_enabled", staticmethod(lambda: True))
|
||||
processor = LTX2AudioVideoAttnProcessor()
|
||||
hidden_states = torch.zeros(1, 3, 4)
|
||||
encoder_hidden_states = torch.zeros(1, 4, 4)
|
||||
key_padding_mask = torch.tensor([[True, True, True, False]])
|
||||
@@ -192,7 +113,7 @@ def _fake_ltx_parallel_attention():
|
||||
|
||||
return forward
|
||||
|
||||
layer = object.__new__(ltx2_transformer._LTX2ParallelAttention)
|
||||
layer = object.__new__(_LTX2ParallelAttention)
|
||||
nn.Module.__init__(layer)
|
||||
layer.attn_backend = Backend
|
||||
layer.attention = SimpleNamespace(forward=run("native"))
|
||||
@@ -202,281 +123,33 @@ def _fake_ltx_parallel_attention():
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)])
|
||||
def test_ltx_masked_attention_uses_sdpa(query_length, key_length):
|
||||
def test_ltx_masked_cross_attention_uses_sdpa(query_length, key_length):
|
||||
layer, calls = _fake_ltx_parallel_attention()
|
||||
query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16)
|
||||
key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16)
|
||||
key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]])
|
||||
metadata = ltx2_transformer.AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
|
||||
metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
|
||||
|
||||
layer._run_local_attention(query, key, value, metadata)
|
||||
|
||||
assert calls == ["sdpa"]
|
||||
|
||||
|
||||
def test_ltx_unmasked_attention_keeps_native_backend():
|
||||
def test_ltx_masked_self_attention_uses_sdpa():
|
||||
layer, calls = _fake_ltx_parallel_attention()
|
||||
query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
|
||||
key_padding_mask = torch.tensor([[True, True, True, False]])
|
||||
metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
|
||||
|
||||
layer._run_local_attention(query, key, value, metadata)
|
||||
|
||||
assert calls == ["sdpa"]
|
||||
|
||||
|
||||
def test_ltx_unmasked_cross_attention_keeps_native_backend():
|
||||
layer, calls = _fake_ltx_parallel_attention()
|
||||
query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
|
||||
|
||||
layer._run_local_attention(query, key, value, None)
|
||||
|
||||
assert calls == ["native"]
|
||||
|
||||
|
||||
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, patch
|
||||
from unittest.mock import Mock
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
@@ -363,18 +363,8 @@ 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(component_path),
|
||||
str(tmp_path / partition_name) for partition_name in source_partitions
|
||||
]
|
||||
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"
|
||||
@@ -946,66 +936,6 @@ 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
|
||||
|
||||
@@ -1048,21 +978,15 @@ def test_model_offload_uses_hooked_text_encoder_call():
|
||||
pipeline.text_encoder.load_to_device.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"offload_flag",
|
||||
["enable_layerwise_offload", "enable_distributed_layerwise_offload"],
|
||||
)
|
||||
def test_layerwise_offload_releases_text_encoder(offload_flag):
|
||||
def test_layerwise_offload_releases_text_encoder():
|
||||
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=False,
|
||||
enable_distributed_layerwise_offload=False,
|
||||
enable_layerwise_offload=True,
|
||||
)
|
||||
setattr(pipeline.od_config, offload_flag, True)
|
||||
pipeline.text_encoder = Mock()
|
||||
expected = torch.ones(2, 3)
|
||||
pipeline.text_encoder.encode_ids.return_value = expected
|
||||
@@ -2064,25 +1988,21 @@ 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()
|
||||
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.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)
|
||||
@@ -2115,7 +2035,6 @@ 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(
|
||||
@@ -2137,10 +2056,10 @@ def test_encoder_load_reports_missing_fused_source_shards(
|
||||
tmp_path, omitted_sources, expected_detail, unreported_target
|
||||
):
|
||||
encoder = _one_layer_text_encoder()
|
||||
weights = _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources)
|
||||
_write_text_encoder_checkpoint(tmp_path, omit=omitted_sources)
|
||||
|
||||
with pytest.raises(RuntimeError) as excinfo:
|
||||
encoder.load_weights(weights.items())
|
||||
encoder._load_weights(str(tmp_path))
|
||||
|
||||
message = str(excinfo.value)
|
||||
assert expected_detail in message
|
||||
@@ -2149,17 +2068,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()
|
||||
weights = _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",))
|
||||
_write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",))
|
||||
|
||||
with pytest.raises(RuntimeError, match=r"weights not loaded.*input_layernorm"):
|
||||
encoder.load_weights(weights.items())
|
||||
encoder._load_weights(str(tmp_path))
|
||||
|
||||
|
||||
def test_encoder_load_places_every_source_in_its_own_rows(tmp_path):
|
||||
encoder = _one_layer_text_encoder()
|
||||
weights = _write_text_encoder_checkpoint(tmp_path)
|
||||
_write_text_encoder_checkpoint(tmp_path)
|
||||
|
||||
encoder.load_weights(weights.items())
|
||||
encoder._load_weights(str(tmp_path))
|
||||
|
||||
params = dict(encoder.named_parameters())
|
||||
q_rows = _ENCODER_NUM_HEADS * _ENCODER_HEAD_DIM
|
||||
|
||||
@@ -36,13 +36,7 @@ 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",
|
||||
"local_sp_prepare",
|
||||
"sp_gather",
|
||||
}
|
||||
assert set(MiniMaxH3DiTModel._sp_plan["sp_prepare"]) == {0, 1, 2}
|
||||
assert set(MiniMaxH3DiTModel._sp_plan["local_sp_prepare"]) == {2}
|
||||
assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"}
|
||||
|
||||
model = object.__new__(MiniMaxH3DiTModel)
|
||||
nn.Module.__init__(model)
|
||||
@@ -95,279 +89,6 @@ 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,49 +191,6 @@ 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,
|
||||
@@ -254,18 +211,3 @@ 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,6 +142,12 @@ 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,28 +8,15 @@ import torch
|
||||
from PIL import Image
|
||||
from torch import nn
|
||||
|
||||
from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
|
||||
from tests.diffusion.models.wan2_2.conftest import 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)
|
||||
@@ -44,29 +31,6 @@ 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(
|
||||
@@ -89,88 +53,6 @@ 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)
|
||||
@@ -244,169 +126,3 @@ 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,29 +93,6 @@ 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"),
|
||||
[
|
||||
@@ -152,10 +129,9 @@ def test_forward_delegates_denoising_to_diffuse(
|
||||
)
|
||||
batch = DiffusionRequestBatch(requests=[mock_req])
|
||||
|
||||
outputs = pipeline.forward(batch)
|
||||
output = pipeline.forward(batch)
|
||||
|
||||
assert len(outputs) == 1
|
||||
assert torch.equal(outputs[0].output, torch.ones((1, 4, 1, 8, 8)))
|
||||
assert torch.equal(output.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
|
||||
@@ -166,134 +142,6 @@ 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,7 +10,6 @@ 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,
|
||||
)
|
||||
@@ -27,17 +26,6 @@ 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,307 +1,14 @@
|
||||
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)
|
||||
@@ -553,37 +260,31 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
|
||||
pipeline._prompt_clean = Wan22S2VPipeline._prompt_clean
|
||||
|
||||
# -- Build request --
|
||||
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,
|
||||
)
|
||||
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
|
||||
|
||||
ref_image = PIL.Image.new("RGB", (1024, 704))
|
||||
audio_data = np.zeros(16000, dtype=np.float32)
|
||||
|
||||
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,
|
||||
)
|
||||
]
|
||||
)
|
||||
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
|
||||
|
||||
# -- Mock methods that use complex internal state --
|
||||
pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 25, 64, 80), 1, 80))
|
||||
@@ -594,7 +295,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(1, 16, 20, 88, 128))
|
||||
pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(16, 20, 88, 128))
|
||||
pipeline.predict_noise = Wan22S2VPipeline.predict_noise.__get__(pipeline)
|
||||
|
||||
# Mock platform methods
|
||||
@@ -605,7 +306,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)[0]
|
||||
result = Wan22S2VPipeline.forward(pipeline, req=req)
|
||||
|
||||
# -- Assertions --
|
||||
# Text encoder was called
|
||||
|
||||
@@ -8,14 +8,13 @@ import torch
|
||||
from PIL import Image
|
||||
from torch import nn
|
||||
|
||||
from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
|
||||
from tests.diffusion.models.wan2_2.conftest import 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]
|
||||
|
||||
@@ -33,27 +32,6 @@ 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")
|
||||
@@ -160,189 +138,3 @@ 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,11 +18,6 @@ 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,
|
||||
@@ -35,10 +30,7 @@ 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]
|
||||
@@ -56,9 +48,6 @@ class DummyEvent:
|
||||
def record(self, _stream) -> None:
|
||||
return None
|
||||
|
||||
def synchronize(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@contextmanager
|
||||
def dummy_stream(_stream):
|
||||
@@ -150,41 +139,6 @@ 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))
|
||||
@@ -362,9 +316,8 @@ 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))
|
||||
|
||||
def transform(tensor):
|
||||
return tensor.flip(0)
|
||||
next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
|
||||
next_block.weight.mmap_weight_transform_pending = True
|
||||
|
||||
hook = DistributedLayerwiseOffloadHook(
|
||||
next_block=next_block,
|
||||
@@ -373,17 +326,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_transform_is_not_applied_without_a_loader_plan(self, dist_group, patched_offload_runtime):
|
||||
def test_regular_loader_weight_does_not_reapply_mmap_transform(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,
|
||||
@@ -397,82 +350,6 @@ 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):
|
||||
@@ -509,59 +386,23 @@ class TestPinnedResidentLayerGroup:
|
||||
|
||||
group.offload()
|
||||
|
||||
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 test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self):
|
||||
module = nn.Linear(2, 2, bias=False)
|
||||
|
||||
group.load()
|
||||
assert torch.equal(block.weight, expected)
|
||||
assert block.weight.stride() == expected_stride
|
||||
group.offload()
|
||||
def transform(tensor):
|
||||
return tensor.flip(0)
|
||||
|
||||
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())
|
||||
module.weight.mmap_weight_transform = transform
|
||||
|
||||
group = PinnedResidentLayerGroup(
|
||||
blocks,
|
||||
device=torch.device("cpu"),
|
||||
copy_stream=DummyStream(),
|
||||
pin_memory=False,
|
||||
rank_local_mmap=True,
|
||||
)
|
||||
backend = object.__new__(DistributedLayerwiseOffloadBackend)
|
||||
backend._remember_mmap_param_attrs(module)
|
||||
module.to_empty(device="meta")
|
||||
|
||||
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
|
||||
replacement = nn.Parameter(torch.arange(4).reshape(2, 2).float())
|
||||
backend._attach_mmap_param_attrs("weight", replacement, module.weight)
|
||||
|
||||
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)
|
||||
assert replacement.mmap_weight_transform is transform
|
||||
assert replacement.mmap_weight_transform_pending is True
|
||||
|
||||
def test_shared_allgather_output_is_narrowed_to_current_block(
|
||||
self, dist_group, patched_offload_runtime, monkeypatch
|
||||
@@ -704,6 +545,7 @@ class TestMmapWeightLoading:
|
||||
OffloadConfig(
|
||||
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||
pin_cpu_memory=False,
|
||||
model_path=str(tmp_path),
|
||||
),
|
||||
torch.device("cpu"),
|
||||
)
|
||||
@@ -711,81 +553,13 @@ 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, plan)
|
||||
backend._load_weights_via_mmap(pipeline, modules)
|
||||
|
||||
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):
|
||||
@@ -1051,123 +825,23 @@ class TestOffloadPlan:
|
||||
|
||||
|
||||
class TestMmapValidation:
|
||||
"""Tests for loader preflight and backend plan realization."""
|
||||
"""Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights."""
|
||||
|
||||
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),
|
||||
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,
|
||||
)
|
||||
|
||||
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"]
|
||||
|
||||
@@ -1185,118 +859,41 @@ class TestMmapValidation:
|
||||
return key
|
||||
|
||||
pipeline = SimplePipeline()
|
||||
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,
|
||||
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}})
|
||||
)
|
||||
|
||||
assert result.plan is None
|
||||
assert result.fallback_reason == "TP=2 requires the ordinary loader"
|
||||
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"])
|
||||
|
||||
@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)
|
||||
with pytest.raises(ValueError, match="Tensor Parallel"):
|
||||
backend._load_weights_via_mmap(pipeline, modules)
|
||||
|
||||
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,
|
||||
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,
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
@@ -1326,67 +923,27 @@ class TestMmapValidation:
|
||||
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=str(tmp_path),
|
||||
tensor_parallel_size=1,
|
||||
use_hsdp=False,
|
||||
online_quantization=False,
|
||||
)
|
||||
|
||||
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.",
|
||||
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"])
|
||||
|
||||
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),
|
||||
)
|
||||
# Should NOT raise — TP=1 is allowed even with custom weight_loader
|
||||
backend._load_weights_via_mmap(pipeline, modules)
|
||||
|
||||
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"]
|
||||
@@ -1428,64 +985,19 @@ class TestMmapValidation:
|
||||
|
||||
backend = DistributedLayerwiseOffloadBackend(
|
||||
OffloadConfig(
|
||||
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||
pin_cpu_memory=False,
|
||||
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"])
|
||||
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._load_weights_via_mmap(pipeline, modules, result.plan)
|
||||
backend._load_weights_via_mmap(pipeline, modules)
|
||||
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(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)
|
||||
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)
|
||||
|
||||
|
||||
class TinyBlock(nn.Module):
|
||||
@@ -128,33 +128,6 @@ 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,11 +16,10 @@ 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, pytest.mark.cpu]
|
||||
pytestmark = [pytest.mark.core_model]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -348,24 +347,6 @@ 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")
|
||||
|
||||
@@ -3,8 +3,6 @@
|
||||
"""Unit tests for WorkerProc async output logic."""
|
||||
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
@@ -31,59 +29,9 @@ 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."""
|
||||
|
||||
@@ -244,87 +192,6 @@ 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,7 +123,6 @@ 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}",
|
||||
|
||||
@@ -305,49 +305,6 @@ 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,9 +740,6 @@ 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
|
||||
|
||||
@@ -782,11 +779,7 @@ 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, host_weight_plan: None,
|
||||
)
|
||||
monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
|
||||
monkeypatch.setattr(
|
||||
model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
|
||||
)
|
||||
|
||||
@@ -217,21 +217,13 @@ 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()
|
||||
@@ -248,57 +240,6 @@ 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:
|
||||
@@ -661,38 +602,6 @@ 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))
|
||||
@@ -716,47 +625,6 @@ 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))
|
||||
|
||||
@@ -1,171 +0,0 @@
|
||||
# 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)
|
||||
@@ -6,7 +6,6 @@ import concurrent.futures
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
@@ -265,119 +264,6 @@ 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:
|
||||
|
||||
@@ -1,55 +0,0 @@
|
||||
# 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)
|
||||
)
|
||||
@@ -1,591 +0,0 @@
|
||||
# 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()
|
||||
@@ -1,753 +0,0 @@
|
||||
# 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
|
||||
@@ -29,14 +29,14 @@ HEIGHT = 512
|
||||
NUM_INFERENCE_STEPS = 20
|
||||
TRUE_CFG_SCALE = 4.0
|
||||
SEED = 42
|
||||
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 = 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_SDPA_FALLBACK = 0.96
|
||||
PSNR_THRESHOLD = 27.0
|
||||
PSNR_THRESHOLD = 30.0
|
||||
|
||||
MODEL_2512_ID = "Qwen/Qwen-Image-2512"
|
||||
MODEL_2512_ENV_VAR = "QWEN_IMAGE_2512_MODEL"
|
||||
|
||||
@@ -46,7 +46,7 @@ REFERENCE_TEXT_IMG2TEXT = (
|
||||
_OMNI_RUNNER_PARAM = (MODEL_NAME, STAGE_CONFIG)
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.slow,
|
||||
pytest.mark.advanced_model,
|
||||
pytest.mark.diffusion,
|
||||
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
|
||||
]
|
||||
|
||||
@@ -72,7 +72,7 @@ def _send_inpaint_with_generator(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.advanced_model
|
||||
@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.slow
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_flux2_klein_inpaint_different_seeds_different_output(omni_runner_handler: OmniRunnerHandler):
|
||||
|
||||
@@ -1,345 +0,0 @@
|
||||
# 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,7 +39,8 @@ def _extract_audio(multimodal_output: dict) -> torch.Tensor:
|
||||
return audio
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_voxcpm2_zero_shot_001(omni_runner: OmniRunner) -> None:
|
||||
@@ -52,7 +53,8 @@ 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.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_voxcpm2_voice_clone_002(omni_runner: OmniRunner) -> None:
|
||||
@@ -93,7 +95,8 @@ 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.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_voxcpm2_prefill_decode_mixed_batch_003(omni_runner: OmniRunner) -> None:
|
||||
|
||||
@@ -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.slow
|
||||
@pytest.mark.advanced_model
|
||||
@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)
|
||||
|
||||
@@ -75,7 +75,8 @@ def _build_img2img_messages(prompt: str, image_b64: str) -> list[dict]:
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
@@ -97,7 +98,8 @@ def test_bagel_text2img_online(omni_server, openai_client) -> None:
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
|
||||
@@ -59,7 +59,8 @@ _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.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_plain_text_wav(self, omni_server, openai_client) -> None:
|
||||
@@ -78,7 +79,7 @@ class TestHiggsAudioV3OnlineHappyPath:
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.core_model
|
||||
@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:
|
||||
@@ -96,7 +97,8 @@ class TestHiggsAudioV3OnlineHappyPath:
|
||||
)
|
||||
|
||||
@pytest.mark.skip(reason="issue#4411")
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_concurrent_plain_text(self, omni_server, openai_client) -> None:
|
||||
@@ -113,7 +115,7 @@ class TestHiggsAudioV3OnlineHappyPath:
|
||||
)
|
||||
|
||||
@pytest.mark.skip(reason="issue#4411")
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None:
|
||||
@@ -145,7 +147,7 @@ class TestHiggsAudioV3OnlineHappyPath:
|
||||
)
|
||||
|
||||
@pytest.mark.skip(reason="issue#4411")
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_concurrent_pcm_streaming(self, omni_server, openai_client) -> None:
|
||||
@@ -187,7 +189,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
|
||||
They do not assert audio quality (out of scope for CI).
|
||||
"""
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_inline_emotion_and_expressive(self, omni_server, openai_client) -> None:
|
||||
@@ -203,7 +205,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_inline_style_whispering(self, omni_server, openai_client) -> None:
|
||||
@@ -219,7 +221,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_inline_prosody_speed_and_pitch(self, omni_server, openai_client) -> None:
|
||||
@@ -237,7 +239,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_inline_pause_mid_text(self, omni_server, openai_client) -> None:
|
||||
@@ -253,7 +255,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_inline_sfx_with_onomatopoeia(self, omni_server, openai_client) -> None:
|
||||
@@ -277,7 +279,8 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
|
||||
class TestHiggsAudioV3OnlineVoiceClone:
|
||||
"""Voice clone via the two payload shapes the model serves."""
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@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:
|
||||
@@ -294,7 +297,7 @@ class TestHiggsAudioV3OnlineVoiceClone:
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_voice_clone_references_alias(self, omni_server, openai_client) -> None:
|
||||
|
||||
@@ -1,89 +0,0 @@
|
||||
# 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,10 +14,9 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.media import get_asset_path, load_test_audio_data_url
|
||||
from tests.helpers.media import load_test_audio_data_url
|
||||
from tests.helpers.runtime import OmniServerParams
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni.entrypoints.openai.serving_speech import _DEFAULT_VOICE_NAME
|
||||
|
||||
try:
|
||||
from transformers import HiggsAudioV2TokenizerModel # noqa: F401
|
||||
@@ -141,121 +140,3 @@ 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, ``slow`` — mixed input
|
||||
- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2511`` only, ``advanced_model`` — 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.slow
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
@@ -82,7 +82,8 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
@@ -113,7 +114,7 @@ def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenA
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
|
||||
@@ -44,7 +44,8 @@ def _get_diffusion_feature_cases(model: str):
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@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):
|
||||
|
||||
@@ -46,7 +46,8 @@ tts_server_params = [
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
|
||||
@@ -71,7 +72,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.slow
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
|
||||
|
||||
@@ -32,7 +32,7 @@ def _get_diffusion_feature_cases():
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.advanced_model
|
||||
@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.slow
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
|
||||
@@ -161,17 +161,6 @@ 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
|
||||
@@ -244,48 +233,6 @@ 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
|
||||
@@ -364,136 +311,6 @@ 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
|
||||
|
||||
|
||||
@@ -270,56 +270,3 @@ 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"}
|
||||
@@ -1138,8 +1138,7 @@ class TestSingleStageReplicaInitialization:
|
||||
runtime._init_visible_devices_baseline = "0"
|
||||
|
||||
mocker.patch.object(runtime_mod, "inject_kv_stage_info")
|
||||
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.build_diffusion_config", return_value="diffusion-config")
|
||||
mock_register = mocker.patch(
|
||||
"vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
|
||||
return_value=StageRegistrationResponse(
|
||||
@@ -1175,7 +1174,6 @@ 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,
|
||||
@@ -1185,7 +1183,7 @@ class TestSingleStageReplicaInitialization:
|
||||
)
|
||||
mock_manager.assert_called_once_with(
|
||||
model="fake-model",
|
||||
od_config=od_config,
|
||||
od_config="diffusion-config",
|
||||
stage_init_timeout=60,
|
||||
handshake_address="tcp://127.0.0.1:26001",
|
||||
addresses=mocker.ANY,
|
||||
@@ -1232,8 +1230,7 @@ class TestSingleStageReplicaInitialization:
|
||||
runtime._init_visible_devices_baseline = "0"
|
||||
|
||||
mocker.patch.object(runtime_mod, "inject_kv_stage_info")
|
||||
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.build_diffusion_config", return_value="diffusion-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, ReferenceImage
|
||||
from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
|
||||
from vllm_omni.entrypoints.openai.storage import LocalStorageManager
|
||||
from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry
|
||||
from vllm_omni.errors import GuardrailViolationError
|
||||
@@ -461,51 +461,6 @@ 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,10 +10,8 @@ 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
|
||||
@@ -23,64 +21,6 @@ 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,9 +49,7 @@ 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, **_kwargs: (
|
||||
None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]}
|
||||
),
|
||||
lambda path, _model: None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]},
|
||||
)
|
||||
|
||||
od_config = OmniDiffusionConfig(
|
||||
|
||||
@@ -22,7 +22,6 @@ 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,
|
||||
@@ -372,48 +371,6 @@ 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):
|
||||
|
||||
@@ -11,16 +11,11 @@ affect in-process Omni (``OmniRunner`` / offline e2e) vs subprocess-based
|
||||
import pytest
|
||||
|
||||
|
||||
def skip_if_gated_repo_inaccessible(
|
||||
repo_id: str,
|
||||
*,
|
||||
revision: str | None = None,
|
||||
filename: str = "config.json",
|
||||
) -> None:
|
||||
def skip_if_gated_repo_inaccessible(repo_id: str) -> None:
|
||||
"""Skip the test if a gated HuggingFace repo is not accessible.
|
||||
|
||||
Tries to download ``filename`` via ``hf_hub_download``, which performs an
|
||||
actual file-access check (unlike ``HfApi().model_info()``
|
||||
Tries to download the model's config.json 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.
|
||||
@@ -31,7 +26,7 @@ def skip_if_gated_repo_inaccessible(
|
||||
except Exception:
|
||||
return
|
||||
try:
|
||||
hf_hub_download(repo_id=repo_id, filename=filename, revision=revision)
|
||||
hf_hub_download(repo_id=repo_id, filename="config.json")
|
||||
except GatedRepoError as exc:
|
||||
pytest.skip(
|
||||
f"Skipping: gated HF repo {repo_id!r} inaccessible to the current "
|
||||
|
||||
@@ -1,363 +0,0 @@
|
||||
# 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"
|
||||
@@ -1,731 +0,0 @@
|
||||
# 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
|
||||
@@ -20,7 +20,6 @@ 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,
|
||||
)
|
||||
|
||||
|
||||
@@ -231,62 +230,18 @@ 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:
|
||||
@@ -469,13 +424,6 @@ 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,8 +28,6 @@ EXCLUDED_MODELS = [
|
||||
"WanPipeline",
|
||||
"WanVACEPipeline",
|
||||
"LTX2TwoStagePipeline",
|
||||
"LTX2DistilledOneStagePipeline",
|
||||
"LTX2DistilledTwoStagePipeline",
|
||||
"LTX2DistilledPipeline",
|
||||
"LTX2T2VDMD2Pipeline",
|
||||
"LTX2I2VDMD2Pipeline",
|
||||
|
||||
@@ -159,10 +159,11 @@ 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 = step.outputs[0].text # exact text, no strip -> truly bit-identical
|
||||
text = out.outputs[0].text # exact text, no strip -> truly bit-identical
|
||||
elif step.final_output_type == "audio":
|
||||
wav = _audio_bytes(step.outputs[0].multimodal_output["audio"])
|
||||
wav = _audio_bytes(out.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:
|
||||
|
||||
@@ -19,10 +19,6 @@ 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(
|
||||
|
||||
@@ -15,7 +15,6 @@ 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
|
||||
@@ -46,39 +45,6 @@ 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,
|
||||
@@ -151,7 +117,7 @@ class StageConfigFactory:
|
||||
"""
|
||||
hf_config = None
|
||||
try:
|
||||
return get_config(_materialize_object_storage_configs(model), trust_remote_code=trust_remote_code)
|
||||
return get_config(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
|
||||
@@ -197,12 +163,10 @@ 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", config_source, revision=None)
|
||||
config_dict = get_hf_file_to_dict("config.json", model, revision=None)
|
||||
if config_dict:
|
||||
if "model_type" in config_dict:
|
||||
return config_dict["model_type"]
|
||||
@@ -219,7 +183,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", config_source, revision=None)
|
||||
model_index = get_hf_file_to_dict("model_index.json", model, revision=None)
|
||||
if model_index and "_class_name" in model_index:
|
||||
class_name = model_index["_class_name"]
|
||||
for obj in OMNI_PIPELINES.values():
|
||||
@@ -239,10 +203,8 @@ 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). 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("_", "")
|
||||
# so "cosyvoice3" (length 10) beats "cosyvoice" (length 9).
|
||||
model_lower = model.lower().replace("-", "").replace("_", "")
|
||||
best: str | None = None
|
||||
best_len = 0
|
||||
for registered_key in OMNI_PIPELINES.keys():
|
||||
|
||||
@@ -51,7 +51,6 @@ 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
|
||||
@@ -139,7 +138,6 @@ 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,
|
||||
|
||||
@@ -1,40 +0,0 @@
|
||||
# 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
Reference in new issue
Block a user