Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
65c86dec44 | ||
|
|
5ce549bc57 | ||
|
|
fa2f5b4571 | ||
|
|
11687f73bd | ||
|
|
88614175c0 | ||
|
|
e2796b6168 |
No files matched your search
@@ -159,7 +159,19 @@ steps:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- export VLLM_USE_DEEP_GEMM="0"
|
||||
- export VLLM_MOE_USE_DEEP_GEMM="0"
|
||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: ":full_moon: TTS · Accuracy Test · VoxCPM2"
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
|
||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts_expansion.py -m 'full_model and cuda and tts' --run-level 'full_model'
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: ":full_moon: TTS · Perf Test"
|
||||
@@ -176,6 +188,35 @@ steps:
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: TTS · VoxCPM2 · Perf Test"
|
||||
key: nightly-tts-performance-voxcpm2
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_voxcpm2.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: TTS · Higgs-Audio-V3 · Perf Test"
|
||||
key: nightly-tts-performance-higgs-audio-v3
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export VLLM_USE_DEEP_GEMM="0"
|
||||
- export VLLM_MOE_USE_DEEP_GEMM="0"
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_higgs_audio_v3.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- group: ":card_index_dividers: Tiny Model Tests [Multi-GPU]"
|
||||
key: nightly-tiny-model-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
@@ -205,7 +246,8 @@ steps:
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2iat-test"
|
||||
steps:
|
||||
# Use explicit file shards so single-card cases do not reserve 2+ GPUs.
|
||||
# Use explicit file shards so single-card cases do not reserve 2+ GPUs,
|
||||
# and Bagel feature coverage is not duplicated via a broad tests/e2e sweep.
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU"
|
||||
timeout_in_minutes: 120
|
||||
@@ -213,6 +255,15 @@ steps:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_qwen_image_expansion.py
|
||||
tests/e2e/online_serving/test_qwen_image_edit_expansion.py
|
||||
tests/e2e/online_serving/test_qwen_image_layered_expansion.py
|
||||
tests/e2e/online_serving/test_boogu_image.py
|
||||
tests/e2e/online_serving/test_boogu_image_edit.py
|
||||
tests/e2e/online_serving/test_boogu_image_expansion.py
|
||||
tests/e2e/online_serving/test_boogu_image_edit_expansion.py
|
||||
tests/e2e/online_serving/test_bagel_expansion.py
|
||||
tests/e2e/online_serving/test_krea2_expansion.py
|
||||
tests/e2e/offline_inference/test_krea2_expansion.py
|
||||
-m "full_model and diffusion and H100 and not distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
@@ -223,14 +274,47 @@ steps:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_qwen_image_expansion.py
|
||||
tests/e2e/online_serving/test_qwen_image_edit_expansion.py
|
||||
tests/e2e/offline_inference/test_krea2_expansion.py
|
||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Layered/LongCat/FLUX"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_qwen_image_layered_expansion.py
|
||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Bagel"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_bagel_expansion.py
|
||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · GLM-Image Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_glm_image_expansion.py
|
||||
tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
|
||||
-m "full_model and diffusion and H100"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
- pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
mirror_hardwares: l4_4
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Doc Test"
|
||||
@@ -243,7 +327,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export VLLM_HTTP_TIMEOUT_KEEP_ALIVE=120
|
||||
- pytest -s -v tests/e2e/accuracy/test_qwen_image.py --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2i_matches_diffusers' -m full_model --run-level full_model
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
@@ -271,6 +355,70 @@ steps:
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Layered"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · BAGEL"
|
||||
key: nightly-diffusion-x2iat-performance-bagel
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- group: ":card_index_dividers: Diffusion X2V Model Test"
|
||||
key: nightly-diffusion-x2v-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
@@ -292,6 +440,12 @@ steps:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda" --run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · LingBot Video MoE Function Test"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_lingbot_video_moe.py -k "video_generation" -m "full_model and cuda" --run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Doc Test"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
@@ -302,8 +456,15 @@ steps:
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_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:
|
||||
@@ -322,8 +483,23 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · LTX Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
|
||||
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx23_two_stage_layer_fused_i2v]' -m full_model --run-level full_model
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · 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
|
||||
@@ -399,6 +575,16 @@ steps:
|
||||
- pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: ":bar_chart: Testcase Statistics"
|
||||
key: nightly-testcase-statistics
|
||||
timeout_in_minutes: 120
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
# NOTE: ":email: Nightly Collection & Email" is deprecated and will be removed soon.
|
||||
# Do not add new depends_on entries or extend this step.
|
||||
- label: ":email: Nightly Collection & Email"
|
||||
@@ -409,20 +595,32 @@ steps:
|
||||
- nightly-omni-performance-vllm-text
|
||||
- nightly-omni-performance-multi-replicas
|
||||
- nightly-tts-performance
|
||||
- nightly-tts-performance-higgs-audio-v3
|
||||
- nightly-diffusion-x2iat-performance-qwen-image
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
- nightly-diffusion-x2iat-performance-bagel
|
||||
- nightly-diffusion-x2v-performance
|
||||
- nightly-testcase-statistics
|
||||
if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
|
||||
commands:
|
||||
- pip install openpyxl
|
||||
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
|
||||
- export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-higgs-audio-v3
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
|
||||
- python tools/nightly/generate_nightly_perf_excel.py
|
||||
- python tools/nightly/generate_nightly_perf_html.py
|
||||
- python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html"
|
||||
|
||||
@@ -149,4 +149,4 @@ steps:
|
||||
pytest -sv
|
||||
tests/e2e/ -m "slow and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_4
|
||||
mirror_hardwares: h100_2
|
||||
@@ -69,12 +69,6 @@ steps:
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
commands:
|
||||
# Pin the Seed-TTS scorer to the second card: _get_eval_device() otherwise
|
||||
# auto-picks npu:0, which the omni server under test already fills with its
|
||||
# KV cache, so Whisper OOMs on `.to(device)` and every clip reports
|
||||
# "ASR / WER failed".
|
||||
- export SEED_TTS_WER_EVAL=1
|
||||
- export SEED_TTS_EVAL_DEVICE=npu:1
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py -m "full_model and npu and A3" --run-level full_model
|
||||
@@ -102,6 +96,33 @@ steps:
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
- group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
|
||||
key: nightly-diffusion-x2iat-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: >-
|
||||
build.env("NIGHTLY") == "1" ||
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2iat-test"
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: a3_npu_8
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
- group: ":card_index_dividers: Diffusion I2I Model Test"
|
||||
key: nightly-HunyuanImage3-x2i-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
|
||||
@@ -1,15 +1,4 @@
|
||||
steps:
|
||||
- label: "Platforms Test"
|
||||
key: ready-platforms-test
|
||||
depends_on: upload-ready-pipeline
|
||||
mirror_hardwares: a3_npu_2
|
||||
commands:
|
||||
- >-
|
||||
pytest -s -v
|
||||
tests/platforms/npu/test_minicpmo_code2wav_npugraph.py
|
||||
-m 'core_model and npu and A3 and omni'
|
||||
--run-level 'core_model'
|
||||
|
||||
- group: ":card_index_dividers: Omni Model Test"
|
||||
key: ready-omni-group
|
||||
depends_on: upload-ready-pipeline
|
||||
|
||||
@@ -66,7 +66,7 @@ repos:
|
||||
entry: python tools/pre_commit/check_tts_adapter.py
|
||||
language: python
|
||||
pass_filenames: false
|
||||
files: ^(vllm_omni/entrypoints/openai/serving_speech\.py|tools/pre_commit/check_tts_adapter\.py)$
|
||||
files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
|
||||
# The budget-is-now-slack reminder is printed on a PASSING run, and
|
||||
# pre-commit hides output from passing hooks unless the hook is verbose.
|
||||
verbose: true
|
||||
|
||||
@@ -43,7 +43,6 @@ nav:
|
||||
- Overview: features/README.md
|
||||
- Runtime and Stage Execution:
|
||||
- Execution Modes and Streaming: user_guide/diffusion/execution_modes.md
|
||||
- Prefill-Decode Disaggregation (experimental): features/pd_disaggregation.md
|
||||
- Sleep Mode: features/sleep_mode.md
|
||||
- Quantization:
|
||||
- Overview: user_guide/quantization/overview.md
|
||||
@@ -117,7 +116,6 @@ nav:
|
||||
- Async Diffusion Output: design/feature/async_diffusion_output.md
|
||||
- Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
|
||||
- Automatic Prefix Caching: design/feature/prefix_caching.md
|
||||
- Realtime AR-Diffusion: design/feature/realtime_ar_diffusion.md
|
||||
- Communication:
|
||||
- OmniConnector Implementations:
|
||||
- Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 1.3 MiB |
Binary file not shown.
|
Before Width: | Height: | Size: 96 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 112 KiB |
@@ -14,6 +14,10 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
|
||||
|
||||
- **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
|
||||
|
||||
## Multi-Stage Recipes
|
||||
|
||||
- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
|
||||
|
||||
## Optimization Features
|
||||
|
||||
- **[Diffusion Features Overview](../user_guide/diffusion_features.md)** - Complete overview of all diffusion model features and supported models
|
||||
@@ -0,0 +1,171 @@
|
||||
# Prefill-Decode (PD) Disaggregation
|
||||
|
||||
PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
|
||||
stages so prompt processing and token generation can run on different workers.
|
||||
|
||||
This is documented as a stage-config recipe instead of a bundled YAML because the
|
||||
deployment-specific values usually change per environment:
|
||||
|
||||
- GPU placement
|
||||
- `tensor_parallel_size`
|
||||
- connector backend and connector ports
|
||||
- connector IPs or bootstrap addresses
|
||||
|
||||
Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
|
||||
and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
|
||||
changes below.
|
||||
|
||||
## Requirements
|
||||
|
||||
- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
|
||||
- A KV connector supported by vLLM, such as `MooncakeConnector`
|
||||
- Matching `tensor_parallel_size` on the prefill and decode thinker stages
|
||||
|
||||
## 1. Split the thinker into prefill and decode stages
|
||||
|
||||
Replace the original thinker stage with two stages:
|
||||
|
||||
```yaml
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_prefill_only: true
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 16
|
||||
model_stage: thinker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.9
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 1
|
||||
kv_transfer_config:
|
||||
kv_connector: "MooncakeConnector"
|
||||
kv_role: "kv_producer"
|
||||
kv_rank: 0
|
||||
kv_parallel_size: 2
|
||||
kv_connector_extra_config:
|
||||
mooncake_bootstrap_port: 25201
|
||||
final_output: false
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
is_decode_only: true
|
||||
runtime:
|
||||
devices: "1"
|
||||
engine_args:
|
||||
max_num_seqs: 64
|
||||
model_stage: thinker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.9
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 1
|
||||
kv_transfer_config:
|
||||
kv_connector: "MooncakeConnector"
|
||||
kv_role: "kv_consumer"
|
||||
kv_rank: 1
|
||||
kv_parallel_size: 2
|
||||
kv_connector_extra_config:
|
||||
mooncake_bootstrap_port: 25202
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
```
|
||||
|
||||
Notes:
|
||||
|
||||
- `is_prefill_only: true` marks the thinker stage that only saves KV.
|
||||
- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
|
||||
- `kv_transfer_config` is required on both stages.
|
||||
- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
|
||||
prefill side only processes the prompt and exports KV.
|
||||
|
||||
## 2. Shift the downstream stages by one index
|
||||
|
||||
After inserting the extra thinker stage, renumber the remaining stages:
|
||||
|
||||
```yaml
|
||||
- stage_id: 2
|
||||
runtime:
|
||||
devices: "2"
|
||||
engine_input_source: [1]
|
||||
sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only
|
||||
|
||||
- stage_id: 3
|
||||
runtime:
|
||||
devices: "2"
|
||||
engine_args:
|
||||
max_num_seqs: 1
|
||||
engine_input_source: [2]
|
||||
sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only
|
||||
```
|
||||
|
||||
Compared with the default Qwen3-Omni config:
|
||||
|
||||
- the talker becomes stage `2` instead of stage `1`
|
||||
- the code2wav stage becomes stage `3` instead of stage `2`
|
||||
- the talker now reads from decode stage `1`
|
||||
|
||||
## 3. Add runtime edges for the four-stage pipeline
|
||||
|
||||
```yaml
|
||||
runtime:
|
||||
enabled: true
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
- from: 1
|
||||
to: 2
|
||||
- from: 2
|
||||
to: 3
|
||||
```
|
||||
|
||||
## 4. Launch with your custom config
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--deploy-config /path/to/qwen3_omni_pd.yaml
|
||||
```
|
||||
|
||||
## Operational Notes
|
||||
|
||||
- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
|
||||
pair. Keep prefill and decode at the same `tensor_parallel_size`.
|
||||
- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
|
||||
separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
|
||||
- Choose connector ports and addresses that match your deployment. The values
|
||||
shown above are examples only.
|
||||
@@ -4,9 +4,6 @@
|
||||
> `experimental/ar_diffusion` contract. It follows the ownership and lifecycle
|
||||
> direction in [#5137](https://github.com/vllm-project/vllm-omni/issues/5137),
|
||||
> but it is not yet part of the stable module-design hierarchy.
|
||||
>
|
||||
> The end-to-end request, identity, failure, and memory contracts are documented
|
||||
> in [Realtime AR-Diffusion sessions](feature/realtime_ar_diffusion.md).
|
||||
|
||||
`ARDiffusionEngine` is an opt-in diffusion backend for pipelines that keep
|
||||
autoregressive attention KV across requests. The engine only selects
|
||||
@@ -54,13 +51,14 @@ time, and its inherited batch/step entry points also fail explicitly. Supporting
|
||||
those paths requires a batch-aware state-binding contract rather than silently
|
||||
bypassing `bind_ar_diffusion_state()`.
|
||||
|
||||
The implementation may retain multiple sessions per model instance. It selects
|
||||
the largest feasible capacity up to the pipeline's `session_capacity`, based on
|
||||
self-attention pages, cross-attention allocations, scratch, and declared
|
||||
model-owned state. One session is admitted whenever it fits actual free device
|
||||
memory; `gpu_memory_fraction` controls additional residency. Requests beyond
|
||||
that capacity evict the least-recently-used session through the same lifecycle
|
||||
hook used by explicit close and failure cleanup.
|
||||
The current implementation retains one session per model instance. A request
|
||||
for a different `session_id` closes the resident session through the same LRU
|
||||
lifecycle used by explicit close and failure cleanup before allocating the new
|
||||
one. Self-attention and cross-attention memory are therefore sized for one
|
||||
resident session even if the pipeline declares a larger `session_capacity`.
|
||||
This prioritizes rollout throughput and prevents block-pool exhaustion before a
|
||||
count-based eviction can run. Multi-session residency can be added later with a
|
||||
memory-backed capacity calculation or allocation-pressure eviction.
|
||||
|
||||
Scratch storage is also capability-driven. `frames_per_block` reserves space for
|
||||
an uncommitted current video block, while `max_scratch_tokens_per_branch`
|
||||
|
||||
@@ -1,213 +0,0 @@
|
||||
# Realtime AR-Diffusion sessions
|
||||
|
||||
> **Status:** experimental. This document describes the internal execution
|
||||
> contract implemented under `vllm_omni.experimental.ar_diffusion`. It does not
|
||||
> define a public HTTP or WebSocket schema.
|
||||
|
||||
## Scope
|
||||
|
||||
Realtime AR-Diffusion serves world models that generate one autoregressive
|
||||
latent block per request while retaining model state across requests. The
|
||||
generic runtime owns identity, event ordering, session lifecycle, routing, and
|
||||
KV allocation. A model adapter owns model-specific controls and conditioning.
|
||||
|
||||
LingBot World v2 is the first single-KV-branch integration. DreamZero uses the
|
||||
same runner capability with two logical CFG branches.
|
||||
|
||||
## Identities
|
||||
|
||||
Four identifiers have distinct lifetimes:
|
||||
|
||||
| Identifier | Lifetime | Owner | Contract |
|
||||
| --- | --- | --- | --- |
|
||||
| `session_id` | One persistent world | Session manager | Selects runner KV and model-owned state. |
|
||||
| `event_id` | One control/prompt update | Transport adapter | Unique and monotonically increasing within a session, including across reset. |
|
||||
| `chunk_index` | One committed AR block position | Session | Contiguous from zero; reset starts again at zero. |
|
||||
| `request_id` | One chunk execution | Session | Correlates the tick with its returned metadata. |
|
||||
|
||||
`AsyncOmni` keeps its normal UUID-suffixed engine routing ID. The tick
|
||||
`request_id` remains inside the immutable AR-Diffusion snapshot and the
|
||||
standard output metadata, while caller-visible `OmniRequestOutput.request_id`
|
||||
continues to use AsyncOmni's existing external-ID contract. Protocol identity
|
||||
therefore does not require a private entrypoint override.
|
||||
|
||||
## Request lifecycle
|
||||
|
||||
1. `ARDiffusionSession.accept_event()` validates and queues an immutable
|
||||
`ARDiffusionSessionEvent`. Prompt and opaque control tracks may be updated in
|
||||
the same event.
|
||||
2. `ARDiffusionSession.next_chunk()` serializes chunk execution and calls
|
||||
`_build_tick_locked()`. The resulting `ARDiffusionTickRequest` is a snapshot
|
||||
of the current prompt, reduced controls, ordered event IDs, and chunk index.
|
||||
3. `ARDiffusionOmniTickConsumer.execute_tick()` copies the stage sampling
|
||||
parameters, stores the typed tick in
|
||||
`sampling_params.extra_args["ar_diffusion_tick"]`, constructs the normal Omni
|
||||
prompt through its `prompt_provider`, and invokes `AsyncOmni.generate()`.
|
||||
4. `AsyncOmni` allocates its normal unique engine routing ID and submits one
|
||||
`OmniDiffusionRequest` through the orchestrator.
|
||||
5. `ARDiffusionModelRunner._request_session()` parses the typed tick.
|
||||
`_get_or_create_session()` obtains runner-owned paged KV, and
|
||||
`bind_ar_diffusion_state()` exposes it to the pipeline only for the duration
|
||||
of `execute_model()`.
|
||||
6. The model pipeline validates the tick and generates exactly one latent
|
||||
block. LingBot interprets camera controls in
|
||||
`LingBotCameraControlReducer` and
|
||||
`LingBotWorldCausalDMDPipeline.forward()`; the generic runtime does not parse
|
||||
camera or action schemas.
|
||||
7. The pipeline returns one `DiffusionOutput` with the generated payload and an
|
||||
`ARDiffusionChunkMetadata` mapping under `metadata["ar_diffusion"]`.
|
||||
`AsyncOmni` exposes it as
|
||||
`OmniRequestOutput.multimodal_output["metadata"]["ar_diffusion"]`.
|
||||
8. The consumer parses the standard envelope. The session commits its reducer
|
||||
state, prompt, controls, event IDs, and next chunk index only when returned
|
||||
metadata exactly equals the submitted tick snapshot.
|
||||
|
||||
The metadata envelope is:
|
||||
|
||||
```json
|
||||
{
|
||||
"metadata": {
|
||||
"ar_diffusion": {
|
||||
"session_id": "world-7",
|
||||
"request_id": "ar-world-7-3-...",
|
||||
"chunk_index": 3,
|
||||
"applied_event_ids": [10, 11]
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Transaction and failure semantics
|
||||
|
||||
Accepted events are not acknowledged as applied until the model output and
|
||||
metadata are validated. Reducer `prepare()` is speculative; reducer `commit()`
|
||||
and the session snapshot form one logical commit. A model, metadata, or reducer
|
||||
failure leaves the chunk index and queued events uncommitted and transitions the
|
||||
session to `FAILED`.
|
||||
|
||||
After a tick failure, worker state is closed because a forward may have
|
||||
partially modified or already released KV. The same chunk is not retried in
|
||||
place. A caller must explicitly reset the session, starting again at chunk
|
||||
zero, or close it.
|
||||
|
||||
Explicit close follows this state machine:
|
||||
|
||||
```text
|
||||
ACTIVE or FAILED
|
||||
|
|
||||
v
|
||||
CLOSING ---- worker cleanup succeeds ----> CLOSED
|
||||
|
|
||||
+------- worker cleanup fails --------> CLEANUP_FAILED
|
||||
|
|
||||
+-- close retry --> CLOSING
|
||||
```
|
||||
|
||||
`CLEANUP_FAILED` is a local tombstone. The session manager retains it, rejects
|
||||
creation of the same `session_id`, and rejects reset. Only an idempotent close
|
||||
retry that succeeds removes the session. Transport disconnect uses the same
|
||||
path.
|
||||
|
||||
## KV ownership and capacity
|
||||
|
||||
`ARDiffusionKVCacheSpec` is the pipeline-owned immutable geometry contract. It
|
||||
declares:
|
||||
|
||||
- TP-local self-attention geometry, AR block size, sink, and recent window;
|
||||
- logical KV branches and their worker-local slots;
|
||||
- fixed-length cross-attention caches;
|
||||
- scratch tokens required by an uncommitted forward;
|
||||
- persistent model-owned CUDA bytes per resident session; and
|
||||
- the requested resident-session capacity.
|
||||
|
||||
The runner computes:
|
||||
|
||||
```text
|
||||
required(capacity) =
|
||||
scratch
|
||||
+ managed_self_attention_pages(capacity)
|
||||
+ capacity * cross_attention_bytes_per_session
|
||||
+ capacity * model_owned_state_bytes_per_session
|
||||
```
|
||||
|
||||
`gpu_memory_fraction` is a soft expansion budget:
|
||||
|
||||
```text
|
||||
configured_budget = available_device_bytes * gpu_memory_fraction
|
||||
effective_budget = max(configured_budget, required(1))
|
||||
```
|
||||
|
||||
If `required(1)` exceeds actual available device memory, initialization fails.
|
||||
Otherwise the runner always admits one viable session and uses the configured
|
||||
fraction to select additional resident sessions up to the model-declared cap.
|
||||
All model-owned reservations are deducted before the paged self-attention pool
|
||||
is allocated.
|
||||
|
||||
LingBot reports its persistent full-horizon image-condition tensor. DreamZero
|
||||
reports the measured 603 MiB per-session Wan VAE causal-convolution state.
|
||||
|
||||
## Concurrency and routing limits
|
||||
|
||||
The current runner executes one request at a time:
|
||||
|
||||
- `max_num_seqs` must be one;
|
||||
- request batching and diffusion step execution are rejected;
|
||||
- one AR block is generated per engine request; and
|
||||
- an AR-Diffusion stage must have exactly one replica.
|
||||
|
||||
The single-replica restriction is intentional. Multi-replica support requires
|
||||
session-affine routing so every tick, reset, and close reaches the worker that
|
||||
owns the session. Random or round-robin routing is not correct.
|
||||
|
||||
Multiple resident sessions may share one runner. Capacity selection accounts
|
||||
for all persistent pools, and the runner performs LRU eviction through the same
|
||||
model lifecycle hook used by explicit close.
|
||||
|
||||
## Model boundary
|
||||
|
||||
The generic runtime transports controls as opaque
|
||||
`ARDiffusionControlInput(track, schema, data)` values. Model adapters validate
|
||||
and reduce those values. This keeps target/velocity, key state, pose,
|
||||
intrinsics, and other model-specific conventions out of the session protocol.
|
||||
|
||||
For LingBot World v2:
|
||||
|
||||
- `lingbot_world/actions.py` parses key-state and trajectory controls;
|
||||
- `lingbot_world/camera.py` loads or constructs camera trajectories and Plücker embeddings;
|
||||
- `lingbot_world/transformer.py` implements checkpoint-compatible causal attention; and
|
||||
- `lingbot_world/pipeline.py` constructs conditioning, owns small non-KV session state, and
|
||||
produces one block plus the standard metadata envelope.
|
||||
|
||||
## Non-goals
|
||||
|
||||
This contract does not currently provide:
|
||||
|
||||
- a public HTTP or WebSocket request schema (the structured-interaction
|
||||
frontend is tracked separately in
|
||||
[#5527](https://github.com/vllm-project/vllm-omni/pull/5527));
|
||||
- camera/action semantics shared by every world model;
|
||||
- session migration or replication across workers;
|
||||
- retry of an ambiguous partially executed chunk;
|
||||
- cross-stage KV transfer; or
|
||||
- stateful streaming VAE decode.
|
||||
|
||||
Those features can be layered on top without adding LingBot-specific identity
|
||||
or lifecycle fields to the generic runtime.
|
||||
|
||||
## Required regression coverage
|
||||
|
||||
CPU contract tests cover:
|
||||
|
||||
- separation of tick correlation IDs from internal engine routing IDs;
|
||||
- event ordering, metadata equality, and reducer fail-closed commit;
|
||||
- close/disconnect failure tombstones and cleanup retry;
|
||||
- capacity one under the default fraction for shipped LingBot geometry;
|
||||
- capacity expansion under a tuned fraction;
|
||||
- DreamZero model-owned state accounting;
|
||||
- runner session reuse, reset, close, LRU eviction, and forward cleanup; and
|
||||
- LingBot model registration, imports, camera/action reduction, and one-block
|
||||
output metadata.
|
||||
|
||||
GPU validation must additionally exercise real model loading, action input,
|
||||
prompt switching, rolling-window replacement, reset, close, finite output
|
||||
latents, and exact metadata/request identity.
|
||||
@@ -22,7 +22,6 @@ implementation contract; it is not, by itself, a general support claim.
|
||||
- [Async Diffusion Output](feature/async_diffusion_output.md)
|
||||
- [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
|
||||
- [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
|
||||
- [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md)
|
||||
|
||||
### Communication
|
||||
|
||||
|
||||
@@ -20,19 +20,13 @@ an implementation contract has a user-facing workflow.
|
||||
| Choose serial, batched, step-wise, or streaming diffusion execution | [Execution Modes and Streaming](../user_guide/diffusion/execution_modes.md) | [Diffusion Continuous Batching](../design/feature/diffusion_continuous_batching.md), [Async Diffusion Output](../design/feature/async_diffusion_output.md) |
|
||||
| Reclaim stage memory without restarting the server | [Sleep Mode](sleep_mode.md) | Runtime lifecycle behavior is documented in the user guide |
|
||||
|
||||
The [Prefill-Decode Disaggregation design reference](pd_disaggregation.md)
|
||||
describes the experimental Qwen3-Omni runtime topology and its current
|
||||
integration boundaries. It is not yet a supported launch recipe.
|
||||
|
||||
Some runtime designs are deliberately not promoted as standalone User Guide
|
||||
features yet:
|
||||
|
||||
- Generic, OmniConnector-based
|
||||
[Disaggregated Inference](../design/feature/disaggregated_inference.md) is not
|
||||
yet a standalone User Guide feature. Its
|
||||
- [Disaggregated Inference](../design/feature/disaggregated_inference.md) and
|
||||
[OmniConnector implementations](../design/index.md#communication) describe
|
||||
transport contracts, while practical connector configuration remains under
|
||||
[Pipeline and deploy configurations](../configuration/stage_configs.md).
|
||||
topology and transport contracts. Practical deployment configuration remains
|
||||
under [Pipeline and deploy configurations](../configuration/stage_configs.md).
|
||||
- [Async Chunk](../design/feature/async_chunk.md) and
|
||||
[Async Omni Output Materialization](../design/feature/omni_async_output_materialization.md)
|
||||
are model- and pipeline-dependent. Use the selected model's deploy
|
||||
|
||||
@@ -1,140 +0,0 @@
|
||||
# Prefill-Decode (PD) Disaggregation (experimental)
|
||||
|
||||
!!! warning "Experimental design reference"
|
||||
PD disaggregation is only partially integrated on the current `main` branch.
|
||||
The runtime contains PD detection and routing scaffolding, but there is no
|
||||
supported deploy overlay or end-to-end validated launch recipe yet. Treat this
|
||||
page as a description of the current design, not as production guidance.
|
||||
|
||||
Prefill-decode disaggregation splits the Qwen3-Omni Thinker into two logical
|
||||
stages: a prefill worker that processes the prompt and produces KV cache, and a
|
||||
decode worker that imports that cache and generates tokens. The Talker and
|
||||
Code2Wav stages remain downstream; PD does not disaggregate the entire
|
||||
Qwen3-Omni pipeline.
|
||||
|
||||
## Current design
|
||||
|
||||
### Runtime topology
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
client["Client request"] --> orchestrator["Omni Orchestrator"]
|
||||
orchestrator -->|"original prompt<br/>max_tokens = 1"| prefill["Thinker prefill<br/>KV producer"]
|
||||
prefill -->|"KV blocks"| connector[("vLLM KV connector")]
|
||||
prefill -->|"KV routing metadata"| orchestrator
|
||||
orchestrator -->|"original prompt<br/>remote KV parameters"| decode["Thinker decode<br/>KV consumer"]
|
||||
connector -->|"remote KV load"| decode
|
||||
decode -->|"generated text"| client
|
||||
```
|
||||
|
||||
The orchestrator owns the request lifecycle, while the vLLM KV connector owns
|
||||
the bulk KV transfer. The original prompt is submitted to both Thinker stages:
|
||||
the decode worker still needs the request tokens and multimodal features to
|
||||
construct its request, while the imported KV cache prevents it from repeating
|
||||
the completed prefill computation.
|
||||
|
||||
For the full Qwen3-Omni pipeline, splitting the Thinker adds one logical stage:
|
||||
|
||||
| Stage | Role | Output or handoff |
|
||||
| --- | --- | --- |
|
||||
| 0 | Thinker prefill | Saves prompt KV and prompt-side multimodal output |
|
||||
| 1 | Thinker decode | Loads remote KV, generates text, and produces decode-side conditioning |
|
||||
| 2 | Talker | Converts combined Thinker conditioning into codec codes |
|
||||
| 3 | Code2Wav | Converts codec codes into audio |
|
||||
|
||||
Text remains a final output of the decode stage, and audio remains a final
|
||||
output of Code2Wav.
|
||||
|
||||
### Request lifecycle
|
||||
|
||||
1. At startup, PD detection looks for one stage marked `is_prefill_only` and
|
||||
one downstream stage marked `is_decode_only`. The decode stage must list the
|
||||
prefill stage as an input source.
|
||||
2. The orchestrator clones the Thinker sampling parameters for prefill, forces
|
||||
`max_tokens=1`, clears stop conditions, and sets producer-side
|
||||
`kv_transfer_params`. This makes prefill finish by length after exporting KV.
|
||||
3. When prefill finishes, the orchestrator captures its KV routing metadata and
|
||||
Qwen multimodal output under the original request ID.
|
||||
4. The orchestrator submits the original prompt to decode with consumer-side
|
||||
`kv_transfer_params`. Those parameters identify the transfer, prefill engine,
|
||||
bootstrap endpoint, and, when required by the connector, remote request.
|
||||
5. The decode engine imports the KV cache and continues generation. The normal
|
||||
Omni stage path then carries Thinker conditioning through Talker and
|
||||
Code2Wav.
|
||||
|
||||
### Control and data paths
|
||||
|
||||
PD uses two different connector layers. They solve different handoffs and are
|
||||
not interchangeable.
|
||||
|
||||
| Path | Payload | Owner |
|
||||
| --- | --- | --- |
|
||||
| PD control path | Request ID, original prompt, `transfer_id`, bootstrap address, engine ID, and connector-specific request metadata | Omni Orchestrator |
|
||||
| PD data path | Thinker attention KV blocks | vLLM KV connector |
|
||||
| Downstream model path | Thinker hidden states and embeddings, Talker codec codes, and audio tensors | vLLM-Omni stage connectors and processors |
|
||||
|
||||
### Detection and validation
|
||||
|
||||
The current runtime accepts one logical prefill/decode pair per pipeline. Each
|
||||
stage may use tensor parallelism, but both sides of the pair must satisfy these
|
||||
invariants:
|
||||
|
||||
- both stages define `kv_transfer_config`
|
||||
- prefill uses `kv_producer` or `kv_both`
|
||||
- decode uses `kv_consumer` or `kv_both`
|
||||
- both stages use the same KV connector
|
||||
- `kv_buffer_device` and `kv_buffer_size` agree when set on both stages
|
||||
- both stages use the same `tensor_parallel_size`
|
||||
|
||||
These checks happen before request processing so an inconsistent PD topology
|
||||
fails during engine initialization.
|
||||
|
||||
## What is implemented today
|
||||
|
||||
| Area | Current state |
|
||||
| --- | --- |
|
||||
| Pair detection and topology validation | Implemented in the shared PD helper |
|
||||
| Prefill sampling preparation | Implemented for sync and async entrypoints |
|
||||
| Prefill-to-decode routing | Implemented in the current Orchestrator |
|
||||
| Deploy-based configuration | Not yet migrated; the supported `stages:` deploy schema cannot currently expand Qwen3-Omni into the four-stage PD topology |
|
||||
| Qwen3 Thinker-to-Talker handoff | Bridge state and embedding-merge helpers exist, but the merge is not wired into the live stage processor |
|
||||
| Connector compatibility | Incomplete; the Orchestrator currently requires `remote_request_id`, while the pinned Mooncake connector completes prefill without returning that metadata |
|
||||
| Validation coverage | PD entrypoint tests and the PD deploy overlay in online end-to-end tests are temporarily disabled during migration |
|
||||
| Performance | No non-regression or throughput/latency result is established for the current path |
|
||||
|
||||
Because these pieces are still converging, the previous legacy `stage_args` YAML
|
||||
and `vllm serve` command have been removed. The serve CLI now accepts the
|
||||
deploy-based `stages:` schema, so retaining the legacy recipe would imply a
|
||||
launch path that current `main` rejects.
|
||||
|
||||
## Configuration contract under development
|
||||
|
||||
The intended deploy integration starts from the
|
||||
[default Qwen3-Omni deploy config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
|
||||
and expands its frozen three-stage pipeline into the four stages shown above.
|
||||
A complete integration must:
|
||||
|
||||
- split the Thinker definition without duplicating the Talker or Code2Wav model
|
||||
contracts
|
||||
- remap downstream stage IDs, input sources, and stage connectors
|
||||
- inject producer and consumer `kv_transfer_config` values into the two Thinker
|
||||
stages
|
||||
- allocate separate device sets for prefill and decode
|
||||
- preserve both final outputs: decode text and Code2Wav audio
|
||||
- provide connector-specific routing metadata without hard-coding one connector
|
||||
contract
|
||||
|
||||
The required number of GPUs depends on tensor-parallel sizes and whether Talker
|
||||
and Code2Wav share a device. A basic TP=1 layout uses one device for prefill, one
|
||||
for decode, and one for the two downstream audio stages.
|
||||
|
||||
## Implementation map
|
||||
|
||||
- PD detection, validation, and sampling-parameter preparation:
|
||||
[`PDDisaggregationMixin`](gh-file:vllm_omni/entrypoints/pd_utils.py)
|
||||
- startup detection and connector metadata:
|
||||
[`AsyncOmniEngine`](gh-file:vllm_omni/engine/async_omni_engine.py)
|
||||
- request-scoped prefill-to-decode routing:
|
||||
[`Orchestrator`](gh-file:vllm_omni/engine/orchestrator.py)
|
||||
- Qwen3 downstream conditioning helpers:
|
||||
[Qwen3-Omni stage input processor](gh-file:vllm_omni/model_executor/stage_input_processors/qwen3_omni.py)
|
||||
@@ -37,7 +37,6 @@ th {
|
||||
| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — |
|
||||
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
|
||||
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
@@ -53,7 +52,6 @@ th {
|
||||
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
|
||||
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
|
||||
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
@@ -119,10 +117,6 @@ support metadata until their recipe is audited.
|
||||
|
||||
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
|
||||
|
||||
See the [LingBot-World 2.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Robbyant/LingBot-World-2.0.md)
|
||||
for offline and experimental realtime usage, supported controls, validation,
|
||||
and current limitations.
|
||||
|
||||
<sup>H3</sup> MiniMax H3 AMD GPU support is validated on gfx942 (MI300X) and
|
||||
gfx950 (MI350) in BF16, using the AITER `FLASH_ATTN` diffusion attention backend.
|
||||
Other AMD SKUs (e.g. MI325X) are not listed until their
|
||||
|
||||
@@ -141,20 +141,7 @@ The response shape depends on the streaming parameters:
|
||||
|
||||
**Non-streaming (default).** With `stream=false` and no `stream_format`, returns the
|
||||
complete clip as binary audio data with an appropriate `Content-Type` header (e.g.
|
||||
`audio/wav`). Because the raw-bytes body has no JSON carrier, successful
|
||||
non-streaming responses from non-diffusion speech servers report usage through
|
||||
response headers:
|
||||
|
||||
| Header | Description |
|
||||
| --- | --- |
|
||||
| `x-vllm-omni-input-tokens` | Total input tokens (`text_tokens` + `audio_tokens`). |
|
||||
| `x-vllm-omni-output-tokens` | Generated codec/audio tokens. |
|
||||
| `x-vllm-omni-total-tokens` | `input_tokens` + `output_tokens`. |
|
||||
| `x-vllm-omni-input-text-tokens` | Tokens from the synthesized text (`input` plus `instructions`). |
|
||||
| `x-vllm-omni-input-audio-tokens` | Reference-audio codec frames, non-zero only for in-context voice cloning. |
|
||||
|
||||
Diffusion-mode speech servers route through a separate response path and do not
|
||||
emit these headers.
|
||||
`audio/wav`). The raw-bytes body has no JSON carrier, so no `usage` is reported.
|
||||
|
||||
**Raw audio stream** (`stream_format="audio"`). Streams raw audio bytes (PCM or
|
||||
WAV) as they are decoded.
|
||||
|
||||
@@ -167,14 +167,6 @@ transfer options such as `control_path`, `control_guidance`,
|
||||
`num_conditional_frames`, `show_control_condition`, and `show_input`; see the
|
||||
Cosmos3 recipe for complete examples.
|
||||
|
||||
HTTP redirects for `image_reference.image_url` follow vLLM's
|
||||
`VLLM_MEDIA_URL_ALLOW_REDIRECTS` setting. Before starting the server, set it to
|
||||
`1` (the default) to allow redirects or `0` to reject them. A redirect target
|
||||
can differ from the original host, and vLLM-Omni does not yet fully implement
|
||||
upstream vLLM's media URL allowlist protection. Deployments should therefore
|
||||
treat remote media URLs as untrusted and choose this setting as part of their
|
||||
URL access policy.
|
||||
|
||||
### Speech-to-Video
|
||||
|
||||
For models that support audio-driven generation (e.g., Wan2.2-S2V), pass both
|
||||
|
||||
@@ -24,6 +24,7 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.model_executor.stage_input_processors.aura_omni import (
|
||||
DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
@@ -227,6 +228,7 @@ def parse_args():
|
||||
parser.add_argument("--init-timeout", type=int, default=2000, help="Overall initialization timeout.")
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=2000, help="Per-stage initialization timeout.")
|
||||
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.")
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -1,288 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Generate a LingBot-World v2 video from an image and camera trajectory.
|
||||
|
||||
The official checkpoint is licensed separately under CC BY-NC-SA and is
|
||||
restricted to non-commercial use. This vLLM-Omni integration remains
|
||||
Apache-2.0 licensed.
|
||||
|
||||
Example:
|
||||
python examples/offline_inference/diffusion/lingbot_world_v2.py \
|
||||
--prompt "The camera moves slowly forward through the scene." \
|
||||
--image /path/to/first_frame.png \
|
||||
--action-dir /path/to/actions/forward \
|
||||
--output lingbot_world_v2.mp4
|
||||
|
||||
``--action-dir`` must contain ``poses.npy`` and ``intrinsics.npy``. The
|
||||
example configures its parent as the trusted action root and sends the
|
||||
contained directory through ``extra_args["action_path"]``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import math
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
|
||||
_NUM_INFERENCE_STEPS = 4
|
||||
_MAX_SEQUENCE_LENGTH = 512
|
||||
_MAX_PIXEL_AREA = 480 * 832
|
||||
_MAX_RAW_FRAMES = 117
|
||||
_MAX_SOURCE_CAMERA_FRAMES = 4096
|
||||
_TEMPORAL_COMPRESSION = 4
|
||||
_LATENT_FRAMES_PER_BLOCK = 3
|
||||
_SPATIAL_ALIGNMENT = 16
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LingBotPaths:
|
||||
image: Path
|
||||
action_dir: Path
|
||||
action_root: Path
|
||||
action_relative: Path
|
||||
camera_frames: int
|
||||
output: Path
|
||||
|
||||
|
||||
def _validate_camera_arrays(action_dir: Path) -> int:
|
||||
"""Validate the bounded camera contract before starting the engine."""
|
||||
|
||||
try:
|
||||
poses = np.load(action_dir / "poses.npy", mmap_mode="r", allow_pickle=False)
|
||||
intrinsics = np.load(action_dir / "intrinsics.npy", mmap_mode="r", allow_pickle=False)
|
||||
except (OSError, TypeError, ValueError):
|
||||
raise ValueError("Camera files must be valid numeric poses.npy and intrinsics.npy arrays.") from None
|
||||
if poses.ndim != 3 or poses.shape[1:] != (4, 4):
|
||||
raise ValueError("poses.npy must have shape [frames, 4, 4].")
|
||||
if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,):
|
||||
raise ValueError("intrinsics.npy must have shape [frames, 4].")
|
||||
if poses.shape[0] != intrinsics.shape[0]:
|
||||
raise ValueError("poses.npy and intrinsics.npy must contain the same frame count.")
|
||||
if not 1 <= poses.shape[0] <= _MAX_SOURCE_CAMERA_FRAMES:
|
||||
raise ValueError("Camera arrays must contain between 1 and 4096 source frames.")
|
||||
if poses.dtype.kind not in "fiu" or intrinsics.dtype.kind not in "fiu":
|
||||
raise ValueError("Camera arrays must contain real numeric values.")
|
||||
if not np.isfinite(poses).all() or not np.isfinite(intrinsics).all():
|
||||
raise ValueError("Camera arrays must contain only finite values.")
|
||||
return int(poses.shape[0])
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate a LingBot-World v2 video from a first-frame image and camera trajectory."
|
||||
)
|
||||
parser.add_argument("--model", default=_MODEL, help="Official Hugging Face model ID or local checkpoint path.")
|
||||
parser.add_argument("--prompt", required=True, help="Text description of the scene and desired motion.")
|
||||
parser.add_argument("--image", required=True, help="Path to the first-frame image.")
|
||||
parser.add_argument(
|
||||
"--action-dir",
|
||||
required=True,
|
||||
help="Camera action directory containing poses.npy and intrinsics.npy.",
|
||||
)
|
||||
parser.add_argument("--height", type=int, default=480, help="Output height; must be divisible by 16.")
|
||||
parser.add_argument("--width", type=int, default=832, help="Output width; must be divisible by 16.")
|
||||
parser.add_argument(
|
||||
"--num-frames",
|
||||
type=int,
|
||||
default=9,
|
||||
help="Raw output frame count. Nine frames are one three-latent-frame causal block.",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument(
|
||||
"--tensor-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of GPUs used for tensor parallelism inside the DiT.",
|
||||
)
|
||||
parser.add_argument("--flow-shift", type=float, default=5.0, help="Positive FlowUniPC scheduler shift.")
|
||||
parser.add_argument("--fps", type=int, default=16, help="Frames per second in the exported MP4.")
|
||||
parser.add_argument("--output", default="lingbot_world_v2.mp4", help="Output MP4 path.")
|
||||
parser.add_argument(
|
||||
"--enforce-eager",
|
||||
action="store_true",
|
||||
help="Disable compilation and force eager execution.",
|
||||
)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def resolve_cli_paths(args: argparse.Namespace) -> LingBotPaths:
|
||||
"""Canonicalize local inputs and construct a safe trusted-root mapping."""
|
||||
|
||||
image = Path(args.image).expanduser().resolve()
|
||||
if not image.is_file():
|
||||
raise ValueError("--image must point to an existing file.")
|
||||
|
||||
action_dir = Path(args.action_dir).expanduser().resolve()
|
||||
if not action_dir.is_dir():
|
||||
raise ValueError("--action-dir must point to an existing directory.")
|
||||
if not all((action_dir / name).is_file() for name in ("poses.npy", "intrinsics.npy")):
|
||||
raise ValueError("--action-dir must contain poses.npy and intrinsics.npy.")
|
||||
camera_frames = _validate_camera_arrays(action_dir)
|
||||
|
||||
action_root = action_dir.parent.resolve()
|
||||
try:
|
||||
action_relative = action_dir.relative_to(action_root)
|
||||
except ValueError:
|
||||
raise ValueError("--action-dir must be contained by its canonical trusted action root.") from None
|
||||
if not action_relative.parts:
|
||||
raise ValueError("--action-dir must name a directory contained by a trusted action root.")
|
||||
|
||||
output = Path(args.output).expanduser().resolve()
|
||||
return LingBotPaths(
|
||||
image=image,
|
||||
action_dir=action_dir,
|
||||
action_root=action_root,
|
||||
action_relative=action_relative,
|
||||
camera_frames=camera_frames,
|
||||
output=output,
|
||||
)
|
||||
|
||||
|
||||
def _positive_finite(value: float, flag: str) -> float:
|
||||
result = float(value)
|
||||
if not math.isfinite(result) or result <= 0:
|
||||
raise ValueError(f"{flag} must be a positive finite number.")
|
||||
return result
|
||||
|
||||
|
||||
def build_omni_kwargs(
|
||||
args: argparse.Namespace,
|
||||
paths: LingBotPaths,
|
||||
*,
|
||||
parallel_config: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""Build Omni configuration without overriding checkpoint class discovery."""
|
||||
|
||||
if args.tensor_parallel_size <= 0:
|
||||
raise ValueError("--tensor-parallel-size must be a positive integer.")
|
||||
flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
|
||||
model_path = Path(args.model).expanduser()
|
||||
model = str(model_path.resolve()) if model_path.exists() else args.model
|
||||
return {
|
||||
"model": model,
|
||||
"flow_shift": flow_shift,
|
||||
"parallel_config": parallel_config,
|
||||
"enforce_eager": args.enforce_eager,
|
||||
"model_config": {"lingbot_action_root": str(paths.action_root)},
|
||||
}
|
||||
|
||||
|
||||
def build_request(
|
||||
args: argparse.Namespace,
|
||||
paths: LingBotPaths,
|
||||
) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
"""Build the prompt and fixed v1 causal-DMD sampling contract."""
|
||||
|
||||
prompt_text = args.prompt.strip()
|
||||
if not prompt_text:
|
||||
raise ValueError("--prompt must contain non-whitespace text.")
|
||||
if args.height <= 0 or args.width <= 0:
|
||||
raise ValueError("--height and --width must be positive integers.")
|
||||
if args.height % _SPATIAL_ALIGNMENT or args.width % _SPATIAL_ALIGNMENT:
|
||||
raise ValueError("--height and --width must each be divisible by 16.")
|
||||
if args.height * args.width > _MAX_PIXEL_AREA:
|
||||
raise ValueError("--height * --width must not exceed 480 * 832 pixels.")
|
||||
if args.num_frames <= 0:
|
||||
raise ValueError("--num-frames must be a positive integer.")
|
||||
if args.num_frames > _MAX_RAW_FRAMES:
|
||||
raise ValueError("--num-frames must not exceed 117 raw frames.")
|
||||
if args.num_frames > paths.camera_frames:
|
||||
raise ValueError("Camera arrays must contain at least --num-frames frames.")
|
||||
if (args.num_frames - 1) % _TEMPORAL_COMPRESSION:
|
||||
raise ValueError("--num-frames must satisfy (num_frames - 1) divisible by 4.")
|
||||
latent_frames = (args.num_frames - 1) // _TEMPORAL_COMPRESSION + 1
|
||||
if latent_frames % _LATENT_FRAMES_PER_BLOCK:
|
||||
raise ValueError("--num-frames must map to whole three-frame latent blocks.")
|
||||
if args.fps <= 0:
|
||||
raise ValueError("--fps must be a positive integer.")
|
||||
flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
|
||||
|
||||
prompt = {
|
||||
"prompt": prompt_text,
|
||||
"multi_modal_data": {"image": str(paths.image)},
|
||||
}
|
||||
sampling_kwargs = {
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_frames": args.num_frames,
|
||||
"num_inference_steps": _NUM_INFERENCE_STEPS,
|
||||
"max_sequence_length": _MAX_SEQUENCE_LENGTH,
|
||||
"seed": args.seed,
|
||||
"fps": args.fps,
|
||||
"extra_args": {
|
||||
"action_path": paths.action_relative.as_posix(),
|
||||
"flow_shift": flow_shift,
|
||||
},
|
||||
}
|
||||
return prompt, sampling_kwargs
|
||||
|
||||
|
||||
def extract_video_array(outputs: Any) -> np.ndarray:
|
||||
"""Unwrap one standard Omni diffusion video into ``[frames, H, W, C]``."""
|
||||
|
||||
if not isinstance(outputs, list) or len(outputs) != 1:
|
||||
raise ValueError("Expected exactly one Omni output for the single LingBot request.")
|
||||
output = outputs[0]
|
||||
request_output = getattr(output, "request_output", None)
|
||||
if request_output is not None and getattr(request_output, "images", None):
|
||||
output = request_output
|
||||
|
||||
images = getattr(output, "images", None)
|
||||
if not images:
|
||||
multimodal_output = getattr(output, "multimodal_output", None) or {}
|
||||
images = multimodal_output.get("video")
|
||||
if images is None:
|
||||
raise ValueError("No video frames found in Omni output.")
|
||||
if isinstance(images, list) and len(images) == 1:
|
||||
images = images[0]
|
||||
|
||||
video: np.ndarray = np.asarray(images)
|
||||
if video.ndim == 5:
|
||||
if video.shape[0] != 1:
|
||||
raise ValueError(f"Expected one generated video, got batch shape {video.shape}.")
|
||||
video = video[0]
|
||||
if video.ndim != 4 or video.shape[-1] not in (3, 4):
|
||||
raise ValueError(f"Expected video shape [frames, height, width, channels], got {video.shape}.")
|
||||
return video
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> Path:
|
||||
args = parse_args(argv)
|
||||
paths = resolve_cli_paths(args)
|
||||
prompt, sampling_kwargs = build_request(args, paths)
|
||||
paths.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Keep vLLM imports below pure CLI/path validation so ``--help`` and the
|
||||
# helper tests work on machines without the native vLLM runtime.
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
parallel_config = DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size)
|
||||
omni_kwargs = build_omni_kwargs(args, paths, parallel_config=parallel_config)
|
||||
omni = Omni(**omni_kwargs)
|
||||
try:
|
||||
outputs = omni.generate(
|
||||
prompt,
|
||||
OmniDiffusionSamplingParams(**sampling_kwargs),
|
||||
use_tqdm=False,
|
||||
)
|
||||
finally:
|
||||
omni.close()
|
||||
|
||||
frames = extract_video_array(outputs)
|
||||
export_to_video(frames, str(paths.output), fps=args.fps)
|
||||
print(f"Saved {frames.shape[0]} frames to {paths.output}")
|
||||
return paths.output
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,239 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Run LingBot-World 2.0 as an in-process realtime AR-Diffusion session.
|
||||
|
||||
This example exercises the internal session API. It is not a public HTTP or
|
||||
WebSocket protocol. Each JSONL event advances the world by one three-latent
|
||||
frame AR block and writes the latent plus its identity metadata. The current
|
||||
scoped implementation supports at most ten ticks per generation epoch because
|
||||
its image condition is bounded to 117 pixel frames; reset or create a session
|
||||
to start a new world.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
|
||||
_CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1"
|
||||
_FRAMES_PER_BLOCK = 3
|
||||
# ((117 pixel frames - 1) / VAE temporal factor 4 + 1) / 3 latent frames.
|
||||
_MAX_REALTIME_TICKS = 10
|
||||
|
||||
|
||||
def _camera_event_data(frames: list[list[str]]) -> dict[str, Any]:
|
||||
"""Build the event-side script consumed by LingBotCameraControlReducer."""
|
||||
return {"mode": "script", "frames": frames}
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Run an in-process realtime LingBot-World 2.0 session.")
|
||||
parser.add_argument("--model", default=_MODEL, help="Hugging Face model ID or local checkpoint path.")
|
||||
parser.add_argument("--image", required=True, help="Initial RGB image.")
|
||||
parser.add_argument("--prompt", required=True, help="Initial scene prompt.")
|
||||
parser.add_argument(
|
||||
"--events",
|
||||
required=True,
|
||||
help="JSONL file with one prompt/action event per AR block; at most 10 events.",
|
||||
)
|
||||
parser.add_argument("--output-dir", required=True, help="Directory for chunk latents and metadata.")
|
||||
parser.add_argument("--session-id", default="lingbot-world", help="Persistent world session identifier.")
|
||||
parser.add_argument("--height", type=int, default=480)
|
||||
parser.add_argument("--width", type=int, default=832)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--gpu-memory-fraction", type=float, default=0.1)
|
||||
parser.add_argument("--tensor-parallel-size", type=int, default=1)
|
||||
parser.add_argument("--enforce-eager", action="store_true")
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def _load_events(path: Path) -> list[dict[str, Any]]:
|
||||
events: list[dict[str, Any]] = []
|
||||
for line_number, raw_line in enumerate(path.read_text().splitlines(), start=1):
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
value = json.loads(line)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"events line {line_number} is not valid JSON: {exc.msg}.") from None
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError(f"events line {line_number} must be a JSON object.")
|
||||
event_id = value.get("event_id")
|
||||
if isinstance(event_id, bool) or not isinstance(event_id, int) or event_id < 0:
|
||||
raise ValueError(f"events line {line_number} requires a non-negative integer event_id.")
|
||||
prompt = value.get("prompt")
|
||||
if prompt is not None and (not isinstance(prompt, str) or not prompt.strip()):
|
||||
raise ValueError(f"events line {line_number} prompt must be a non-empty string.")
|
||||
frames = value.get("frames")
|
||||
if frames is not None:
|
||||
if not isinstance(frames, list) or len(frames) != _FRAMES_PER_BLOCK:
|
||||
raise ValueError(f"events line {line_number} frames must contain exactly three lists.")
|
||||
for frame in frames:
|
||||
if not isinstance(frame, list) or any(not isinstance(action, str) for action in frame):
|
||||
raise ValueError(f"events line {line_number} frames must contain only action strings.")
|
||||
if prompt is None and frames is None:
|
||||
raise ValueError(f"events line {line_number} must update prompt and/or frames.")
|
||||
events.append({"event_id": event_id, "prompt": prompt, "frames": frames})
|
||||
if not events:
|
||||
raise ValueError("events file must contain at least one event.")
|
||||
if len(events) > _MAX_REALTIME_TICKS:
|
||||
raise ValueError(
|
||||
"events file must contain at most "
|
||||
f"{_MAX_REALTIME_TICKS} events because the current LingBot realtime "
|
||||
"image-condition horizon is 117 pixel frames."
|
||||
)
|
||||
event_ids = [event["event_id"] for event in events]
|
||||
if event_ids != sorted(set(event_ids)):
|
||||
raise ValueError("event_id values must be unique and strictly increasing.")
|
||||
return events
|
||||
|
||||
|
||||
def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path]:
|
||||
image = Path(args.image).expanduser().resolve()
|
||||
events = Path(args.events).expanduser().resolve()
|
||||
output_dir = Path(args.output_dir).expanduser().resolve()
|
||||
if not image.is_file():
|
||||
raise ValueError("--image must point to an existing file.")
|
||||
if not events.is_file():
|
||||
raise ValueError("--events must point to an existing JSONL file.")
|
||||
if not args.prompt.strip():
|
||||
raise ValueError("--prompt must contain non-whitespace text.")
|
||||
if args.height <= 0 or args.width <= 0 or args.height % 16 or args.width % 16:
|
||||
raise ValueError("--height and --width must be positive multiples of 16.")
|
||||
if args.tensor_parallel_size <= 0:
|
||||
raise ValueError("--tensor-parallel-size must be positive.")
|
||||
if not math.isfinite(args.gpu_memory_fraction) or not 0 < args.gpu_memory_fraction <= 1:
|
||||
raise ValueError("--gpu-memory-fraction must be in (0, 1].")
|
||||
return image, events, output_dir
|
||||
|
||||
|
||||
async def run(argv: Sequence[str] | None = None) -> Path:
|
||||
args = parse_args(argv)
|
||||
image, events_path, output_dir = _validate_args(args)
|
||||
events = _load_events(events_path)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Imports stay below pure input validation so --help and helper tests do not
|
||||
# require a CUDA-enabled vLLM installation.
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.models.lingbot_world.actions import LingBotCameraControlReducer
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
|
||||
from vllm_omni.experimental.ar_diffusion.session import (
|
||||
ARDiffusionSessionEvent,
|
||||
ARDiffusionSessionManager,
|
||||
ARDiffusionWorkerLifecycle,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionControlInput
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
engine = AsyncOmni(
|
||||
model=args.model,
|
||||
engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine",
|
||||
enforce_eager=args.enforce_eager,
|
||||
parallel_config=DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size),
|
||||
max_num_seqs=1,
|
||||
model_config={
|
||||
"ar_diffusion_height": args.height,
|
||||
"ar_diffusion_width": args.width,
|
||||
"ar_diffusion_kv_config": {
|
||||
"gpu_memory_fraction": args.gpu_memory_fraction,
|
||||
"warmup_cudagraph": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
sampling = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=9,
|
||||
num_inference_steps=4,
|
||||
max_sequence_length=512,
|
||||
seed=args.seed,
|
||||
output_type="latent",
|
||||
extra_args={"flow_shift": 5.0},
|
||||
)
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
engine,
|
||||
prompt_provider=lambda tick: {
|
||||
"prompt": tick.prompt,
|
||||
"multi_modal_data": {"image": str(image)},
|
||||
},
|
||||
sampling_params_list=[sampling],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=consumer,
|
||||
lifecycle=ARDiffusionWorkerLifecycle(engine, stage_ids=[0], timeout=180.0),
|
||||
max_pending_events=32,
|
||||
control_reducer_factory=LingBotCameraControlReducer,
|
||||
)
|
||||
session = await manager.create_session(args.session_id)
|
||||
measurements: list[dict[str, Any]] = []
|
||||
try:
|
||||
for chunk_index, event in enumerate(events):
|
||||
controls = ()
|
||||
if event["frames"] is not None:
|
||||
controls = (
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema=_CAMERA_ACTION_SCHEMA,
|
||||
data=_camera_event_data(event["frames"]),
|
||||
),
|
||||
)
|
||||
await session.accept_event(
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=event["event_id"],
|
||||
prompt=event["prompt"]
|
||||
if event["prompt"] is not None
|
||||
else (args.prompt if chunk_index == 0 else None),
|
||||
controls=controls,
|
||||
)
|
||||
)
|
||||
started = time.perf_counter()
|
||||
output = await session.next_chunk()
|
||||
elapsed = time.perf_counter() - started
|
||||
if len(output.images) != 1 or not isinstance(output.images[0], torch.Tensor):
|
||||
raise RuntimeError("Expected one latent tensor from each realtime LingBot chunk.")
|
||||
latent = output.images[0].detach().float().cpu()
|
||||
metadata = consumer.chunk_metadata(output).to_dict()
|
||||
torch.save(latent, output_dir / f"chunk_{chunk_index:03d}.pt")
|
||||
(output_dir / f"chunk_{chunk_index:03d}.json").write_text(
|
||||
json.dumps(metadata, indent=2, sort_keys=True) + "\n"
|
||||
)
|
||||
measurements.append(
|
||||
{
|
||||
"chunk_index": chunk_index,
|
||||
"latency_seconds": elapsed,
|
||||
"shape": list(latent.shape),
|
||||
"finite": bool(torch.isfinite(latent).all()),
|
||||
"metadata": metadata,
|
||||
}
|
||||
)
|
||||
print(json.dumps(measurements[-1], sort_keys=True), flush=True)
|
||||
finally:
|
||||
try:
|
||||
await manager.close_session(args.session_id)
|
||||
finally:
|
||||
engine.shutdown()
|
||||
|
||||
summary_path = output_dir / "summary.json"
|
||||
summary_path.write_text(json.dumps({"chunks": measurements}, indent=2, sort_keys=True) + "\n")
|
||||
return summary_path
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> Path:
|
||||
return asyncio.run(run(argv))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -82,6 +82,9 @@ def parse_args():
|
||||
# ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
|
||||
# and ``create_default_diffusion`` materializes the stage config.
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(p)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -67,7 +67,7 @@ def main() -> None:
|
||||
parser.add_argument("--output-dir", default="results/nemotron_voicechat")
|
||||
parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
"--stage-configs-path",
|
||||
default=None,
|
||||
help=(
|
||||
"Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml "
|
||||
@@ -134,8 +134,8 @@ def main() -> None:
|
||||
]
|
||||
|
||||
omni_kwargs: dict = {}
|
||||
if args.deploy_config:
|
||||
omni_kwargs["deploy_config"] = args.deploy_config
|
||||
if args.stage_configs_path:
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path
|
||||
omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs)
|
||||
outputs = omni.generate(inputs, sampling_params_list)
|
||||
|
||||
|
||||
@@ -38,6 +38,17 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_MODEL = "fishaudio/s2-pro"
|
||||
DEFAULT_DEPLOY_CONFIG = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
"..",
|
||||
"..",
|
||||
"..",
|
||||
"..",
|
||||
"vllm_omni",
|
||||
"model_executor",
|
||||
"stage_configs",
|
||||
"fish_speech_s2_pro.yaml",
|
||||
)
|
||||
|
||||
|
||||
def build_prompt(
|
||||
@@ -141,8 +152,7 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
if args.deploy_config:
|
||||
omni_kwargs["deploy_config"] = args.deploy_config
|
||||
omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
|
||||
return omni_kwargs
|
||||
|
||||
|
||||
|
||||
@@ -32,10 +32,11 @@ from end2end import _estimate_prompt_len
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config
|
||||
|
||||
|
||||
def _default_deploy_config() -> str:
|
||||
def _default_stage_config() -> str:
|
||||
repo_root = Path(__file__).resolve().parents[4]
|
||||
return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
|
||||
|
||||
@@ -108,7 +109,7 @@ def main(args: Any) -> None:
|
||||
"output_dir",
|
||||
):
|
||||
omni_kwargs.pop(key, None)
|
||||
omni_kwargs["deploy_config"] = args.deploy_config or _default_deploy_config()
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
|
||||
omni_kwargs["log_stats"] = args.log_stats
|
||||
omni = Omni(**omni_kwargs)
|
||||
|
||||
@@ -155,7 +156,7 @@ def parse_args() -> Any:
|
||||
default=None,
|
||||
help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
|
||||
)
|
||||
parser.add_argument("--deploy-config", default=None, help="Qwen3-TTS deploy YAML")
|
||||
parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
|
||||
parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align")
|
||||
parser.add_argument("--language", default="English", help="Qwen3-TTS language field")
|
||||
parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name")
|
||||
@@ -163,6 +164,7 @@ def parse_args() -> Any:
|
||||
parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens")
|
||||
parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar")
|
||||
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging")
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -1,104 +0,0 @@
|
||||
# π0 (Pi-Zero) VLA — OpenPI realtime serving
|
||||
|
||||
[π0](https://www.physicalintelligence.company/blog/pi0) is a Vision-Language-Action
|
||||
model from Physical Intelligence: multi-camera images + a language instruction +
|
||||
robot proprioceptive state → a continuous action chunk via flow-matching denoising
|
||||
(it does **not** emit text tokens). This example serves π0 over the OpenPI realtime
|
||||
websocket protocol at `/v1/realtime/robot/openpi`.
|
||||
|
||||
## Install extras
|
||||
|
||||
The core `pip install -e .` does not include the OpenPI client used here:
|
||||
|
||||
- `openpi-client` (from the [openpi](https://github.com/Physical-Intelligence/openpi)
|
||||
repo: `pip install -e packages/openpi-client`), `websockets`, `msgpack`, `msgpack-numpy`
|
||||
|
||||
## Weights
|
||||
|
||||
π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the server
|
||||
download it (`MODEL=lerobot/pi0_base`) or point at a local copy
|
||||
(`MODEL=/path/to/pi0_base`).
|
||||
|
||||
## Run the server
|
||||
|
||||
```bash
|
||||
vllm serve lerobot/pi0_base --omni --port 8000 \
|
||||
--served-model-name pi0 \
|
||||
--deploy-config vllm_omni/deploy/pi0.yaml \
|
||||
--enforce-eager --disable-log-stats
|
||||
```
|
||||
|
||||
The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
|
||||
(`Pi0Pipeline`), float32, `max_num_seqs: 1`, and the `policy_server_config` handshake
|
||||
metadata (3 cameras, `joint_position`, action horizon 50, action dim 32).
|
||||
|
||||
### Commonly adjusted `model_config` keys
|
||||
|
||||
These live under `stages[0].model_config` in `vllm_omni/deploy/pi0.yaml`:
|
||||
|
||||
| Key | Default | Meaning |
|
||||
|---|---|---|
|
||||
| `chunk_size` | `50` | Action-chunk length (timesteps) the model predicts per inference. |
|
||||
| `num_inference_steps` | `10` | Flow-matching Euler denoising steps. |
|
||||
| `max_action_dim` | `32` | Action dimensionality (state/action are padded to this). |
|
||||
| `max_state_dim` | `32` | Proprioceptive-state dimensionality (zero-padded to this). |
|
||||
| `image_resolution` | `[224, 224]` | Per-camera input size (square; SigLIP). |
|
||||
| `tokenizer_max_length` | `48` | Max PaliGemma prompt tokens. |
|
||||
| `max_cameras` | `3` | Camera slots the model attends to (real + `-1`-padded). |
|
||||
| `image_feature_keys` | 3 `observation.images.*` keys | Camera order the model attends to. |
|
||||
| `image_key_map` | `{}` | Map raw obs camera keys → `image_feature_keys` (empty = verbatim). |
|
||||
|
||||
The `policy_server_config` block below them is the OpenPI handshake metadata
|
||||
advertised to the client; keep its `action_horizon` / `action_dim` /
|
||||
`image_resolution` in sync with the `model_config` values above. The e2e test
|
||||
`tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online` connects to a
|
||||
live server and asserts the advertised metadata matches these `pi0.yaml` values.
|
||||
|
||||
## Run the client
|
||||
|
||||
```bash
|
||||
python examples/online_serving/pi0/openpi_client.py --host 127.0.0.1 --port 8000 \
|
||||
--prompt "pick up the red block and place it in the bin"
|
||||
```
|
||||
|
||||
It connects, prints the server metadata, sends robot observations, and prints the
|
||||
returned `[action_horizon, action_dim] = [50, 32]` action chunks. Replace the blank
|
||||
cameras / zero state in `_make_dummy_obs` with real frames (HWC uint8) and
|
||||
proprioceptive state to drive a robot.
|
||||
|
||||
### Observation format
|
||||
|
||||
The client sends a flat dict per inference:
|
||||
|
||||
```python
|
||||
{
|
||||
"observation.images.base_0_rgb": np.uint8[H, W, 3],
|
||||
"observation.images.left_wrist_0_rgb": np.uint8[H, W, 3],
|
||||
"observation.images.right_wrist_0_rgb":np.uint8[H, W, 3],
|
||||
"state": np.float32[state_dim], # zero-padded to max_state_dim=32 server-side
|
||||
"prompt": "pick up the red block",
|
||||
"session_id": "<uuid>", # accepted but ignored (π0 is stateless)
|
||||
}
|
||||
```
|
||||
|
||||
Camera keys must match the server's `image_feature_keys` (the checkpoint's
|
||||
`input_features` order). If your robot uses different camera names, set
|
||||
`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
|
||||
|
||||
## Correctness
|
||||
|
||||
π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy` reference
|
||||
(`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
|
||||
`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`) and is
|
||||
version-stable across transformers releases (the version-stability checks in
|
||||
`tests/diffusion/models/pi0/test_pi0_units.py`). An OpenPI websocket online-serving e2e
|
||||
lives in `tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online`.
|
||||
|
||||
## Limitations
|
||||
|
||||
- **Normalization stats**: `lerobot/pi0_base` uses identity normalization, which is
|
||||
fully supported (state passes through, actions are returned in the model's space).
|
||||
Per-dataset `norm_stats` declared in a checkpoint's `config.json` are honored, but
|
||||
stats stored only in LeRobot's `policy_preprocessor.json` companion safetensors are
|
||||
not yet loaded — a fine-tuned checkpoint that relies on those would need that bridge
|
||||
added before its actions are in real-world units.
|
||||
@@ -1,119 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Minimal OpenPI client for the π0 (Pi-Zero) VLA policy server.
|
||||
|
||||
Connects to a running π0 server (see README.md), sends a single robot
|
||||
observation (3 cameras + state + language instruction), and prints the returned
|
||||
continuous action chunk ``[action_horizon, action_dim]``.
|
||||
|
||||
The server speaks the OpenPI websocket protocol:
|
||||
connect -> server sends msgpack(policy_server_config metadata)
|
||||
infer -> client sends msgpack(obs), server replies msgpack(actions ndarray)
|
||||
|
||||
Usage:
|
||||
python openpi_client.py --host 127.0.0.1 --port 8000 \
|
||||
--prompt "pick up the red block"
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
# Avoid shadowing the installed ``openpi_client`` package with this dir.
|
||||
example_dir = str(Path(__file__).resolve().parent)
|
||||
removed = sys.path and sys.path[0] == example_dir
|
||||
if removed:
|
||||
sys.path.pop(0)
|
||||
try:
|
||||
from openpi_client.websocket_client_policy import WebsocketClientPolicy
|
||||
finally:
|
||||
if removed:
|
||||
sys.path.insert(0, example_dir)
|
||||
except ImportError as exc: # pragma: no cover - runtime dependency guard
|
||||
raise ImportError("π0 OpenPI example requires `openpi-client`.") from exc
|
||||
|
||||
logger = logging.getLogger("pi0_openpi_client")
|
||||
|
||||
DEFAULT_HOST = "127.0.0.1"
|
||||
DEFAULT_PORT = 8000
|
||||
DEFAULT_PATH = "/v1/realtime/robot/openpi"
|
||||
DEFAULT_PROMPT = "pick up the red block and place it in the bin"
|
||||
|
||||
# π0 / pi0_base camera identities (must match the server's image_feature_keys).
|
||||
CAMERA_KEYS = (
|
||||
"observation.images.base_0_rgb",
|
||||
"observation.images.left_wrist_0_rgb",
|
||||
"observation.images.right_wrist_0_rgb",
|
||||
)
|
||||
STATE_DIM = 32
|
||||
|
||||
|
||||
def _make_dummy_obs(prompt: str, session_id: str, image_size: int = 224) -> dict[str, Any]:
|
||||
"""Build a single π0 observation with blank cameras + zero state.
|
||||
|
||||
Replace the zeros with real camera frames (HWC uint8) and proprioceptive
|
||||
state to drive an actual robot.
|
||||
"""
|
||||
obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in CAMERA_KEYS}
|
||||
obs["state"] = np.zeros(STATE_DIM, dtype=np.float32)
|
||||
obs["prompt"] = prompt
|
||||
obs["session_id"] = session_id
|
||||
return obs
|
||||
|
||||
|
||||
def _ws_uri(host: str, port: int, path: str) -> str:
|
||||
"""Build the full websocket URI. WebsocketClientPolicy takes the whole URI as
|
||||
``host`` (it has no separate path arg)."""
|
||||
if host.startswith(("ws://", "wss://")):
|
||||
return host
|
||||
return f"ws://{host}:{port}{path}"
|
||||
|
||||
|
||||
def _as_action_array(response: Any) -> np.ndarray:
|
||||
if isinstance(response, dict) and response.get("type") == "error":
|
||||
raise RuntimeError(f"Inference failed: {response.get('message', response)}")
|
||||
return np.asarray(response, dtype=np.float32)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description="π0 OpenPI client")
|
||||
parser.add_argument("--host", default=DEFAULT_HOST)
|
||||
parser.add_argument("--port", type=int, default=DEFAULT_PORT)
|
||||
parser.add_argument("--path", default=DEFAULT_PATH)
|
||||
parser.add_argument("--prompt", default=DEFAULT_PROMPT)
|
||||
parser.add_argument("--num-steps", type=int, default=2, help="Inference observations to send.")
|
||||
args = parser.parse_args()
|
||||
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
||||
|
||||
policy = WebsocketClientPolicy(host=_ws_uri(args.host, args.port, args.path))
|
||||
metadata = dict(policy.get_server_metadata())
|
||||
logger.info("Server metadata: %s", metadata)
|
||||
|
||||
session_id = str(uuid.uuid4())
|
||||
for step in range(args.num_steps):
|
||||
obs = _make_dummy_obs(args.prompt, session_id)
|
||||
actions = _as_action_array(policy.infer(obs))
|
||||
logger.info(
|
||||
"[step %d] actions shape=%s mean=%.4f std=%.4f",
|
||||
step,
|
||||
actions.shape,
|
||||
float(actions.mean()),
|
||||
float(actions.std()),
|
||||
)
|
||||
if not np.isfinite(actions).all():
|
||||
raise RuntimeError("Server returned non-finite actions.")
|
||||
logger.info("Done — π0 OpenPI round-trip OK.")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -79,7 +79,6 @@ plugins:
|
||||
user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
|
||||
user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
|
||||
user_guide/quantization/index.md: user_guide/quantization/overview.md
|
||||
configuration/pd_disaggregation.md: features/pd_disaggregation.md
|
||||
design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
|
||||
design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
|
||||
design/feature/cpu_offloading.md: design/feature/offloader/README.md
|
||||
|
||||
@@ -45,7 +45,7 @@ MusicGen-style delay pattern).
|
||||
**Online serving:**
|
||||
|
||||
```bash
|
||||
vllm serve bosonai/higgs-audio-v3-tts-4b \
|
||||
vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
|
||||
--host 0.0.0.0 --port 8095 \
|
||||
--trust-remote-code --omni
|
||||
```
|
||||
|
||||
@@ -81,7 +81,7 @@ MammothModa2 generation parameters as a JSON object through `--extra-body`:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model ./MammothModa2-Preview \
|
||||
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
|
||||
--stage-configs-path vllm_omni/deploy/mammoth_moda2.yaml \
|
||||
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
@@ -132,44 +132,6 @@ ls -lh mammoth_t2i.png
|
||||
python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)"
|
||||
```
|
||||
|
||||
### 1x AMD MI300X, MammothModa2 Preview
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux 6.8.0-134-generic, x86_64
|
||||
- Container: official ROCm image built from `docker/Dockerfile.rocm`
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.11.0+gitd0c8b1f
|
||||
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
|
||||
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
|
||||
- vLLM version: 0.27.0+rocm723
|
||||
- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
|
||||
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
|
||||
|
||||
#### Offline Commands
|
||||
|
||||
The checked run used the committed stage split, with `gpu_memory_utilization` set to 0.5 for AR and 0.3 for DiT:
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model bytedance-research/MammothModa2-Preview \
|
||||
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
|
||||
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--seed 42 \
|
||||
--extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
|
||||
--enable-diffusion-pipeline-profiler \
|
||||
--log-stats \
|
||||
--output mammoth_t2i.png
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting.
|
||||
|
||||
The output was a valid 1024 by 1024 RGB PNG.
|
||||
|
||||
## MammothModa2-Dev unified inference
|
||||
|
||||
MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses
|
||||
|
||||
@@ -51,7 +51,7 @@ NemotronH (~3B active parameters). Pick the pipeline by task:
|
||||
`model_type` auto-resolves `audex_tts.yaml`):
|
||||
|
||||
```bash
|
||||
vllm serve nvidia/Nemotron-Labs-Audex-2B \
|
||||
vllm-omni serve nvidia/Nemotron-Labs-Audex-2B \
|
||||
--host 0.0.0.0 --port 8097 \
|
||||
--trust-remote-code --omni
|
||||
# other modes: add --stage-configs-path vllm_omni/deploy/audex_{tta,thinker_only,s2s}.yaml
|
||||
@@ -109,7 +109,7 @@ The 30B REQUIRES its explicit deploy yaml (the shared HF `model_type`
|
||||
auto-resolves to the 2B-tuned config otherwise):
|
||||
|
||||
```bash
|
||||
vllm serve nvidia/Nemotron-Labs-Audex-30B-A3B \
|
||||
vllm-omni serve nvidia/Nemotron-Labs-Audex-30B-A3B \
|
||||
--host 0.0.0.0 --port 8097 \
|
||||
--trust-remote-code --omni \
|
||||
--stage-configs-path vllm_omni/deploy/audex_tts_30b.yaml
|
||||
|
||||
@@ -71,35 +71,6 @@ single-GPU entry point; `minicpmo_4_5_2gpu.yaml` is the recommended
|
||||
two-GPU profile. The removed fused two-stage implementation is not retained as
|
||||
a fallback because it would duplicate state machines and correctness paths.
|
||||
|
||||
### Graph execution
|
||||
|
||||
Graph boundaries follow each stage's state model:
|
||||
|
||||
| Stage | Graph mode on Ascend | Eager boundary |
|
||||
| --- | --- | --- |
|
||||
| 0 Thinker | vLLM `PIECEWISE` | multimodal preprocessing and output routing |
|
||||
| 1 Talker | vLLM `PIECEWISE` | conditioning preprocessing, codec sampling, request-state updates |
|
||||
| 2 Code2Wav | inner exact-shape NPUGraph for the CFM DiT estimator | encoder, timestep embedding, HiFT/RNG, request parsing, stream-state commit |
|
||||
|
||||
Stage 2 keeps `enforce_eager: true` for the generation runner because its
|
||||
Python request metadata and per-request cache dictionaries are not valid
|
||||
outer-graph inputs. The backend captures only the deterministic CFM DiT
|
||||
estimator, with an exact graph key for each tensor shape. Timestep embedding
|
||||
stays eager because the upstream implementation creates a CPU frequency tensor;
|
||||
HiFT stays eager because it generates random phase. The graph cache stores up
|
||||
to 32 entries by default, after which unseen shapes run eagerly.
|
||||
|
||||
An ACL graph capture failure is fatal to that Stage-2 process because older
|
||||
torch-npu releases can leave allocator and RNG capture state invalid. Restart
|
||||
the service after a capture failure. To run without capture, set
|
||||
`code2wav_enable_npu_graph: false` under the Stage-2
|
||||
`platforms.npu.stages[].additional_config` block before startup. Tune the cache
|
||||
limit there with `code2wav_max_npu_graphs`. Graph mode also requires
|
||||
`ASCEND_LAUNCH_BLOCKING` to be unset or set to `0`.
|
||||
|
||||
The inner NPUGraph is independent of the outer runner setting, so do not use a
|
||||
global `--enforce-eager` override when Stage 0/1 `PIECEWISE` replay is desired.
|
||||
|
||||
## GPU
|
||||
|
||||
### 1 x GPU (default — single command)
|
||||
@@ -240,9 +211,9 @@ speech output (TTS)"** checkbox on / off.
|
||||
model processes still share one CUDA device.
|
||||
- `--trust-remote-code` is required — the HF repo ships a custom
|
||||
`MiniCPMO` config / model class.
|
||||
- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 keeps its
|
||||
request orchestration eager while using inner CFM and HiFT CUDA Graphs. On
|
||||
Ascend, the exact-shape CFM DiT estimator instead uses inner NPUGraph replay.
|
||||
- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 remains
|
||||
eager because its request-owned Flow/HiFT caches and variable chunk/cache
|
||||
shapes are not yet exposed through a static exact-shape graph wrapper.
|
||||
- All default stages use `max_num_seqs: 4` to reduce cross-process GPU
|
||||
contention. Talker AR
|
||||
state and Code2Wav caches are request-owned; Code2Wav batches only
|
||||
@@ -253,8 +224,8 @@ speech output (TTS)"** checkbox on / off.
|
||||
- `StageRequestStats.batch_size` is a request-scoped placeholder, not the
|
||||
scheduler's execution batch.
|
||||
- Single-GPU co-location trades throughput for hardware density: Stage 0/1
|
||||
CUDA Graph replay and Stage 2 vocoder kernels compete across three CUDA
|
||||
contexts. Use the 8x4090 config or a custom multi-GPU mapping for
|
||||
CUDA Graph replay and eager Stage 2 vocoder kernels compete across three
|
||||
CUDA contexts. Use the 8x4090 config or a custom multi-GPU mapping for
|
||||
throughput-sensitive serving.
|
||||
|
||||
### 8 x RTX 4090 24GB (consumer-GPU layout)
|
||||
|
||||
@@ -252,53 +252,3 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
- stage_id: 1
|
||||
gpu_memory_utilization: 0.15
|
||||
```
|
||||
|
||||
### 1x AMD MI300X, 1.7B checkpoints
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux 6.8.0-134-generic, x86_64
|
||||
- Container: official ROCm image built from `docker/Dockerfile.rocm`
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.11.0+gitd0c8b1f
|
||||
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
|
||||
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
|
||||
- vLLM version: 0.27.0+rocm723
|
||||
- vLLM Omni version or commit: `3fecb6953ca8dc51210cc0421ef24552267a41ef`
|
||||
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
|
||||
- ONNX Runtime: onnxruntime-rocm 1.22.2.post3 with `ROCMExecutionProvider`
|
||||
- transformers: 5.15.0
|
||||
|
||||
#### Command
|
||||
|
||||
The MI300X tests also worked without `enforce_eager`. To use that setup, remove `enforce_eager: true` from stage 1 under `platforms.rocm` in `vllm_omni/deploy/qwen3_tts.yaml`.
|
||||
|
||||
Set this environment variable before you start Qwen3-TTS:
|
||||
|
||||
```bash
|
||||
export MIOPEN_FIND_MODE=FAST
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
Each row records one offline request with one prompt and batch size one. Request time is the logged `e2e_wall_time_ms`, and stage 1 time is the logged `e2e_stage_1_wall_time_ms`.
|
||||
|
||||
| Checkpoint | Code2Wav configuration | `MIOPEN_FIND_MODE` | Request time | Stage 1 time | Output duration | Real time factor | Maximum sampled device memory |
|
||||
| --- | --- | --- | ---: | ---: | ---: | ---: | ---: |
|
||||
| CustomVoice | Eager | Unset | 1.683 s | 1.613 s | 5.68 s | 0.30 | 60.91 GiB |
|
||||
| CustomVoice | Graphs allowed | `FAST` | 1.326 s | 1.259 s | 5.68 s | 0.23 | 62.19 GiB |
|
||||
| VoiceDesign | Eager | Unset | 1.195 s | 1.180 s | 4.40 s | 0.27 | 60.91 GiB |
|
||||
| VoiceDesign | Graphs allowed | `FAST` | 1.127 s | 1.054 s | 4.40 s | 0.26 | 62.21 GiB |
|
||||
| Base | Eager | Unset | 13.330 s | 13.263 s | 4.32 s | 3.09 | 61.74 GiB |
|
||||
| Base | Graphs allowed | `FAST` | 14.314 s | 14.245 s | 4.80 s | 2.98 | 62.84 GiB |
|
||||
|
||||
All six processes finished successfully, and each output passed the 24 kHz mono WAV checks.
|
||||
|
||||
#### Notes
|
||||
|
||||
- The current ROCm config keeps Code2Wav in eager mode.
|
||||
- The tested alternative allows graphs and sets [`MIOPEN_FIND_MODE=FAST`](https://rocm.docs.amd.com/projects/MIOpen/en/develop/reference/env_variables.html). MIOpen uses a saved kernel choice when available and uses its immediate fallback otherwise.
|
||||
- CustomVoice used a graph for one of four logged Code2Wav batches. VoiceDesign used a graph for two of four batches. Base used no graphs for its four logged batches.
|
||||
- The recorded request time with graphs allowed was lower for CustomVoice and VoiceDesign and higher for Base. The Base runs generated different amounts of audio, so their request times are not a direct comparison.
|
||||
- Each setup was tested with one request, so the timing is not a performance benchmark.
|
||||
- GPU memory was sampled once per second with `rocm-smi`.
|
||||
+3
-5
@@ -49,7 +49,6 @@ recipes/
|
||||
| [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
|
||||
| [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
|
||||
| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
|
||||
| [`k2-fsa/OmniVoice.md`](./k2-fsa/OmniVoice.md) | Offline and online multilingual text-to-speech | 1x AMD MI300X (ROCm 7.2) |
|
||||
| [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
|
||||
| [`IndexTeam/IndexTTS-2_5.md`](./IndexTeam/IndexTTS-2_5.md) | Online + offline multilingual voice-cloned TTS with native speed and emotion control | 1x NVIDIA H20 96GB |
|
||||
| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
|
||||
@@ -61,7 +60,6 @@ recipes/
|
||||
| [`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 |
|
||||
| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) |
|
||||
| [`Robbyant/LingBot-World-2.0.md`](./Robbyant/LingBot-World-2.0.md) | Offline and experimental realtime interactive world generation | 1x H200/B200 or 2x B200 for TP=2 |
|
||||
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
|
||||
| [`cosmos3/Cosmos3-Edge.md`](./cosmos3/Cosmos3-Edge.md) | T2I / T2V / I2V + action (Physical-AI) generation on the Nemotron-based Edge transformer | 1x GPU (~8 GiB) |
|
||||
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
|
||||
@@ -74,9 +72,9 @@ recipes/
|
||||
| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
|
||||
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
|
||||
| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
|
||||
| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB / 1x AMD MI300X |
|
||||
| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X |
|
||||
| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X |
|
||||
| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB |
|
||||
| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) |
|
||||
| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB |
|
||||
| [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
|
||||
| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
|
||||
| [`Wan-AI/Wan2.2-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 1x datacenter Blackwell (B200/B300) |
|
||||
|
||||
@@ -1,122 +0,0 @@
|
||||
# LingBot-World 2.0
|
||||
|
||||
> Offline and experimental realtime interactive world generation
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Robbyant
|
||||
- Model: `robbyant/lingbot-world-v2-14b-causal-fast-diffusers`
|
||||
- Task: image-conditioned interactive world generation
|
||||
- Modes: offline trajectory replay and in-process realtime AR-Diffusion ticks
|
||||
- Hardware validated: NVIDIA H200 and B200
|
||||
- Maintainer: Community
|
||||
|
||||
The checkpoint is separately licensed under CC BY-NC-SA and restricted to
|
||||
non-commercial use. The vLLM-Omni integration code remains Apache-2.0.
|
||||
|
||||
## Offline generation
|
||||
|
||||
The offline path consumes one source image and an action directory containing
|
||||
`poses.npy` with shape `[frames, 4, 4]` and `intrinsics.npy` with shape
|
||||
`[frames, 4]`.
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/diffusion/lingbot_world_v2.py \
|
||||
--prompt "The camera moves slowly forward through the scene." \
|
||||
--image /path/to/first_frame.png \
|
||||
--action-dir /path/to/actions/forward \
|
||||
--num-frames 81 \
|
||||
--output lingbot_world_v2.mp4
|
||||
```
|
||||
|
||||
The current causal-fast checkpoint uses four DMD sampling steps. Raw frame
|
||||
counts must be `9 + 12k`, up to 117 frames. Camera files may contain a
|
||||
longer bounded source trajectory (the official example contains 269 frames);
|
||||
the runtime consumes only the prefix needed by the request.
|
||||
|
||||
## Realtime in-process generation
|
||||
|
||||
The realtime example keeps AR-Diffusion KV and model-owned state across
|
||||
requests. Each JSONL line describes the prompt and/or three latent-frame camera
|
||||
actions applied at the next chunk boundary:
|
||||
|
||||
```json
|
||||
{"event_id":1,"prompt":"A road through a forest","frames":[["j"],[],[]]}
|
||||
{"event_id":2,"frames":[["w"],["w"],["w"]]}
|
||||
{"event_id":3,"prompt":"The road enters a snowy valley","frames":[[],[],[]]}
|
||||
```
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/diffusion/lingbot_world_v2_realtime.py \
|
||||
--image /path/to/first_frame.png \
|
||||
--events /path/to/events.jsonl \
|
||||
--output-dir /tmp/lingbot-realtime \
|
||||
--gpu-memory-fraction 0.6
|
||||
```
|
||||
|
||||
The example writes one latent tensor and one metadata JSON file per chunk. It
|
||||
exercises the same `ARDiffusionSessionManager -> ARDiffusionOmniTickConsumer
|
||||
-> AsyncOmni -> ARDiffusionEngine` path used by a future HTTP or WebSocket
|
||||
transport.
|
||||
|
||||
This PR intentionally does not define a public realtime HTTP/WebSocket schema.
|
||||
An online serving client should be added together with the public transport
|
||||
API rather than exposing LingBot-specific event fields from the generic model
|
||||
runtime. The structured camera-interaction frontend and transport work is
|
||||
tracked separately in
|
||||
[vllm-project/vllm-omni#5527](https://github.com/vllm-project/vllm-omni/pull/5527).
|
||||
|
||||
## Realtime identity and controls
|
||||
|
||||
- `session_id` identifies the persistent world and its worker-owned state.
|
||||
- `event_id` identifies a prompt/control update and remains monotonic across reset.
|
||||
- `chunk_index` is contiguous from zero and restarts from zero after reset.
|
||||
- `request_id` correlates one chunk snapshot with its output metadata.
|
||||
- AsyncOmni uses a separate UUID-suffixed internal engine routing ID.
|
||||
|
||||
The generic runtime transports controls as opaque snapshots. LingBot's adapter
|
||||
accepts:
|
||||
|
||||
- `lingbot.camera_actions.v1` for per-latent-frame key states such as `w`, `a`,
|
||||
`s`, `d`, `i`, `j`, `k`, and `l`;
|
||||
- `lingbot.camera_trajectory.v1` for explicit pose/intrinsics trajectories.
|
||||
|
||||
## Validation
|
||||
|
||||
Real-checkpoint validation uses 480x832 output, four DMD steps, and seed 42.
|
||||
The exercised matrix includes:
|
||||
|
||||
- TP=1 and TP=2 execution;
|
||||
- two interleaved resident sessions;
|
||||
- action input and prompt switching;
|
||||
- seven contiguous chunks crossing the sink plus recent rolling window;
|
||||
- direct versus paged replay;
|
||||
- CUDA Graph execution;
|
||||
- reset, close, failure cleanup, and exact metadata matching; and
|
||||
- nine-frame VAE decode.
|
||||
|
||||
Official `generate.py` accuracy and performance numbers, together with
|
||||
generation artifacts, are recorded in the PR validation section for the exact
|
||||
tested commit.
|
||||
|
||||
## Current limitations
|
||||
|
||||
- Only the 14B causal-fast checkpoint is supported.
|
||||
- The realtime control plane is internal; there is no public server transport yet.
|
||||
- AR-Diffusion stages currently require one replica because session-affine
|
||||
routing across replicas is not implemented.
|
||||
- One AR block is generated per request and `max_num_seqs` must be one.
|
||||
- Stateful streaming VAE decode is not implemented; the realtime example emits
|
||||
latent chunks.
|
||||
- SP/USP, pipeline/CFG parallelism, HSDP, VAE parallelism, quantization,
|
||||
Cache-DiT, TeaCache, causal-pretrain, and the 1.3B checkpoint are not claimed.
|
||||
- No AMD GPU, Ascend NPU, or Intel GPU support is claimed.
|
||||
|
||||
## References
|
||||
|
||||
- Checkpoint: <https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers>
|
||||
- Official implementation: <https://github.com/robbyant/lingbot-world-v2>
|
||||
- Offline example: [`examples/offline_inference/diffusion/lingbot_world_v2.py`](../../examples/offline_inference/diffusion/lingbot_world_v2.py)
|
||||
- Realtime design: [`docs/design/feature/realtime_ar_diffusion.md`](../../docs/design/feature/realtime_ar_diffusion.md)
|
||||
@@ -178,48 +178,6 @@ with the expected 1536×2720 resolution.
|
||||
- The LLM transformer uses `QKVParallelLinear` and `MergedColumnParallelLinear`
|
||||
for fused QKV and gate/up projections with TP support.
|
||||
|
||||
### 1x AMD MI300X 192GB
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux 6.8.0-134-generic, x86_64
|
||||
- Container: official ROCm image built from `docker/Dockerfile.rocm`
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.11.0+gitd0c8b1f
|
||||
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
|
||||
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
|
||||
- vLLM version: 0.27.0+rocm723
|
||||
- vLLM Omni version or commit: `a704c8759c96e123c0d7c89b11f120b1c0f120cf`
|
||||
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model SenseNova/SenseNova-U1-8B-MoT \
|
||||
--prompt "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light, emotional documentary look." \
|
||||
--width 1536 \
|
||||
--height 2720 \
|
||||
--seed 42 \
|
||||
--num-inference-steps 50 \
|
||||
--cfg-scale 4.0 \
|
||||
--extra-body '{"think": true, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
|
||||
--enable-diffusion-pipeline-profiler \
|
||||
--log-stats \
|
||||
--output sensenova_u1_mi300x.png
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
The command completed and wrote a valid 1536 by 2720 RGB PNG.
|
||||
|
||||
#### Notes
|
||||
|
||||
- Generation with 50 inference steps took 34.011 seconds.
|
||||
- Model loading used 32.774 GiB and took 13.987 seconds.
|
||||
- The internal profiler recorded 35.67 GB reserved and 35.10 GB allocated for the request.
|
||||
- The highest one second whole device memory sample was 37.93 GiB.
|
||||
|
||||
## Online Serving
|
||||
|
||||
SenseNova-U1 supports all four modalities via the OpenAI-compatible
|
||||
|
||||
@@ -78,7 +78,7 @@ python examples/offline_inference/text_to_audio/text_to_audio.py \
|
||||
Start the online serving endpoint:
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/stable-audio-open-1.0 \
|
||||
vllm-omni serve /path/to/stable-audio-open-1.0 \
|
||||
--host 0.0.0.0 \
|
||||
--port 8091 \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
@@ -164,50 +164,7 @@ PY
|
||||
`torchaudio==2.11.0+cu130`.
|
||||
- The `NIXL is not available`, `GLOO_SOCKET_IFNAME`, and `torchsde` boundary
|
||||
warnings observed during validation did not prevent successful generation.
|
||||
- The RTX 4090 entry was validated on one 24 GB GPU. The MI300X entry below covers one 192 GB GPU.
|
||||
- This recipe was validated on one RTX 4090 24GB GPU only. Other GPU counts,
|
||||
ROCm, XPU, and NPU setups are not covered here.
|
||||
- Long generations, higher inference-step counts, and non-WAV response formats
|
||||
were not benchmarked in this recipe.
|
||||
|
||||
#### 1x AMD MI300X 192GB
|
||||
|
||||
##### Environment
|
||||
|
||||
- OS: Linux 6.8.0-134-generic, x86_64
|
||||
- Container: official ROCm image built from `docker/Dockerfile.rocm`
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.11.0+gitd0c8b1f
|
||||
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
|
||||
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
|
||||
- vLLM version: 0.27.0+rocm723
|
||||
- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
|
||||
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
|
||||
|
||||
##### Commands
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_audio/text_to_audio.py \
|
||||
--model stabilityai/stable-audio-open-1.0 \
|
||||
--prompt "A gentle piano melody with soft room ambience" \
|
||||
--negative-prompt "Low quality, distorted, noisy" \
|
||||
--seed 42 \
|
||||
--guidance-scale 7.0 \
|
||||
--audio-length 10.0 \
|
||||
--num-inference-steps 50 \
|
||||
--cache-backend tea_cache \
|
||||
--enable-diffusion-pipeline-profiler \
|
||||
--output stable_audio_10s.wav
|
||||
```
|
||||
|
||||
##### Verification
|
||||
|
||||
The command completed and wrote a valid 44.1 kHz stereo WAV with 10.00 seconds of audio.
|
||||
|
||||
##### Notes
|
||||
|
||||
- TeaCache ran with `rel_l1_thresh=0.2`.
|
||||
- Model loading used 2.7891 GiB and took 3.706 seconds.
|
||||
- Generation took 4.750 seconds, which gives a real time factor of 0.475 for the 10.00 second output.
|
||||
- The internal profiler recorded 15.65 GB reserved and 9.68 GB allocated for the request.
|
||||
- The highest one second whole device memory sample was 19.61 GiB.
|
||||
- The output RMS was 0.0887, and its peak absolute amplitude was 0.5761.
|
||||
- The full process took 384 seconds, including startup and compilation.
|
||||
@@ -39,9 +39,6 @@ platform section, document one or more tested hardware configurations.
|
||||
|
||||
#### Command
|
||||
|
||||
Serve commands use `vllm serve <model> --omni`. The `--omni` flag is what
|
||||
selects the Omni pipeline, so it is required and cannot be dropped.
|
||||
|
||||
```bash
|
||||
# Add the exact command(s) here
|
||||
```
|
||||
|
||||
@@ -39,7 +39,7 @@ Other hardware is welcome as community validation lands.
|
||||
Launch the two-stage talker:
|
||||
|
||||
```bash
|
||||
vllm serve inclusionAI/Ming-omni-tts-0.5B \
|
||||
vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
@@ -170,35 +170,6 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text`
|
||||
for the dialect case.
|
||||
|
||||
**Offline smoke result on ROCm 7.2:**
|
||||
|
||||
A later offline smoke run used the same MI300X model with the following software:
|
||||
|
||||
- GPU: one AMD Instinct MI300X with 191.69 GiB visible HBM and gfx942
|
||||
- Kernel: Linux 6.8.0-134-generic, x86_64
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.11.0+gitd0c8b1f
|
||||
- ROCm or HIP: 7.2.53211
|
||||
- vLLM: 0.27.0+rocm723
|
||||
- vLLM Omni checkout: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--model inclusionAI/Ming-omni-tts-0.5B \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--case basic \
|
||||
--text "你好,这是 AMD MI300X ROCm 配方验证。" \
|
||||
--output-dir ./ming_tts_output \
|
||||
--output-name ming_omni_tts_mi300x.wav \
|
||||
--enforce-eager \
|
||||
--log-stats \
|
||||
--metadata-json ./ming_tts_manifest.json
|
||||
```
|
||||
|
||||
The checked run produced a valid 44.1 kHz mono WAV with 4.80 seconds of audio, RMS 0.1449, and peak absolute amplitude 0.8621. The first request took 20.406 seconds, which gives a real time factor of 4.25. The stage weight loads used 1.31 GiB and 1.47 GiB and took 1.756 seconds and 1.264 seconds. The largest one second whole device memory sample was 91.05 GiB, including the reserved KV cache. Both stages used eager mode, and the AR stage selected `TRITON_ATTN`.
|
||||
|
||||
The run used the official ROCm image built from `docker/Dockerfile.rocm`.
|
||||
|
||||
## Notes
|
||||
|
||||
- The official ROCm image includes the platform dependencies.
|
||||
|
||||
@@ -1,98 +0,0 @@
|
||||
# OmniVoice
|
||||
|
||||
> Offline and online text to speech with automatic voice, voice design, and reference voice cloning
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: k2-fsa
|
||||
- Model: `k2-fsa/OmniVoice`
|
||||
- Task: Multilingual text to speech
|
||||
- Mode: Offline inference for all modes and online serving for automatic voice
|
||||
- Maintainer: Community
|
||||
|
||||
## References
|
||||
|
||||
- [Hugging Face model](https://huggingface.co/k2-fsa/OmniVoice)
|
||||
- [Offline text to speech examples](../../docs/user_guide/examples/offline_inference/text_to_speech.md#omnivoice)
|
||||
- [Online text to speech examples](../../docs/user_guide/examples/online_serving/text_to_speech.md#omnivoice)
|
||||
- [Deploy config](../../vllm_omni/deploy/omnivoice.yaml)
|
||||
|
||||
## Requirements
|
||||
|
||||
Download the checkpoint before the first run:
|
||||
|
||||
```bash
|
||||
huggingface-cli download k2-fsa/OmniVoice
|
||||
```
|
||||
|
||||
Automatic voice and voice design require `transformers>=4.57.0`. Reference voice cloning requires `transformers>=5.3.0`.
|
||||
|
||||
The vLLM Omni deploy config runs the OmniVoice generator and decoder inside one float32 diffusion stage.
|
||||
|
||||
## ROCm
|
||||
|
||||
### 1x AMD MI300X
|
||||
|
||||
#### Environment
|
||||
|
||||
- GPU: one AMD Instinct MI300X with 191.69 GiB visible HBM and gfx942
|
||||
- Kernel: Linux 6.8.0-134-generic, x86_64
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.11.0+gitd0c8b1f
|
||||
- ROCm or HIP: 7.2.53211
|
||||
- vLLM: 0.27.0+rocm723
|
||||
- vLLM Omni checkout: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
|
||||
- transformers: 5.15.0
|
||||
|
||||
The run used the official ROCm image built from `docker/Dockerfile.rocm`.
|
||||
|
||||
#### Automatic voice
|
||||
|
||||
Run the offline example from the repository root:
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
--deploy-config vllm_omni/deploy/omnivoice.yaml \
|
||||
--text "Hello, this is OmniVoice running on one AMD MI300X." \
|
||||
--seed 42 \
|
||||
--output omnivoice_mi300x.wav
|
||||
```
|
||||
|
||||
The checked run produced a valid 24 kHz mono WAV with 3.60 seconds of audio, RMS 0.1194, and peak absolute amplitude 0.9129. The first request took 0.685 seconds, which gives a real time factor of 0.19. Model loading used 3.5739 GiB and took 9.564 seconds. The request used 4.56 GB reserved and 3.64 GB allocated according to the internal profiler, while the largest one second whole device sample was 8.33 GiB. The log states that the OmniVoice Triton kernels loaded.
|
||||
|
||||
#### Reference voice cloning
|
||||
|
||||
The offline interface also accepts a reference WAV and its transcript:
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
--text "Hello, this is a cloned voice." \
|
||||
--ref-audio ref.wav \
|
||||
--ref-text "This is the reference transcription." \
|
||||
--output cloned.wav
|
||||
```
|
||||
|
||||
#### Voice design
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
--text "Hello, this is a designed voice." \
|
||||
--instruct "female, low pitch, British accent" \
|
||||
--output designed.wav
|
||||
```
|
||||
|
||||
## Online serving
|
||||
|
||||
Online serving currently supports automatic voice:
|
||||
|
||||
```bash
|
||||
vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
|
||||
```
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/text_to_speech/omnivoice
|
||||
python speech_client.py --text "Hello, how are you?"
|
||||
```
|
||||
@@ -1,108 +0,0 @@
|
||||
# Pi0 (π0, Pi-Zero) VLA
|
||||
|
||||
> Vision-Language-Action policy serving over the OpenPI realtime websocket
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Physical Intelligence (checkpoint redistributed by LeRobot)
|
||||
- Model: `lerobot/pi0_base`
|
||||
- Task: Vision-Language-Action — multi-camera images + a language instruction +
|
||||
robot proprioceptive state → a continuous action chunk
|
||||
`[action_horizon, action_dim] = [50, 32]` via flow-matching denoising. π0 does
|
||||
**not** emit text tokens.
|
||||
- Mode: Online serving via the OpenPI realtime robot API
|
||||
(`/v1/realtime/robot/openpi` websocket)
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe to deploy `lerobot/pi0_base` as a robot policy server: a robot
|
||||
client connects over the OpenPI websocket and, per control step, sends an
|
||||
observation (camera frames + state + instruction) and receives an action chunk.
|
||||
|
||||
A single pipeline class (`Pi0Pipeline`) serves the policy. π0 is stateless across
|
||||
calls (no KV reuse), so the protocol's `session_id` / `reset` are accepted but
|
||||
ignored.
|
||||
|
||||
## Weights
|
||||
|
||||
π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the
|
||||
server download it (`MODEL=lerobot/pi0_base`) or point at a local copy.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
### 1x H200 141GB (Online serving)
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu 22.04+
|
||||
- Python: 3.12+
|
||||
- Driver / runtime: NVIDIA CUDA environment
|
||||
- vLLM / vLLM-Omni: use the commit you are deploying from
|
||||
- OpenPI client extras (for the example client): `openpi-client`
|
||||
(`pip install -e packages/openpi-client` from the
|
||||
[openpi](https://github.com/Physical-Intelligence/openpi) repo), `websockets`,
|
||||
`msgpack`, `msgpack-numpy`
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
vllm serve lerobot/pi0_base \
|
||||
--omni \
|
||||
--host 0.0.0.0 --port 8000 \
|
||||
--served-model-name pi0 \
|
||||
--deploy-config vllm_omni/deploy/pi0.yaml \
|
||||
--enforce-eager --disable-log-stats
|
||||
```
|
||||
|
||||
The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
|
||||
(`Pi0Pipeline`), `float32`, `max_num_seqs: 1`, and the `policy_server_config`
|
||||
handshake metadata (3 cameras, `joint_position`, action horizon 50, action dim
|
||||
32) that enables the OpenPI realtime serving layer.
|
||||
|
||||
#### Verification
|
||||
|
||||
Run the example client (sends a synthetic observation and prints the returned
|
||||
action chunk):
|
||||
|
||||
```bash
|
||||
python examples/online_serving/pi0/openpi_client.py \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--prompt "pick up the red block and place it in the bin"
|
||||
```
|
||||
|
||||
It connects, prints the server metadata, sends robot observations, and prints the
|
||||
returned `[action_horizon, action_dim] = [50, 32]` action chunks. The observation
|
||||
is a flat dict per inference:
|
||||
|
||||
```python
|
||||
{
|
||||
"observation.images.base_0_rgb": np.uint8[H, W, 3],
|
||||
"observation.images.left_wrist_0_rgb": np.uint8[H, W, 3],
|
||||
"observation.images.right_wrist_0_rgb": np.uint8[H, W, 3],
|
||||
"state": np.float32[state_dim], # zero-padded to max_state_dim=32 server-side
|
||||
"prompt": "pick up the red block",
|
||||
"session_id": "<uuid>", # accepted but ignored (π0 is stateless)
|
||||
}
|
||||
```
|
||||
|
||||
Camera keys must match the server's `image_feature_keys` (the checkpoint's
|
||||
`input_features` order). If your robot uses different camera names, set
|
||||
`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
|
||||
|
||||
## Correctness
|
||||
|
||||
π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy`
|
||||
reference (`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
|
||||
`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`).
|
||||
|
||||
## References
|
||||
|
||||
- Model blog: <https://www.physicalintelligence.company/blog/pi0>
|
||||
- Checkpoint: <https://huggingface.co/lerobot/pi0_base>
|
||||
- Pipeline: [`vllm_omni/diffusion/models/pi0/pipeline_pi0.py`](../../vllm_omni/diffusion/models/pi0/pipeline_pi0.py)
|
||||
- Deploy config: [`vllm_omni/deploy/pi0.yaml`](../../vllm_omni/deploy/pi0.yaml)
|
||||
- Example client: [`examples/online_serving/pi0/`](../../examples/online_serving/pi0/)
|
||||
- Tests: CPU units + LeRobot parity in [`tests/diffusion/models/pi0/`](../../tests/diffusion/models/pi0/),
|
||||
OpenPI websocket e2e in
|
||||
[`tests/e2e/online_serving/test_pi0_expansion.py`](../../tests/e2e/online_serving/test_pi0_expansion.py)
|
||||
@@ -44,20 +44,20 @@ conditioned on reference audio and its transcript.
|
||||
Start the server from the repository root:
|
||||
|
||||
```bash
|
||||
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
|
||||
vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
|
||||
```
|
||||
|
||||
Async chunking is enabled by default in the bundled deployment config. For
|
||||
the sync (non-streaming) path:
|
||||
|
||||
```bash
|
||||
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
|
||||
vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
|
||||
```
|
||||
|
||||
Use a custom deploy config for advanced cases:
|
||||
|
||||
```bash
|
||||
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
|
||||
vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
|
||||
--deploy-config /path/to/your_glm_tts_overrides.yaml
|
||||
```
|
||||
|
||||
|
||||
@@ -1102,13 +1102,6 @@ class TestResolveScheduler:
|
||||
|
||||
|
||||
class TestDeployConfigLoading:
|
||||
def test_rejects_legacy_stage_args_schema(self, tmp_path):
|
||||
deploy_path = tmp_path / "legacy.yaml"
|
||||
deploy_path.write_text("stage_args:\n - stage_id: 0\n", encoding="utf-8")
|
||||
|
||||
with pytest.raises(ValueError, match=r"stage_args.*PipelineConfig.*stages"):
|
||||
load_deploy_config(deploy_path)
|
||||
|
||||
def test_load_minicpmo_duplex_deploy_config(self):
|
||||
deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml"))
|
||||
|
||||
@@ -1384,8 +1377,7 @@ stages:
|
||||
assert async_stages[1].custom_process_input_func is None
|
||||
|
||||
def test_no_bundled_legacy_stage_config_yamls(self):
|
||||
repo_root = Path(__file__).resolve().parents[2]
|
||||
stage_config_dir = repo_root / "vllm_omni" / "model_executor" / "stage_configs"
|
||||
stage_config_dir = Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "stage_configs"
|
||||
assert not list(stage_config_dir.glob("*.yaml"))
|
||||
|
||||
def test_merge_pipeline_deploy(self):
|
||||
|
||||
@@ -4,10 +4,7 @@ from types import SimpleNamespace
|
||||
import pytest
|
||||
from vllm.v1.engine import FinishReason
|
||||
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import (
|
||||
DEFAULT_INPUT_WAIT_TIMEOUT_S,
|
||||
OmniSchedulerMixin,
|
||||
)
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
||||
from vllm_omni.core.sched.output import OmniChunkRecvHandle
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -25,25 +22,13 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
|
||||
scheduler.requests = {"request": object()}
|
||||
scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
|
||||
scheduler._process_pending_input_timeouts = lambda: calls.append(("timeouts",))
|
||||
|
||||
def _collect_timed_out(timeout_s):
|
||||
calls.append(("chunk-timeouts", timeout_s))
|
||||
return set()
|
||||
|
||||
def _collect_failed_sends():
|
||||
calls.append(("failed-sends",))
|
||||
return {}
|
||||
|
||||
scheduler.chunk_transfer_adapter = SimpleNamespace(
|
||||
receives_chunks=True,
|
||||
process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
|
||||
("process", waiting, running, scheduler_requests)
|
||||
),
|
||||
restore_queues=lambda waiting, running, scheduler_requests: calls.append(
|
||||
("restore-chunks", waiting, running, scheduler_requests)
|
||||
),
|
||||
collect_timed_out_request_ids=_collect_timed_out,
|
||||
collect_failed_send_request_ids=_collect_failed_sends,
|
||||
)
|
||||
scheduler.input_coordinator = SimpleNamespace(
|
||||
restore_queues=lambda waiting: calls.append(("restore-full", waiting))
|
||||
@@ -56,10 +41,6 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
|
||||
("consume", "ar"),
|
||||
("timeouts",),
|
||||
("process", scheduler.waiting, scheduler.running, scheduler.requests),
|
||||
# The chunk deadline runs after chunks are applied, so a chunk that
|
||||
# arrived this cycle resets the clock before it is measured (R1.1).
|
||||
("chunk-timeouts", DEFAULT_INPUT_WAIT_TIMEOUT_S),
|
||||
("failed-sends",),
|
||||
("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
|
||||
("restore-full", scheduler.waiting),
|
||||
]
|
||||
|
||||
@@ -1,10 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit coverage for the two stage-input timeout hooks.
|
||||
|
||||
``_process_pending_input_timeouts`` covers the full-payload transport;
|
||||
``_process_pending_chunk_timeouts`` covers async-chunk. Both share
|
||||
``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S``.
|
||||
"""Unit coverage for _process_pending_input_timeouts.
|
||||
|
||||
Verifies that the mixin correctly *delegates* timed-out requests to the
|
||||
base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``.
|
||||
@@ -16,10 +12,6 @@ this file only asserts the wiring from the mixin to that API.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
@@ -110,248 +102,5 @@ def test_process_pending_input_timeouts_disabled_when_timeout_zero(monkeypatch):
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
class _FakeChunkAdapter:
|
||||
"""Stands in for OmniChunkTransferAdapter's timeout-collection surface."""
|
||||
|
||||
def __init__(self, timed_out_ids, *, receives_chunks=True):
|
||||
self._timed_out_ids = set(timed_out_ids)
|
||||
self.receives_chunks = receives_chunks
|
||||
self.calls = []
|
||||
|
||||
def collect_timed_out_request_ids(self, timeout_s):
|
||||
self.calls.append(timeout_s)
|
||||
return set(self._timed_out_ids)
|
||||
|
||||
|
||||
class _FakeChunkScheduler(OmniSchedulerMixin):
|
||||
def __init__(self, requests, adapter):
|
||||
self.requests = requests
|
||||
self.chunk_transfer_adapter = adapter
|
||||
self.finish_calls = []
|
||||
|
||||
def finish_requests(self, req_ids, status):
|
||||
self.finish_calls.append((set(req_ids), status))
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_delegates_to_finish_requests():
|
||||
"""A request stalled in WAITING_FOR_CHUNK is failed, not left parked.
|
||||
|
||||
This is the async-chunk half of the safety net (RFC #4855 R1.1); before it
|
||||
existed, the chunk path had no deadline at all and issue #3833 could park a
|
||||
request forever.
|
||||
"""
|
||||
req_id = "stalled-stream"
|
||||
adapter = _FakeChunkAdapter(timed_out_ids={req_id})
|
||||
scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id)}, adapter)
|
||||
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
|
||||
assert len(adapter.calls) == 1
|
||||
assert adapter.calls[0] > 0
|
||||
assert len(scheduler.finish_calls) == 1
|
||||
finished_ids, status = scheduler.finish_calls[0]
|
||||
assert finished_ids == {req_id}
|
||||
assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_skips_already_freed_request():
|
||||
adapter = _FakeChunkAdapter(timed_out_ids={"already-freed"})
|
||||
scheduler = _FakeChunkScheduler(requests={}, adapter=adapter)
|
||||
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_noop_on_a_producer_stage():
|
||||
"""Stage 0 does not receive chunks, so it has nothing to expire."""
|
||||
adapter = _FakeChunkAdapter(timed_out_ids={"r1"}, receives_chunks=False)
|
||||
scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter)
|
||||
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
|
||||
assert adapter.calls == []
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_noop_without_adapter():
|
||||
class _NoAdapter(OmniSchedulerMixin):
|
||||
def __init__(self):
|
||||
self.requests = {}
|
||||
self.chunk_transfer_adapter = None
|
||||
self.finish_calls = []
|
||||
|
||||
def finish_requests(self, req_ids, status):
|
||||
self.finish_calls.append((set(req_ids), status))
|
||||
|
||||
scheduler = _NoAdapter()
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_disabled_when_timeout_zero(monkeypatch):
|
||||
"""The shared knob disables both transports' nets together."""
|
||||
from vllm_omni.core.sched import omni_scheduler_mixin
|
||||
|
||||
monkeypatch.setattr(omni_scheduler_mixin, "DEFAULT_INPUT_WAIT_TIMEOUT_S", 0.0)
|
||||
|
||||
adapter = _FakeChunkAdapter(timed_out_ids={"r1"})
|
||||
scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter)
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
|
||||
assert adapter.calls == []
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
class _FakeSendFailureAdapter:
|
||||
"""Stands in for the adapter's drained send-failure map (R1.2)."""
|
||||
|
||||
def __init__(self, failures):
|
||||
self._failures = dict(failures)
|
||||
self.collect_calls = 0
|
||||
|
||||
def collect_failed_send_request_ids(self):
|
||||
self.collect_calls += 1
|
||||
failures, self._failures = self._failures, {}
|
||||
return failures
|
||||
|
||||
|
||||
class _FakeSendFailureScheduler(OmniSchedulerMixin):
|
||||
def __init__(self, requests, adapter):
|
||||
self.requests = requests
|
||||
self.chunk_transfer_adapter = adapter
|
||||
|
||||
|
||||
def test_log_failed_chunk_sends_reports_a_live_request(caplog):
|
||||
adapter = _FakeSendFailureAdapter({"req-live": "OSError: No space left on device"})
|
||||
scheduler = _FakeSendFailureScheduler({"req-live": SimpleNamespace(request_id="req-live")}, adapter)
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
scheduler._log_failed_chunk_sends()
|
||||
|
||||
assert "req-live" in caplog.text
|
||||
assert "No space left on device" in caplog.text
|
||||
|
||||
|
||||
def test_log_failed_chunk_sends_still_reports_a_departed_request(caplog):
|
||||
"""``collect_*`` drains the map, so a failure for a request that already left
|
||||
this scheduler is the one place a drop could still go unrecorded."""
|
||||
adapter = _FakeSendFailureAdapter({"req-gone": "OSError: No space left on device"})
|
||||
scheduler = _FakeSendFailureScheduler({}, adapter)
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
scheduler._log_failed_chunk_sends()
|
||||
|
||||
assert adapter.collect_calls == 1
|
||||
assert "req-gone" in caplog.text
|
||||
assert "No space left on device" in caplog.text
|
||||
assert "finished or aborted" in caplog.text
|
||||
|
||||
|
||||
def test_log_failed_chunk_sends_noop_without_adapter():
|
||||
scheduler = _FakeSendFailureScheduler({}, None)
|
||||
scheduler._log_failed_chunk_sends()
|
||||
|
||||
|
||||
# --- R1.4: VLLM_OMNI_INPUT_WAIT_TIMEOUT_S is parsed at import time, so each of
|
||||
# these reloads the module under a different environment. One knob now gates the
|
||||
# deadline on both transports, which is why a silent 0 matters more than it did
|
||||
# when only the full-payload path used it.
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def reload_with(monkeypatch):
|
||||
"""Reload the mixin module under a given env value, then put it back.
|
||||
|
||||
``importlib.reload`` re-executes into the live module dict, so whatever the
|
||||
last test parsed is what every later test in the session sees -- ``-k
|
||||
zero_warns`` alone would leave ``DEFAULT_INPUT_WAIT_TIMEOUT_S = 0.0`` and
|
||||
silently disable both deadlines. The rejection cases are worse: the
|
||||
ValueError fires *after* the rebind, so ``nan`` lingers in the module.
|
||||
``monkeypatch`` restores the environment but not the module, so reload once
|
||||
more on the way out.
|
||||
"""
|
||||
import vllm_omni.core.sched.omni_scheduler_mixin as mod
|
||||
|
||||
def _reload(raw):
|
||||
if raw is None:
|
||||
monkeypatch.delenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raw)
|
||||
return importlib.reload(mod)
|
||||
|
||||
yield _reload
|
||||
|
||||
monkeypatch.undo()
|
||||
importlib.reload(mod)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"raw,expected",
|
||||
[
|
||||
(None, 600.0),
|
||||
("30", 30.0),
|
||||
("0.5", 0.5),
|
||||
("0", 0.0), # explicit opt-out is still honoured
|
||||
],
|
||||
)
|
||||
def test_timeout_accepts_valid_values(reload_with, raw, expected):
|
||||
mod = reload_with(raw)
|
||||
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["-1", "-600", "-0.5"])
|
||||
def test_negative_timeout_is_rejected(reload_with, raw):
|
||||
"""`-1` is a common "no limit" idiom; substituting 600 would give an operator
|
||||
who meant "never time out" mysterious failures ten minutes in. Fail loudly."""
|
||||
with pytest.raises(ValueError, match="negative"):
|
||||
reload_with(raw)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["nan", "inf", "-inf", "1e999"])
|
||||
def test_non_finite_timeout_is_rejected(reload_with, raw):
|
||||
"""nan compares false against every deadline check and inf is unreachable;
|
||||
both disable the net while looking like a number."""
|
||||
with pytest.raises(ValueError, match="not finite"):
|
||||
reload_with(raw)
|
||||
|
||||
|
||||
def test_zero_warns_and_names_both_transports(reload_with, caplog):
|
||||
with caplog.at_level(logging.WARNING):
|
||||
mod = reload_with("0")
|
||||
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 0.0
|
||||
text = caplog.text
|
||||
assert "DISABLED" in text
|
||||
assert "WAITING_FOR_INPUT" in text and "WAITING_FOR_CHUNK" in text
|
||||
|
||||
|
||||
def test_unparsable_timeout_falls_back(reload_with):
|
||||
mod = reload_with("not-a-number")
|
||||
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 600.0
|
||||
|
||||
|
||||
def test_a_rejected_value_does_not_outlive_its_test(monkeypatch):
|
||||
"""The reason ``reload_with`` reloads on the way out, done here by hand.
|
||||
|
||||
``nan`` is bound before the ValueError fires, so the reload that rejects it
|
||||
still leaves it in the module -- and nan compares false against every
|
||||
deadline check, which is exactly the silent-disable this item exists to
|
||||
stop. Restoring the environment is not enough; the module has to be reloaded
|
||||
too.
|
||||
"""
|
||||
import vllm_omni.core.sched.omni_scheduler_mixin as mod
|
||||
|
||||
ambient = float(os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600"))
|
||||
|
||||
monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "nan")
|
||||
with pytest.raises(ValueError, match="not finite"):
|
||||
importlib.reload(mod)
|
||||
assert math.isnan(mod.DEFAULT_INPUT_WAIT_TIMEOUT_S)
|
||||
|
||||
monkeypatch.undo()
|
||||
importlib.reload(mod)
|
||||
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == ambient
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task",
|
||||
@@ -61,7 +61,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task",
|
||||
@@ -113,7 +113,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task",
|
||||
@@ -164,7 +164,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task",
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"tts"
|
||||
],
|
||||
"server_params": {
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline with enforce-eager (no torch.compile)",
|
||||
@@ -64,7 +64,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device with torch.compile (default, no enforce-eager)",
|
||||
@@ -117,7 +117,7 @@
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "CFG-parallel=2 with enforce-eager",
|
||||
@@ -172,7 +172,7 @@
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "CFG-parallel=2 with torch.compile",
|
||||
@@ -226,7 +226,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "CacheDiT with enforce-eager",
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline (two input images)",
|
||||
@@ -76,7 +76,7 @@
|
||||
"num_cards": 4
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
|
||||
@@ -125,7 +125,7 @@
|
||||
"num_cards": 4
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Ulysses SP=2 + CFG=2 + CacheDiT",
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline",
|
||||
@@ -73,7 +73,7 @@
|
||||
"num_cards": 4
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
|
||||
@@ -121,7 +121,7 @@
|
||||
"num_cards": 4
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Ulysses SP=2 + CFG=2 + CacheDiT",
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline",
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"tts"
|
||||
],
|
||||
"server_params": {
|
||||
|
||||
@@ -18,7 +18,6 @@ from vllm_omni.experimental.ar_diffusion.capability import (
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVConfig
|
||||
from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest
|
||||
|
||||
BLOCK = 16
|
||||
POS = "positive"
|
||||
@@ -57,21 +56,6 @@ def dreamzero_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
|
||||
)
|
||||
|
||||
|
||||
def tiny_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
|
||||
return ARDiffusionKVCacheSpec(
|
||||
num_layers=1,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
tokens_per_frame=1,
|
||||
frames_per_block=3,
|
||||
window_frames=3,
|
||||
sink_frames=3,
|
||||
kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
|
||||
session_capacity=capacity,
|
||||
cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 2),),
|
||||
)
|
||||
|
||||
|
||||
class CapablePipeline:
|
||||
def __init__(self, spec: ARDiffusionKVCacheSpec) -> None:
|
||||
self.spec = spec
|
||||
@@ -119,7 +103,6 @@ def make_runner(
|
||||
*,
|
||||
available_bytes: int = 1 << 28,
|
||||
step_execution: bool = False,
|
||||
gpu_memory_fraction: float = 0.1,
|
||||
) -> ARDiffusionModelRunner:
|
||||
runner = object.__new__(ARDiffusionModelRunner)
|
||||
runner.od_config = SimpleNamespace(
|
||||
@@ -130,10 +113,7 @@ def make_runner(
|
||||
)
|
||||
runner.device = torch.device("cpu")
|
||||
runner.pipeline = pipeline
|
||||
runner.ar_diffusion_kv_config = ARDiffusionKVConfig(
|
||||
enable=True,
|
||||
gpu_memory_fraction=gpu_memory_fraction,
|
||||
)
|
||||
runner.ar_diffusion_kv_config = ARDiffusionKVConfig(enable=True)
|
||||
runner.kv_cache = None
|
||||
runner._ar_diffusion_capability = None
|
||||
runner._ar_diffusion_kv_cache_spec = None
|
||||
@@ -152,16 +132,6 @@ def commit_one_frame(runner: ARDiffusionModelRunner, session_id: str, kv_branch:
|
||||
return state
|
||||
|
||||
|
||||
def commit_one_block(runner: ARDiffusionModelRunner, session_id: str, kv_branch: str):
|
||||
state = runner._get_or_create_session(session_id)
|
||||
assert runner.kv_cache is not None
|
||||
seq_len = BLOCK * runner.kv_cache.frames_per_block
|
||||
ctx = state.get_kv_caches(kv_branch, seq_len=seq_len, commit_current=True)[0].forward_ctx
|
||||
ctx.ensure_video_slots(torch.device("cpu"))
|
||||
state.commit_paged_context(kv_branch)
|
||||
return state
|
||||
|
||||
|
||||
def test_ar_runner_rejects_pipeline_without_capability():
|
||||
runner = object.__new__(ARDiffusionModelRunner)
|
||||
runner.od_config = SimpleNamespace(max_num_seqs=1)
|
||||
@@ -170,43 +140,6 @@ def test_ar_runner_rejects_pipeline_without_capability():
|
||||
runner._preallocate_kv_cache(available_bytes=1 << 20)
|
||||
|
||||
|
||||
def test_runner_uses_typed_tick_as_authoritative_session_contract():
|
||||
tick = ARDiffusionTickRequest(
|
||||
session_id="world-7",
|
||||
request_id="request-3",
|
||||
chunk_index=3,
|
||||
reset=True,
|
||||
)
|
||||
req = SimpleNamespace(
|
||||
request_id="request-3",
|
||||
sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
|
||||
)
|
||||
|
||||
session_id, extra_args, parsed = ARDiffusionModelRunner._request_session(req)
|
||||
|
||||
assert session_id == "world-7"
|
||||
assert extra_args == tick.to_extra_args()
|
||||
assert parsed == tick
|
||||
|
||||
|
||||
def test_runner_keeps_engine_request_id_separate_from_tick_correlation_id():
|
||||
tick = ARDiffusionTickRequest(
|
||||
session_id="world-7",
|
||||
request_id="client-request-3",
|
||||
chunk_index=3,
|
||||
)
|
||||
req = SimpleNamespace(
|
||||
request_id="engine-request-uuid",
|
||||
sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
|
||||
)
|
||||
|
||||
session_id, _, parsed = ARDiffusionModelRunner._request_session(req)
|
||||
|
||||
assert session_id == "world-7"
|
||||
assert req.request_id == "engine-request-uuid"
|
||||
assert parsed.request_id == "client-request-3"
|
||||
|
||||
|
||||
def test_lingbot_like_single_branch_session_reuse_reset_and_close():
|
||||
pipeline = CapablePipeline(lingbot_like_spec())
|
||||
runner = make_runner(pipeline)
|
||||
@@ -237,31 +170,6 @@ def test_lingbot_like_single_branch_session_reuse_reset_and_close():
|
||||
assert pipeline.closes == ["s1"]
|
||||
|
||||
|
||||
def test_lingbot_like_interleaved_sessions_keep_independent_kv_partitions():
|
||||
runner = make_runner(CapablePipeline(lingbot_like_spec(capacity=2)))
|
||||
kv = runner.kv_cache
|
||||
assert kv is not None
|
||||
|
||||
session_a = commit_one_block(runner, "world-a", "main")
|
||||
a0_blocks = kv.window_block_ids(session_a.adapter("main"))
|
||||
session_b = commit_one_block(runner, "world-b", "main")
|
||||
b0_blocks = kv.window_block_ids(session_b.adapter("main"))
|
||||
session_a_again = commit_one_block(runner, "world-a", "main")
|
||||
a1_blocks = kv.window_block_ids(session_a_again.adapter("main"))
|
||||
|
||||
assert session_a_again is session_a
|
||||
assert session_a.adapter("main").request_id == "ar::world-a::main"
|
||||
assert session_b.adapter("main").request_id == "ar::world-b::main"
|
||||
assert session_a.adapter("main").completed_chunks == 6
|
||||
assert session_b.adapter("main").completed_chunks == 3
|
||||
assert len(a0_blocks) == 3
|
||||
assert len(b0_blocks) == 3
|
||||
assert len(a1_blocks) == 6
|
||||
assert set(a0_blocks) <= set(a1_blocks)
|
||||
assert set(a1_blocks).isdisjoint(b0_blocks)
|
||||
assert tuple(runner._sessions) == ("world-b", "world-a")
|
||||
|
||||
|
||||
def test_lingbot_like_sink_survives_sliding_window_eviction():
|
||||
runner = make_runner(CapablePipeline(lingbot_like_spec()))
|
||||
kv = runner.kv_cache
|
||||
@@ -290,67 +198,24 @@ def test_dreamzero_like_two_branches_are_independent():
|
||||
|
||||
|
||||
def test_lru_eviction_releases_blocks_and_notifies_pipeline():
|
||||
pipeline = CapablePipeline(lingbot_like_spec(capacity=2))
|
||||
pipeline = CapablePipeline(lingbot_like_spec(capacity=8))
|
||||
runner = make_runner(pipeline)
|
||||
kv = runner.kv_cache
|
||||
assert kv is not None
|
||||
assert runner._session_capacity == 2
|
||||
assert kv.session_capacity == 2
|
||||
assert runner._session_capacity == 1
|
||||
assert kv.session_capacity == 1
|
||||
free_total = kv.manager.block_pool.get_num_free_blocks()
|
||||
old = commit_one_frame(runner, "old", "main")
|
||||
k = torch.randn(1, 8, 4, 64)
|
||||
old.populate_cross_attention("main", "text", [(k, k)] * old.num_layers)
|
||||
assert kv.manager.block_pool.get_num_free_blocks() < free_total
|
||||
|
||||
runner._get_or_create_session("new")
|
||||
assert tuple(runner._sessions) == ("old", "new")
|
||||
runner._get_or_create_session("newest")
|
||||
|
||||
assert tuple(runner._sessions) == ("new", "newest")
|
||||
assert pipeline.closes == ["old"]
|
||||
assert "old" not in kv._cross_sessions
|
||||
assert kv.manager.block_pool.get_num_free_blocks() == free_total
|
||||
|
||||
|
||||
def test_budget_reduced_capacity_drives_runner_lru():
|
||||
pipeline = CapablePipeline(tiny_spec(capacity=2))
|
||||
# One tiny LingBot-like session requires 128 bytes; two require 192.
|
||||
runner = make_runner(
|
||||
pipeline,
|
||||
available_bytes=128,
|
||||
gpu_memory_fraction=1.0,
|
||||
)
|
||||
kv = runner.kv_cache
|
||||
assert kv is not None
|
||||
assert kv.requested_session_capacity == 2
|
||||
assert kv.session_capacity == 1
|
||||
assert runner._session_capacity == 1
|
||||
|
||||
runner._get_or_create_session("old")
|
||||
runner._get_or_create_session("new")
|
||||
|
||||
assert tuple(runner._sessions) == ("new",)
|
||||
assert pipeline.closes == ["old"]
|
||||
|
||||
|
||||
def test_dreamzero_like_requested_capacity_is_capped_by_budget():
|
||||
spec = dreamzero_like_spec(capacity=64)
|
||||
# Per all-layer self-KV page: 65,536 bytes. Two resident sessions need:
|
||||
# managed=(2 * (2 * 6 + 4) + 2)=34 pages, scratch=8 pages,
|
||||
# cross-attention=1 page/session, for 44 pages total.
|
||||
page_bytes = 65_536
|
||||
pipeline = CapablePipeline(spec)
|
||||
runner = make_runner(
|
||||
pipeline,
|
||||
available_bytes=44 * page_bytes,
|
||||
gpu_memory_fraction=1.0,
|
||||
)
|
||||
kv = runner.kv_cache
|
||||
assert kv is not None
|
||||
assert kv.requested_session_capacity == 64
|
||||
assert kv.session_capacity == 2
|
||||
assert runner._session_capacity == 2
|
||||
assert kv.cross_attention_reserved_bytes == 2 * page_bytes
|
||||
assert "old" not in kv._cross_sessions
|
||||
assert kv.manager.block_pool.get_num_free_blocks() == free_total
|
||||
|
||||
|
||||
def test_forward_exception_releases_pending_allocation_and_model_state(monkeypatch):
|
||||
@@ -367,10 +232,7 @@ def test_forward_exception_releases_pending_allocation_and_model_state(monkeypat
|
||||
raise RuntimeError("layer exploded")
|
||||
|
||||
monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
|
||||
request = SimpleNamespace(
|
||||
request_id="broken-request",
|
||||
sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
|
||||
)
|
||||
request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
|
||||
|
||||
with pytest.raises(RuntimeError, match="layer exploded"):
|
||||
runner.execute_model(request)
|
||||
@@ -401,10 +263,7 @@ def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
|
||||
monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
|
||||
monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom)
|
||||
runner.device = torch.device("cuda")
|
||||
request = SimpleNamespace(
|
||||
request_id="broken-request",
|
||||
sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
|
||||
)
|
||||
request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
|
||||
|
||||
with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
|
||||
runner.execute_model(request)
|
||||
|
||||
@@ -1,231 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Sequence
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
|
||||
from vllm_omni.experimental.ar_diffusion.session import (
|
||||
ARDiffusionSession,
|
||||
ARDiffusionSessionEvent,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
|
||||
ARDiffusionChunkMetadata,
|
||||
ARDiffusionControlInput,
|
||||
ARDiffusionTickRequest,
|
||||
)
|
||||
from vllm_omni.inputs.data import (
|
||||
OmniDiffusionSamplingParams,
|
||||
OmniPromptType,
|
||||
OmniSamplingParams,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeStagePool:
|
||||
def __init__(self, num_replicas: int = 1) -> None:
|
||||
self.num_replicas = num_replicas
|
||||
|
||||
|
||||
class FakeEngine:
|
||||
def __init__(self, num_replicas: int = 1) -> None:
|
||||
self.stage_pools = [FakeStagePool(num_replicas)]
|
||||
|
||||
|
||||
class FakeGenerateClient:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
include_metadata: bool = True,
|
||||
applied_event_ids: list[int] | None = None,
|
||||
output_error: str | None = None,
|
||||
num_replicas: int = 1,
|
||||
) -> None:
|
||||
self.engine = FakeEngine(num_replicas)
|
||||
self.include_metadata = include_metadata
|
||||
self.applied_event_ids = applied_event_ids
|
||||
self.output_error = output_error
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
prompt: OmniPromptType,
|
||||
*,
|
||||
request_id: str,
|
||||
sampling_params_list: Sequence[OmniSamplingParams],
|
||||
output_modalities: list[str] | None = None,
|
||||
) -> AsyncIterator[OmniRequestOutput]:
|
||||
self.calls.append(
|
||||
{
|
||||
"prompt": prompt,
|
||||
"request_id": request_id,
|
||||
"sampling_params_list": sampling_params_list,
|
||||
"output_modalities": output_modalities,
|
||||
}
|
||||
)
|
||||
params = sampling_params_list[0]
|
||||
assert isinstance(params, OmniDiffusionSamplingParams)
|
||||
tick = ARDiffusionTickRequest.from_extra_args(params.extra_args)
|
||||
assert tick is not None
|
||||
chunk_metadata = ARDiffusionChunkMetadata.from_tick(tick).to_dict()
|
||||
if self.applied_event_ids is not None:
|
||||
chunk_metadata["applied_event_ids"] = self.applied_event_ids
|
||||
metadata = {"metadata": {"ar_diffusion": chunk_metadata}} if self.include_metadata else {}
|
||||
yield OmniRequestOutput(
|
||||
request_id=request_id,
|
||||
stage_id=0,
|
||||
final_output_type="video",
|
||||
images=["chunk"], # type: ignore[list-item]
|
||||
_multimodal_output=metadata,
|
||||
finished=True,
|
||||
error=self.output_error,
|
||||
)
|
||||
|
||||
|
||||
class FakeLifecycle:
|
||||
async def reset_session(self, session_id: str) -> None:
|
||||
pass
|
||||
|
||||
async def close_session(self, session_id: str) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def make_tick() -> ARDiffusionTickRequest:
|
||||
return ARDiffusionTickRequest(
|
||||
session_id="world-1",
|
||||
request_id="world-1-chunk-3",
|
||||
chunk_index=3,
|
||||
applied_event_ids=(7, 8),
|
||||
prompt="turn left",
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data={"samples": [{"x": 1.0}]},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_consumer_rejects_replicated_ar_diffusion_stage() -> None:
|
||||
client = FakeGenerateClient(num_replicas=2)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"require exactly one replica.*num_replicas=2",
|
||||
):
|
||||
ARDiffusionOmniTickConsumer(
|
||||
client,
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
|
||||
assert client.calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consumer_submits_typed_tick_and_returns_standard_output() -> None:
|
||||
client = FakeGenerateClient()
|
||||
template = OmniDiffusionSamplingParams(
|
||||
output_type="latent",
|
||||
extra_args={"deployment_option": "kept"},
|
||||
)
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
client,
|
||||
prompt_provider=lambda tick: {
|
||||
"prompt": tick.prompt,
|
||||
"multi_modal_data": {"image": "initial-image"},
|
||||
},
|
||||
sampling_params_list=[template],
|
||||
diffusion_stage_id=0,
|
||||
output_modalities=["video"],
|
||||
)
|
||||
tick = make_tick()
|
||||
|
||||
output = await consumer.execute_tick(tick)
|
||||
|
||||
assert output.images == ["chunk"]
|
||||
assert consumer.chunk_metadata(output) == ARDiffusionChunkMetadata.from_tick(tick)
|
||||
assert len(client.calls) == 1
|
||||
call = client.calls[0]
|
||||
assert call["request_id"] == tick.request_id
|
||||
assert call["output_modalities"] == ["video"]
|
||||
assert call["prompt"] == {
|
||||
"prompt": "turn left",
|
||||
"multi_modal_data": {"image": "initial-image"},
|
||||
}
|
||||
submitted_params = call["sampling_params_list"][0]
|
||||
assert submitted_params.extra_args["deployment_option"] == "kept"
|
||||
assert ARDiffusionTickRequest.from_extra_args(submitted_params.extra_args) == tick
|
||||
assert "ar_diffusion_tick" not in template.extra_args
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consumer_rejects_missing_chunk_metadata() -> None:
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
FakeGenerateClient(include_metadata=False),
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="multimodal_output.metadata"):
|
||||
await consumer.execute_tick(make_tick())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consumer_rejects_non_monotonic_applied_event_metadata() -> None:
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
FakeGenerateClient(applied_event_ids=[8, 7]),
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="unique and strictly increasing"):
|
||||
await consumer.execute_tick(make_tick())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consumer_preserves_engine_failure_before_metadata_validation() -> None:
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
FakeGenerateClient(include_metadata=False, output_error="model failed"),
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="model failed"):
|
||||
await consumer.execute_tick(make_tick())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_commits_only_after_standard_engine_output_metadata() -> None:
|
||||
client = FakeGenerateClient()
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
client,
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
session = ARDiffusionSession(
|
||||
"world-1",
|
||||
tick_consumer=consumer,
|
||||
lifecycle=FakeLifecycle(),
|
||||
max_pending_events=4,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-chunk-{chunk_index}",
|
||||
)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=7, prompt="turn left"))
|
||||
|
||||
output = await session.next_chunk()
|
||||
|
||||
assert output.request_id == "world-1-chunk-0"
|
||||
assert session.chunk_index == 1
|
||||
assert session.pending_event_count == 0
|
||||
assert client.calls[0]["request_id"] == "world-1-chunk-0"
|
||||
@@ -30,8 +30,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
BLOCK = 16
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
|
||||
return ChunkWindowSpec(
|
||||
@@ -97,6 +95,7 @@ def test_sliding_replace_snaps_to_chunk_boundary():
|
||||
def test_sink_chunks_protected():
|
||||
# sink = 1 chunk (16 tokens) is never skipped.
|
||||
skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
|
||||
# raw overflow = 80 - 32 - 16 = 32 -> snapped 32
|
||||
assert skip == 32
|
||||
|
||||
|
||||
@@ -202,149 +201,12 @@ def test_cross_attn_pool_deducted_from_self_attn_budget():
|
||||
kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0)),
|
||||
session_capacity=1,
|
||||
cross_attention_lengths={"text": L},
|
||||
)
|
||||
cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2 # K+V, pos+neg, layers
|
||||
page_bytes = kv.spec.page_size_bytes * 2
|
||||
scratch_bytes = kv.scratch_num_blocks * page_bytes
|
||||
expected = (int(avail * 0.5) - cross_bytes - scratch_bytes) // page_bytes
|
||||
assert kv.num_blocks == expected
|
||||
assert kv.cross_attention_reserved_bytes == cross_bytes
|
||||
assert (kv.num_blocks_total * page_bytes) + cross_bytes <= int(avail * 0.5)
|
||||
assert expected > 1 * (2 + 1) + 2 # above the local-slot minimum, so the cross deduction is what's tested
|
||||
|
||||
|
||||
def _make_tiny_capacity_kv(
|
||||
*,
|
||||
requested_capacity: int,
|
||||
available_bytes: int,
|
||||
gpu_memory_fraction: float = 1.0,
|
||||
model_owned_state_bytes_per_session: int = 0,
|
||||
) -> ARDiffusionKVCache:
|
||||
return ARDiffusionKVCache(
|
||||
ARDiffusionKVConfig(
|
||||
enable=True,
|
||||
chunk_size=1,
|
||||
window_chunks=3,
|
||||
sink_chunks=3,
|
||||
gpu_memory_fraction=gpu_memory_fraction,
|
||||
),
|
||||
num_layers=1,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
block_size=1,
|
||||
max_model_len=128,
|
||||
available_bytes=available_bytes,
|
||||
kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
|
||||
session_capacity=requested_capacity,
|
||||
cross_attention_lengths={"text": 2},
|
||||
frames_per_block=3,
|
||||
model_owned_state_bytes_per_session=model_owned_state_bytes_per_session,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
|
||||
def test_capacity_two_retains_both_windows_and_allocates_next_block():
|
||||
# LingBot-like single-branch geometry: page=8 bytes; capacity=2 requires
|
||||
# 17 managed + 3 scratch blocks, plus two 16-byte cross-attention
|
||||
# reservations: 192 bytes exactly.
|
||||
kv = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192)
|
||||
assert kv.requested_session_capacity == 2
|
||||
assert kv.session_capacity == 2
|
||||
assert kv.managed_num_blocks == 17
|
||||
|
||||
adapters = [kv.begin_request(session_id) for session_id in ("first", "second")]
|
||||
for adapter in adapters:
|
||||
for _ in range(6): # three sink + three window blocks
|
||||
kv.allocate_token_slots(adapter, 1)
|
||||
kv.commit_chunk(adapter)
|
||||
assert len(kv.window_block_ids(adapter)) == 6
|
||||
|
||||
# Only one request is in flight, so its three-frame block is counted once.
|
||||
kv.allocate_token_slots(adapters[0], 3)
|
||||
|
||||
|
||||
def test_requested_capacity_is_capped_and_cross_reservation_uses_effective_capacity():
|
||||
kv = _make_tiny_capacity_kv(requested_capacity=64, available_bytes=192)
|
||||
|
||||
assert kv.requested_session_capacity == 64
|
||||
assert kv.session_capacity == 2
|
||||
assert kv.cross_attention_bytes_per_session == 16
|
||||
assert kv.cross_attention_reserved_bytes == 32
|
||||
assert kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes == kv.memory_budget_bytes
|
||||
|
||||
|
||||
def test_capacity_rejects_budget_that_cannot_fit_one_session():
|
||||
# One session needs 11 managed + 3 scratch blocks and 16 cross bytes:
|
||||
# 14 * 8 + 16 = 128 bytes.
|
||||
with pytest.raises(ValueError, match="cannot fit one session.*available=127 bytes.*required=128 bytes"):
|
||||
_make_tiny_capacity_kv(requested_capacity=1, available_bytes=127)
|
||||
|
||||
|
||||
def test_fraction_is_soft_floor_when_one_session_fits_actual_memory():
|
||||
kv = _make_tiny_capacity_kv(
|
||||
requested_capacity=2,
|
||||
available_bytes=1280,
|
||||
gpu_memory_fraction=0.05,
|
||||
)
|
||||
|
||||
assert kv.configured_memory_budget_bytes == 64
|
||||
assert kv.memory_budget_bytes == 128
|
||||
assert kv.session_capacity == 1
|
||||
|
||||
|
||||
def test_model_owned_state_reduces_effective_session_capacity():
|
||||
kv = _make_tiny_capacity_kv(
|
||||
requested_capacity=2,
|
||||
available_bytes=200,
|
||||
model_owned_state_bytes_per_session=16,
|
||||
)
|
||||
|
||||
assert kv.session_capacity == 1
|
||||
assert kv.model_owned_state_reserved_bytes == 16
|
||||
assert (
|
||||
kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes + kv.model_owned_state_reserved_bytes
|
||||
<= kv.memory_budget_bytes
|
||||
)
|
||||
|
||||
|
||||
def _make_shipped_lingbot_geometry(*, gpu_memory_fraction: float) -> ARDiffusionKVCache:
|
||||
return ARDiffusionKVCache(
|
||||
ARDiffusionKVConfig(
|
||||
enable=True,
|
||||
chunk_size=1560,
|
||||
window_chunks=9,
|
||||
sink_chunks=9,
|
||||
gpu_memory_fraction=gpu_memory_fraction,
|
||||
),
|
||||
num_layers=40,
|
||||
num_kv_heads=40,
|
||||
head_size=128,
|
||||
dtype=torch.bfloat16,
|
||||
block_size=1560,
|
||||
max_model_len=1 << 20,
|
||||
available_bytes=100 * (1 << 30),
|
||||
kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
|
||||
session_capacity=2,
|
||||
cross_attention_lengths={"text": 512},
|
||||
frames_per_block=3,
|
||||
model_owned_state_bytes_per_session=7_488_000,
|
||||
)
|
||||
|
||||
|
||||
def test_shipped_lingbot_geometry_default_fraction_admits_one_session():
|
||||
kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.1)
|
||||
|
||||
assert kv.configured_memory_budget_bytes == 10 * (1 << 30)
|
||||
assert kv.memory_budget_bytes == 33_653_670_400
|
||||
assert kv.session_capacity == 1
|
||||
|
||||
|
||||
def test_shipped_lingbot_geometry_tuned_fraction_admits_two_sessions():
|
||||
kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.6)
|
||||
|
||||
assert kv.memory_budget_bytes == 60 * (1 << 30)
|
||||
assert kv.session_capacity == 2
|
||||
cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2 # K+V, pos+neg, layers
|
||||
expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
|
||||
assert kv.num_blocks == expected
|
||||
assert expected > 1 * (2 + 1) + 2 # above the local-slot floor, so the cross deduction is what's tested
|
||||
|
||||
|
||||
def _make_kv(
|
||||
@@ -359,24 +221,15 @@ def _make_kv(
|
||||
if local_branches == 1
|
||||
else (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1))
|
||||
)
|
||||
page_bytes = 2 * BLOCK * 4 * 64 * torch.float32.itemsize
|
||||
declared_scratch_blocks = (max_scratch_tokens_per_branch + BLOCK - 1) // BLOCK
|
||||
scratch_blocks = local_branches * (num_frame_per_block + declared_scratch_blocks)
|
||||
managed_blocks = local_branches * (window_chunks + num_frame_per_block) + 2
|
||||
return ARDiffusionKVCache(
|
||||
ARDiffusionKVConfig(
|
||||
enable=True,
|
||||
chunk_size=BLOCK,
|
||||
window_chunks=window_chunks,
|
||||
gpu_memory_fraction=1.0,
|
||||
),
|
||||
ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
|
||||
num_layers=1,
|
||||
num_kv_heads=4,
|
||||
head_size=64,
|
||||
dtype=torch.float32,
|
||||
block_size=BLOCK,
|
||||
max_model_len=4096,
|
||||
available_bytes=(managed_blocks + scratch_blocks) * page_bytes,
|
||||
available_bytes=1 << 16, # tiny -> the floor binds
|
||||
device=torch.device("cpu"),
|
||||
kv_branches=kv_branches,
|
||||
session_capacity=1,
|
||||
|
||||
@@ -54,15 +54,13 @@ def test_full_request_lifecycle_and_eviction():
|
||||
used = {int(s) // kv.block_size for s in slots}
|
||||
assert kv.null_block_id not in used
|
||||
assert used <= set(block_table)
|
||||
# Allocation may transiently add one in-flight chunk.
|
||||
# Resident window stays bounded by window + the in-flight chunk.
|
||||
assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
|
||||
kv.commit_chunk(adapter)
|
||||
# A successful commit immediately prunes physical ownership.
|
||||
assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks
|
||||
free_after.append(kv.manager.block_pool.get_num_free_blocks())
|
||||
kv.commit_chunk(adapter)
|
||||
|
||||
# Pool memory plateaus once the window is full (eviction recycles blocks).
|
||||
assert free_after[-1] == free_after[kv.spec.window_chunks - 1]
|
||||
assert free_after[-1] == free_after[kv.spec.window_chunks]
|
||||
|
||||
kv.end_request(adapter)
|
||||
assert kv.manager.block_pool.get_num_free_blocks() == free_total
|
||||
|
||||
@@ -332,7 +332,7 @@ def test_block_table_padded_to_fixed_width():
|
||||
|
||||
def test_prepare_is_idempotent_and_layers_share_metadata():
|
||||
device = torch.device("cpu")
|
||||
kv, st = make_state(num_layers=2)
|
||||
_, st = make_state(num_layers=2)
|
||||
contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)
|
||||
fctx = contexts[0].forward_ctx
|
||||
fctx.prepare(device=device, action_len=0, query_len=BLOCK)
|
||||
@@ -349,10 +349,6 @@ def test_prepare_is_idempotent_and_layers_share_metadata():
|
||||
assert i0.block_table is i1.block_table
|
||||
assert i0.seq_lens is i1.seq_lens
|
||||
assert i0.video_slots is i1.video_slots
|
||||
assert i0.key_pool is kv._k_pools[0]
|
||||
assert i0.value_pool is kv._v_pools[0]
|
||||
assert i1.key_pool is kv._k_pools[1]
|
||||
assert i1.value_pool is kv._v_pools[1]
|
||||
|
||||
|
||||
def test_layer_inputs_before_prepare_raises():
|
||||
|
||||
@@ -174,19 +174,6 @@ def test_window_reset_can_keep_named_cross_attention_cache():
|
||||
assert not st.is_cross_attention_populated(POS, "text")
|
||||
|
||||
|
||||
def test_clear_cross_attention_preserves_self_attention_history():
|
||||
kv, st = make_state(num_layers=1, cross_attn_length=8)
|
||||
_prepare_and_commit(st, POS, 2)
|
||||
k = torch.randn(1, 8, N_HEADS, HEAD_DIM)
|
||||
st.populate_cross_attention(POS, "text", [(k, k)])
|
||||
resident_before = kv.window_block_ids(st.adapter(POS))
|
||||
|
||||
st.clear_cross_attention()
|
||||
|
||||
assert kv.window_block_ids(st.adapter(POS)) == resident_before
|
||||
assert not st.is_cross_attention_populated(POS, "text")
|
||||
|
||||
|
||||
def test_cross_attention_uses_logical_branch_when_local_slot_is_shared():
|
||||
_, st = make_state(num_layers=2, cross_attn_length=8, shared_local_index=True)
|
||||
assert st.kv_cache.num_local_kv_branches == 1
|
||||
|
||||
@@ -1,505 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.ar_diffusion.session import (
|
||||
ARDiffusionChunkMetadataError,
|
||||
ARDiffusionEventAcceptanceStatus,
|
||||
ARDiffusionEventQueueFullError,
|
||||
ARDiffusionPreparedControls,
|
||||
ARDiffusionSession,
|
||||
ARDiffusionSessionError,
|
||||
ARDiffusionSessionEvent,
|
||||
ARDiffusionSessionManager,
|
||||
ARDiffusionSessionStateError,
|
||||
ARDiffusionSessionStatus,
|
||||
ARDiffusionStaleEventError,
|
||||
ARDiffusionWorkerLifecycle,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
|
||||
ARDiffusionChunkMetadata,
|
||||
ARDiffusionControlInput,
|
||||
ARDiffusionTickRequest,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeTickConsumer:
|
||||
def __init__(self) -> None:
|
||||
self.ticks: list[ARDiffusionTickRequest] = []
|
||||
self.failures_remaining = 0
|
||||
self.started: asyncio.Event | None = None
|
||||
self.release: asyncio.Event | None = None
|
||||
self.metadata_override: ARDiffusionChunkMetadata | None = None
|
||||
|
||||
async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
|
||||
self.ticks.append(tick)
|
||||
if self.started is not None:
|
||||
self.started.set()
|
||||
if self.release is not None:
|
||||
await self.release.wait()
|
||||
if self.failures_remaining:
|
||||
self.failures_remaining -= 1
|
||||
raise RuntimeError("fake tick failed")
|
||||
return self.metadata_override or ARDiffusionChunkMetadata.from_tick(tick)
|
||||
|
||||
def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata:
|
||||
return output
|
||||
|
||||
|
||||
class FakeLifecycle:
|
||||
def __init__(self) -> None:
|
||||
self.resets: list[str] = []
|
||||
self.closes: list[str] = []
|
||||
self.close_failures_remaining = 0
|
||||
|
||||
async def reset_session(self, session_id: str) -> None:
|
||||
self.resets.append(session_id)
|
||||
|
||||
async def close_session(self, session_id: str) -> None:
|
||||
self.closes.append(session_id)
|
||||
if self.close_failures_remaining:
|
||||
self.close_failures_remaining -= 1
|
||||
raise RuntimeError("fake close failed")
|
||||
|
||||
|
||||
class FailingCommitReducer:
|
||||
def __init__(self) -> None:
|
||||
self.commit_calls = 0
|
||||
self.reset_calls = 0
|
||||
|
||||
def prepare(
|
||||
self,
|
||||
*,
|
||||
current_controls: dict[str, ARDiffusionControlInput],
|
||||
events: tuple[ARDiffusionSessionEvent, ...],
|
||||
chunk_index: int,
|
||||
) -> ARDiffusionPreparedControls:
|
||||
del current_controls, events, chunk_index
|
||||
return ARDiffusionPreparedControls(controls=())
|
||||
|
||||
def commit(self, prepared: ARDiffusionPreparedControls) -> None:
|
||||
del prepared
|
||||
self.commit_calls += 1
|
||||
raise RuntimeError("fake reducer commit failed")
|
||||
|
||||
def reset(self) -> None:
|
||||
self.reset_calls += 1
|
||||
|
||||
|
||||
def make_session(
|
||||
*,
|
||||
max_pending_events: int = 8,
|
||||
consumer: FakeTickConsumer | None = None,
|
||||
lifecycle: FakeLifecycle | None = None,
|
||||
) -> tuple[ARDiffusionSession, FakeTickConsumer, FakeLifecycle]:
|
||||
resolved_consumer = consumer or FakeTickConsumer()
|
||||
resolved_lifecycle = lifecycle or FakeLifecycle()
|
||||
session = ARDiffusionSession(
|
||||
"world-1",
|
||||
tick_consumer=resolved_consumer,
|
||||
lifecycle=resolved_lifecycle,
|
||||
max_pending_events=max_pending_events,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
|
||||
)
|
||||
return session, resolved_consumer, resolved_lifecycle
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_events_are_accepted_then_sorted_and_applied_at_chunk_boundary() -> None:
|
||||
session, consumer, _ = make_session()
|
||||
event_two = ARDiffusionSessionEvent(event_id=2, prompt="turn")
|
||||
event_one = ARDiffusionSessionEvent(
|
||||
event_id=1,
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data={"samples": [{"x": 1.0}]},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
accepted_two = await session.accept_event(event_two)
|
||||
accepted_one = await session.accept_event(event_one)
|
||||
duplicate = await session.accept_event(event_two)
|
||||
|
||||
assert accepted_two.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
|
||||
assert accepted_one.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
|
||||
assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
|
||||
assert consumer.ticks == []
|
||||
|
||||
metadata = await session.next_chunk()
|
||||
|
||||
assert metadata.applied_event_ids == (1, 2)
|
||||
assert consumer.ticks[0] == ARDiffusionTickRequest(
|
||||
session_id="world-1",
|
||||
request_id="world-1-request-0",
|
||||
chunk_index=0,
|
||||
applied_event_ids=(1, 2),
|
||||
prompt="turn",
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data={"samples": [{"x": 1.0}]},
|
||||
),
|
||||
),
|
||||
)
|
||||
assert session.pending_event_count == 0
|
||||
|
||||
applied_duplicate = await session.accept_event(event_two)
|
||||
assert applied_duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
|
||||
|
||||
next_metadata = await session.next_chunk()
|
||||
assert next_metadata.chunk_index == 1
|
||||
assert next_metadata.applied_event_ids == ()
|
||||
assert consumer.ticks[1].prompt == "turn"
|
||||
assert consumer.ticks[1].controls == consumer.ticks[0].controls
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_composite_event_is_isolated_and_queue_admission_is_atomic() -> None:
|
||||
session, consumer, _ = make_session(max_pending_events=1)
|
||||
payload = {"samples": [{"x": 1.0}]}
|
||||
composite = ARDiffusionSessionEvent(
|
||||
event_id=5,
|
||||
prompt="new prompt",
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data=payload,
|
||||
),
|
||||
ARDiffusionControlInput(
|
||||
track="action",
|
||||
schema="example.action.v1",
|
||||
data={"keys": ["left"]},
|
||||
),
|
||||
),
|
||||
)
|
||||
await session.accept_event(composite)
|
||||
payload["samples"][0]["x"] = 99.0
|
||||
|
||||
duplicate = await session.accept_event(composite)
|
||||
assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
|
||||
with pytest.raises(ARDiffusionEventQueueFullError, match="queue is full"):
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=6, prompt="rejected"))
|
||||
|
||||
await session.next_chunk()
|
||||
assert consumer.ticks[0].applied_event_ids == (5,)
|
||||
assert consumer.ticks[0].prompt == "new prompt"
|
||||
assert consumer.ticks[0].controls[0].data["keys"] == ("left",)
|
||||
assert consumer.ticks[0].controls[1].data["samples"] == ({"x": 1.0},)
|
||||
assert consumer.ticks[0].controls[0].to_dict()["data"] == {"keys": ["left"]}
|
||||
assert consumer.ticks[0].controls[1].to_dict()["data"] == {"samples": [{"x": 1.0}]}
|
||||
|
||||
|
||||
def test_event_rejects_duplicate_tracks_and_non_transport_payloads() -> None:
|
||||
with pytest.raises(ValueError, match="at most once"):
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=1,
|
||||
controls=(
|
||||
ARDiffusionControlInput(track="camera", schema="one", data={}),
|
||||
ARDiffusionControlInput(track="camera", schema="two", data={}),
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="transport-safe"):
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data={"object": object()},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_tick_closes_worker_and_requires_reset() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
consumer.failures_remaining = 1
|
||||
lifecycle = FakeLifecycle()
|
||||
session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=3, prompt="retry me"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="fake tick failed"):
|
||||
await session.next_chunk()
|
||||
|
||||
assert session.status is ARDiffusionSessionStatus.FAILED
|
||||
assert session.chunk_index == 0
|
||||
assert session.pending_event_count == 1
|
||||
assert lifecycle.closes == ["world-1"]
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="failed"):
|
||||
await session.next_chunk()
|
||||
|
||||
await session.reset()
|
||||
assert session.status is ARDiffusionSessionStatus.ACTIVE
|
||||
assert session.pending_event_count == 0
|
||||
assert lifecycle.resets == ["world-1"]
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="after reset"))
|
||||
metadata = await session.next_chunk()
|
||||
assert metadata.chunk_index == 0
|
||||
assert metadata.applied_event_ids == (4,)
|
||||
assert [tick.applied_event_ids for tick in consumer.ticks] == [(3,), (4,)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inflight_snapshot_is_immutable_and_new_event_waits_for_next_chunk() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
consumer.started = asyncio.Event()
|
||||
consumer.release = asyncio.Event()
|
||||
session, _, _ = make_session(consumer=consumer)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="first"))
|
||||
|
||||
first_chunk = asyncio.create_task(session.next_chunk())
|
||||
await consumer.started.wait()
|
||||
accepted = await session.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="second"))
|
||||
assert accepted.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
|
||||
assert consumer.ticks[0].prompt == "first"
|
||||
assert consumer.ticks[0].applied_event_ids == (10,)
|
||||
|
||||
with pytest.raises(ARDiffusionStaleEventError, match="event floor is 10"):
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=9, prompt="stale"))
|
||||
|
||||
consumer.release.set()
|
||||
first_metadata = await first_chunk
|
||||
second_metadata = await session.next_chunk()
|
||||
assert first_metadata.applied_event_ids == (10,)
|
||||
assert second_metadata.applied_event_ids == (11,)
|
||||
assert consumer.ticks[1].prompt == "second"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metadata_mismatch_does_not_commit_snapshot() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
consumer.metadata_override = ARDiffusionChunkMetadata(
|
||||
session_id="another-world",
|
||||
request_id="wrong",
|
||||
chunk_index=0,
|
||||
)
|
||||
lifecycle = FakeLifecycle()
|
||||
session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
|
||||
|
||||
with pytest.raises(ARDiffusionChunkMetadataError, match="does not match"):
|
||||
await session.next_chunk()
|
||||
|
||||
assert session.chunk_index == 0
|
||||
assert session.pending_event_count == 1
|
||||
assert session.status is ARDiffusionSessionStatus.FAILED
|
||||
assert lifecycle.closes == ["world-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reducer_commit_failure_is_fail_closed_without_partial_session_commit() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
lifecycle = FakeLifecycle()
|
||||
reducer = FailingCommitReducer()
|
||||
session = ARDiffusionSession(
|
||||
"world-1",
|
||||
tick_consumer=consumer,
|
||||
lifecycle=lifecycle,
|
||||
max_pending_events=4,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
|
||||
control_reducer=reducer,
|
||||
)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="fake reducer commit failed"):
|
||||
await session.next_chunk()
|
||||
|
||||
assert reducer.commit_calls == 1
|
||||
assert session.status is ARDiffusionSessionStatus.FAILED
|
||||
assert session.chunk_index == 0
|
||||
assert session.pending_event_count == 1
|
||||
assert lifecycle.closes == ["world-1"]
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="failed"):
|
||||
await session.next_chunk()
|
||||
|
||||
await session.reset()
|
||||
assert reducer.reset_calls == 1
|
||||
assert session.status is ARDiffusionSessionStatus.ACTIVE
|
||||
assert session.chunk_index == 0
|
||||
assert session.pending_event_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_close_and_disconnect_reach_lifecycle_boundary() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
lifecycle = FakeLifecycle()
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=consumer,
|
||||
lifecycle=lifecycle,
|
||||
max_pending_events=4,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}",
|
||||
)
|
||||
session = await manager.create_session("world-reset")
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="before reset"))
|
||||
await session.next_chunk()
|
||||
await manager.reset_session("world-reset")
|
||||
|
||||
assert lifecycle.resets == ["world-reset"]
|
||||
assert session.status is ARDiffusionSessionStatus.ACTIVE
|
||||
assert session.chunk_index == 0
|
||||
duplicate = await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="duplicate"))
|
||||
assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
|
||||
|
||||
reset_tick = await session.next_chunk()
|
||||
assert reset_tick.chunk_index == 0
|
||||
assert consumer.ticks[-1].prompt is None
|
||||
assert consumer.ticks[-1].controls == ()
|
||||
|
||||
await manager.close_session("world-reset")
|
||||
assert lifecycle.closes == ["world-reset"]
|
||||
assert session.status is ARDiffusionSessionStatus.CLOSED
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
|
||||
await manager.get_session("world-reset")
|
||||
|
||||
disconnected = await manager.create_session("world-disconnect")
|
||||
await manager.disconnect("world-disconnect")
|
||||
assert disconnected.status is ARDiffusionSessionStatus.CLOSED
|
||||
assert lifecycle.closes[-1] == "world-disconnect"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manager_keeps_interleaved_sessions_independent() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=consumer,
|
||||
lifecycle=FakeLifecycle(),
|
||||
max_pending_events=4,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
|
||||
)
|
||||
session_a = await manager.create_session("world-a")
|
||||
session_b = await manager.create_session("world-b")
|
||||
|
||||
await session_a.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="a0"))
|
||||
await session_a.next_chunk()
|
||||
await session_b.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="b0"))
|
||||
await session_b.next_chunk()
|
||||
await session_a.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="a1"))
|
||||
await session_a.next_chunk()
|
||||
|
||||
assert [
|
||||
(
|
||||
tick.session_id,
|
||||
tick.request_id,
|
||||
tick.chunk_index,
|
||||
tick.applied_event_ids,
|
||||
)
|
||||
for tick in consumer.ticks
|
||||
] == [
|
||||
("world-a", "world-a-request-0", 0, (10,)),
|
||||
("world-b", "world-b-request-0", 0, (1,)),
|
||||
("world-a", "world-a-request-1", 1, (11,)),
|
||||
]
|
||||
assert session_a.chunk_index == 2
|
||||
assert session_b.chunk_index == 1
|
||||
with pytest.raises(ARDiffusionStaleEventError, match="event floor is 11"):
|
||||
await session_a.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="stale only for a"))
|
||||
assert (
|
||||
await session_b.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="b1"))
|
||||
).status is ARDiffusionEventAcceptanceStatus.ACCEPTED
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_close_retains_tombstone_until_cleanup_retry_succeeds() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
lifecycle = FakeLifecycle()
|
||||
lifecycle.close_failures_remaining = 1
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=consumer,
|
||||
lifecycle=lifecycle,
|
||||
max_pending_events=4,
|
||||
)
|
||||
session = await manager.create_session("world-close-retry")
|
||||
|
||||
with pytest.raises(RuntimeError, match="fake close failed"):
|
||||
await manager.close_session("world-close-retry")
|
||||
|
||||
assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
|
||||
assert await manager.get_session("world-close-retry") is session
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="already exists"):
|
||||
await manager.create_session("world-close-retry")
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="cleanup_failed"):
|
||||
await manager.reset_session("world-close-retry")
|
||||
|
||||
await manager.close_session("world-close-retry")
|
||||
|
||||
assert lifecycle.closes == ["world-close-retry", "world-close-retry"]
|
||||
assert session.status is ARDiffusionSessionStatus.CLOSED
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
|
||||
await manager.get_session("world-close-retry")
|
||||
recreated = await manager.create_session("world-close-retry")
|
||||
assert recreated is not session
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_disconnect_retains_session_for_explicit_cleanup_retry() -> None:
|
||||
lifecycle = FakeLifecycle()
|
||||
lifecycle.close_failures_remaining = 1
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=FakeTickConsumer(),
|
||||
lifecycle=lifecycle,
|
||||
max_pending_events=4,
|
||||
)
|
||||
session = await manager.create_session("world-disconnect-retry")
|
||||
|
||||
with pytest.raises(RuntimeError, match="fake close failed"):
|
||||
await manager.disconnect("world-disconnect-retry")
|
||||
|
||||
assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
|
||||
assert await manager.get_session("world-disconnect-retry") is session
|
||||
|
||||
await manager.close_session("world-disconnect-retry")
|
||||
assert session.status is ARDiffusionSessionStatus.CLOSED
|
||||
|
||||
|
||||
class FakeRPCClient:
|
||||
def __init__(self, result: list[Any]) -> None:
|
||||
self.result = result
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
async def collective_rpc(self, **kwargs: Any) -> list[Any]:
|
||||
self.calls.append(kwargs)
|
||||
return self.result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worker_lifecycle_routes_to_selected_stage_rpc() -> None:
|
||||
client = FakeRPCClient([[True]])
|
||||
lifecycle = ARDiffusionWorkerLifecycle(client, stage_ids=[2], timeout=3.0)
|
||||
|
||||
await lifecycle.reset_session("world-1")
|
||||
await lifecycle.close_session("world-1")
|
||||
|
||||
assert client.calls == [
|
||||
{
|
||||
"method": "reset_ar_diffusion_session",
|
||||
"timeout": 3.0,
|
||||
"args": ("world-1",),
|
||||
"stage_ids": [2],
|
||||
},
|
||||
{
|
||||
"method": "close_ar_diffusion_session",
|
||||
"timeout": 3.0,
|
||||
"args": ("world-1",),
|
||||
"stage_ids": [2],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worker_lifecycle_rejects_unsupported_stage_result() -> None:
|
||||
lifecycle = ARDiffusionWorkerLifecycle(
|
||||
FakeRPCClient([{"supported": False, "error": "not an AR stage"}]),
|
||||
stage_ids=[1],
|
||||
)
|
||||
|
||||
with pytest.raises(ARDiffusionSessionError, match="not an AR stage"):
|
||||
await lifecycle.close_session("world-1")
|
||||
@@ -1,108 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
|
||||
ARDiffusionChunkMetadata,
|
||||
ARDiffusionControlInput,
|
||||
ARDiffusionTickRequest,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_tick_round_trip_preserves_identity_and_control_snapshot() -> None:
|
||||
tick = ARDiffusionTickRequest(
|
||||
session_id="world-1",
|
||||
request_id="chunk-request-7",
|
||||
chunk_index=7,
|
||||
applied_event_ids=(10, 11),
|
||||
prompt="turn left",
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="lingbot.camera_trajectory.v1",
|
||||
data={"poses": [[1.0, 0.0]], "intrinsics": [[2.0]]},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
restored = ARDiffusionTickRequest.from_extra_args(tick.to_extra_args())
|
||||
|
||||
assert restored == tick
|
||||
assert ARDiffusionChunkMetadata.from_tick(tick).to_dict() == {
|
||||
"session_id": "world-1",
|
||||
"request_id": "chunk-request-7",
|
||||
"chunk_index": 7,
|
||||
"applied_event_ids": [10, 11],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value", "message"),
|
||||
[
|
||||
("session_id", "", "session_id"),
|
||||
("request_id", "", "request_id"),
|
||||
("chunk_index", -1, "chunk_index"),
|
||||
("applied_event_ids", (2, 1), "strictly increasing"),
|
||||
("applied_event_ids", (1, 1), "strictly increasing"),
|
||||
],
|
||||
)
|
||||
def test_tick_rejects_invalid_identity(
|
||||
field: str,
|
||||
value: object,
|
||||
message: str,
|
||||
) -> None:
|
||||
values = {
|
||||
"session_id": "world-1",
|
||||
"request_id": "request-1",
|
||||
"chunk_index": 0,
|
||||
"applied_event_ids": (),
|
||||
}
|
||||
values[field] = value
|
||||
|
||||
with pytest.raises(ValueError, match=message):
|
||||
ARDiffusionTickRequest(**values)
|
||||
|
||||
|
||||
def test_tick_snapshot_is_detached_and_deeply_immutable() -> None:
|
||||
payload = {"samples": [{"keys": ["w"]}]}
|
||||
event_ids = [1, 2]
|
||||
controls = [
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data=payload,
|
||||
)
|
||||
]
|
||||
tick = ARDiffusionTickRequest(
|
||||
session_id="world-1",
|
||||
request_id="request-a",
|
||||
chunk_index=0,
|
||||
applied_event_ids=event_ids, # type: ignore[arg-type]
|
||||
controls=controls, # type: ignore[arg-type]
|
||||
)
|
||||
serialized = tick.to_extra_args()
|
||||
|
||||
payload["samples"][0]["keys"].append("j")
|
||||
event_ids.append(3)
|
||||
controls.clear()
|
||||
serialized["ar_diffusion_tick"]["controls"][0]["data"]["samples"][0]["keys"].append("k")
|
||||
|
||||
assert tick.applied_event_ids == (1, 2)
|
||||
assert isinstance(tick.controls, tuple)
|
||||
assert tick.to_extra_args()["ar_diffusion_tick"]["controls"] == [
|
||||
{
|
||||
"track": "camera",
|
||||
"schema": "example.camera.v1",
|
||||
"data": {"samples": [{"keys": ["w"]}]},
|
||||
}
|
||||
]
|
||||
with pytest.raises(TypeError):
|
||||
tick.controls[0].data["samples"] = () # type: ignore[index]
|
||||
with pytest.raises(AttributeError):
|
||||
tick.controls[0].data["samples"][0]["keys"].append("j")
|
||||
|
||||
|
||||
def test_legacy_request_has_no_typed_tick() -> None:
|
||||
assert ARDiffusionTickRequest.from_extra_args({"session_id": "legacy"}) is None
|
||||
@@ -1,53 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeARDiffusionRunner:
|
||||
def __init__(self) -> None:
|
||||
self.resets: list[str] = []
|
||||
self.closes: list[str] = []
|
||||
|
||||
def reset_session(self, session_id: str) -> None:
|
||||
self.resets.append(session_id)
|
||||
|
||||
def close_session(self, session_id: str) -> None:
|
||||
self.closes.append(session_id)
|
||||
|
||||
|
||||
class FakeRegularRunner:
|
||||
pass
|
||||
|
||||
|
||||
def make_worker(runner: object) -> DiffusionWorker:
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
worker.model_runner = runner
|
||||
return worker
|
||||
|
||||
|
||||
def test_worker_delegates_ar_session_lifecycle_to_runner() -> None:
|
||||
runner = FakeARDiffusionRunner()
|
||||
worker = make_worker(runner)
|
||||
|
||||
assert worker.reset_ar_diffusion_session("world-1") is True
|
||||
assert worker.close_ar_diffusion_session("world-1") is True
|
||||
assert runner.resets == ["world-1"]
|
||||
assert runner.closes == ["world-1"]
|
||||
|
||||
|
||||
def test_worker_reports_unsupported_runner_without_model_assumptions() -> None:
|
||||
worker = make_worker(FakeRegularRunner())
|
||||
|
||||
assert worker.reset_ar_diffusion_session("world-1") is False
|
||||
assert worker.close_ar_diffusion_session("world-1") is False
|
||||
|
||||
|
||||
def test_worker_rejects_empty_session_id() -> None:
|
||||
worker = make_worker(FakeARDiffusionRunner())
|
||||
|
||||
with pytest.raises(ValueError, match="session_id"):
|
||||
worker.close_ar_diffusion_session("")
|
||||
@@ -6,10 +6,7 @@ from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import (
|
||||
DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION,
|
||||
DreamZeroPipeline,
|
||||
)
|
||||
from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
|
||||
from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -35,37 +32,6 @@ def test_dreamzero_pipeline_state_is_session_keyed() -> None:
|
||||
assert session_b.call_count == 3
|
||||
|
||||
|
||||
def test_dreamzero_kv_spec_accounts_for_model_owned_cuda_state(monkeypatch) -> None:
|
||||
pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
|
||||
pipeline.transformer = SimpleNamespace(
|
||||
frame_seqlen=16,
|
||||
blocks=[
|
||||
SimpleNamespace(
|
||||
self_attn=SimpleNamespace(
|
||||
max_attention_size=96,
|
||||
tp_num_heads=4,
|
||||
)
|
||||
)
|
||||
],
|
||||
text_len=8,
|
||||
model_type="t2v",
|
||||
num_layers=2,
|
||||
num_heads=4,
|
||||
dim=256,
|
||||
num_frame_per_block=4,
|
||||
num_action_per_block=2,
|
||||
num_state_per_block=1,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero.get_classifier_free_guidance_world_size",
|
||||
lambda: 1,
|
||||
)
|
||||
|
||||
spec = pipeline.ar_diffusion_kv_cache_spec()
|
||||
|
||||
assert spec.model_owned_state_bytes_per_session == DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION
|
||||
|
||||
|
||||
def test_dreamzero_pipeline_state_follows_runner_lifecycle_notifications() -> None:
|
||||
pipeline = _empty_pipeline()
|
||||
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
@@ -1,13 +0,0 @@
|
||||
{
|
||||
"_class_name": "LingBotWorldCausalDMDPipeline",
|
||||
"_diffusers_version": "0.35.0.dev0",
|
||||
"boundary_ratio": 0.9,
|
||||
"image_encoder": [null, null],
|
||||
"image_processor": [null, null],
|
||||
"scheduler": ["diffusers", "UniPCMultistepScheduler"],
|
||||
"text_encoder": ["transformers", "UMT5EncoderModel"],
|
||||
"tokenizer": ["transformers", "T5TokenizerFast"],
|
||||
"transformer": ["diffusers", "CausalLingBotWorldTransformer3DModel"],
|
||||
"transformer_2": [null, null],
|
||||
"vae": ["diffusers", "AutoencoderKLWan"]
|
||||
}
|
||||
-42
@@ -1,42 +0,0 @@
|
||||
{
|
||||
"sources": {
|
||||
"config": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/config.json",
|
||||
"shard_1_header": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model-00001-of-00008.safetensors",
|
||||
"shard_8_header": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model-00008-of-00008.safetensors"
|
||||
},
|
||||
"dtype": "F32",
|
||||
"shard_representatives": {
|
||||
"diffusion_pytorch_model-00001-of-00008.safetensors": "blocks.0.cam_injector_layer1.bias",
|
||||
"diffusion_pytorch_model-00002-of-00008.safetensors": "blocks.4.cam_injector_layer1.bias",
|
||||
"diffusion_pytorch_model-00003-of-00008.safetensors": "blocks.10.cam_injector_layer1.bias",
|
||||
"diffusion_pytorch_model-00004-of-00008.safetensors": "blocks.14.cam_scale_layer.bias",
|
||||
"diffusion_pytorch_model-00005-of-00008.safetensors": "blocks.20.cam_injector_layer1.bias",
|
||||
"diffusion_pytorch_model-00006-of-00008.safetensors": "blocks.25.cam_injector_layer1.bias",
|
||||
"diffusion_pytorch_model-00007-of-00008.safetensors": "blocks.30.cam_injector_layer1.bias",
|
||||
"diffusion_pytorch_model-00008-of-00008.safetensors": "blocks.35.cam_injector_layer1.bias"
|
||||
},
|
||||
"parameter_shapes": {
|
||||
"blocks.0.cam_injector_layer1.bias": [5120],
|
||||
"blocks.4.cam_injector_layer1.bias": [5120],
|
||||
"blocks.10.cam_injector_layer1.bias": [5120],
|
||||
"blocks.14.cam_scale_layer.bias": [5120],
|
||||
"blocks.20.cam_injector_layer1.bias": [5120],
|
||||
"blocks.25.cam_injector_layer1.bias": [5120],
|
||||
"blocks.30.cam_injector_layer1.bias": [5120],
|
||||
"blocks.35.cam_injector_layer1.bias": [5120],
|
||||
"patch_embedding.weight": [5120, 36, 1, 2, 2],
|
||||
"patch_embedding_wancamctrl.weight": [5120, 1536],
|
||||
"c2ws_hidden_states_layer1.weight": [5120, 5120],
|
||||
"c2ws_hidden_states_layer2.weight": [5120, 5120],
|
||||
"blocks.0.modulation": [1, 6, 5120],
|
||||
"blocks.0.self_attn.q.weight": [5120, 5120],
|
||||
"blocks.0.ffn.0.weight": [13824, 5120],
|
||||
"blocks.0.ffn.2.weight": [5120, 13824],
|
||||
"blocks.0.cam_injector_layer1.weight": [5120, 5120],
|
||||
"blocks.0.cam_injector_layer2.weight": [5120, 5120],
|
||||
"blocks.0.cam_scale_layer.weight": [5120, 5120],
|
||||
"blocks.0.cam_shift_layer.weight": [5120, 5120],
|
||||
"head.modulation": [1, 2, 5120],
|
||||
"head.head.weight": [64, 5120]
|
||||
}
|
||||
}
|
||||
-30
@@ -1,30 +0,0 @@
|
||||
{
|
||||
"_class_name": "UniPCMultistepScheduler",
|
||||
"_diffusers_version": "0.35.0.dev0",
|
||||
"beta_end": 0.02,
|
||||
"beta_schedule": "linear",
|
||||
"beta_start": 0.0001,
|
||||
"disable_corrector": [],
|
||||
"dynamic_thresholding_ratio": 0.995,
|
||||
"final_sigmas_type": "zero",
|
||||
"flow_shift": 5.0,
|
||||
"lower_order_final": true,
|
||||
"num_train_timesteps": 1000,
|
||||
"predict_x0": true,
|
||||
"prediction_type": "flow_prediction",
|
||||
"rescale_betas_zero_snr": false,
|
||||
"sample_max_value": 1.0,
|
||||
"solver_order": 2,
|
||||
"solver_p": null,
|
||||
"solver_type": "bh2",
|
||||
"steps_offset": 0,
|
||||
"thresholding": false,
|
||||
"time_shift_type": "exponential",
|
||||
"timestep_spacing": "linspace",
|
||||
"trained_betas": null,
|
||||
"use_beta_sigmas": false,
|
||||
"use_dynamic_shifting": false,
|
||||
"use_exponential_sigmas": false,
|
||||
"use_flow_sigmas": true,
|
||||
"use_karras_sigmas": false
|
||||
}
|
||||
-112
@@ -1,112 +0,0 @@
|
||||
{
|
||||
"source": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model.safetensors.index.json",
|
||||
"index_parameter_count": 1421,
|
||||
"canonical_name_set_sha256": "5d31b7f64fd5e5e271b789b2dbe50fb7bbf146bd8b750646bdc9abb4e5b8139d",
|
||||
"canonical_name_to_shard_sha256": "e1addc27d2b1fad6f10226a23a265ee3e3c61e74dce4d5859727bdf180a67992",
|
||||
"canonical_serialization": "sorted lines formatted as parameter_name=shard_filename with LF separators and no trailing LF",
|
||||
"shards": {
|
||||
"diffusion_pytorch_model-00001-of-00008.safetensors": {
|
||||
"parameter_count": 180,
|
||||
"parameter_names_sha256": "81a3d4778de604e75609aae7dac3d8dcbc081811e7748267fe953bbeb375e79d"
|
||||
},
|
||||
"diffusion_pytorch_model-00002-of-00008.safetensors": {
|
||||
"parameter_count": 180,
|
||||
"parameter_names_sha256": "4e36e2596a6468f29179fa5795396f2e9a7f7d8debc912b34a6f6c2c7c7c610d"
|
||||
},
|
||||
"diffusion_pytorch_model-00003-of-00008.safetensors": {
|
||||
"parameter_count": 180,
|
||||
"parameter_names_sha256": "6f74a7b2cf7444c2daa2b3fc32bd3787589b35543c230c5c13b89f0f0f2a0998"
|
||||
},
|
||||
"diffusion_pytorch_model-00004-of-00008.safetensors": {
|
||||
"parameter_count": 180,
|
||||
"parameter_names_sha256": "ec5c967ed5271ba7834337c4ee93a73f8ed0eead4492d4920624251c65c058ac"
|
||||
},
|
||||
"diffusion_pytorch_model-00005-of-00008.safetensors": {
|
||||
"parameter_count": 180,
|
||||
"parameter_names_sha256": "ad1fdab5830dd9ce43ea007f112b4be223327b7c8197c766d9cf9219e74d619b"
|
||||
},
|
||||
"diffusion_pytorch_model-00006-of-00008.safetensors": {
|
||||
"parameter_count": 180,
|
||||
"parameter_names_sha256": "fd58f0d86be42d8afcf55d3e07d8a22f7f6ef5064da931a600c890b71b6dfbae"
|
||||
},
|
||||
"diffusion_pytorch_model-00007-of-00008.safetensors": {
|
||||
"parameter_count": 180,
|
||||
"parameter_names_sha256": "305a1714b5f4fb54d5877c593b451273ab1490d90fd0896a299d96e140295047"
|
||||
},
|
||||
"diffusion_pytorch_model-00008-of-00008.safetensors": {
|
||||
"parameter_count": 161,
|
||||
"parameter_names_sha256": "4c18d94d96168682eb212f93e3cdf7313928776d8d9e6e4de0255fe899decb25"
|
||||
}
|
||||
},
|
||||
"top_level_modules": [
|
||||
"blocks",
|
||||
"c2ws_hidden_states_layer1",
|
||||
"c2ws_hidden_states_layer2",
|
||||
"head",
|
||||
"patch_embedding",
|
||||
"patch_embedding_wancamctrl",
|
||||
"text_embedding",
|
||||
"time_embedding",
|
||||
"time_projection"
|
||||
],
|
||||
"non_block_parameters": [
|
||||
"c2ws_hidden_states_layer1.bias",
|
||||
"c2ws_hidden_states_layer1.weight",
|
||||
"c2ws_hidden_states_layer2.bias",
|
||||
"c2ws_hidden_states_layer2.weight",
|
||||
"head.head.bias",
|
||||
"head.head.weight",
|
||||
"head.modulation",
|
||||
"patch_embedding.bias",
|
||||
"patch_embedding.weight",
|
||||
"patch_embedding_wancamctrl.bias",
|
||||
"patch_embedding_wancamctrl.weight",
|
||||
"text_embedding.0.bias",
|
||||
"text_embedding.0.weight",
|
||||
"text_embedding.2.bias",
|
||||
"text_embedding.2.weight",
|
||||
"time_embedding.0.bias",
|
||||
"time_embedding.0.weight",
|
||||
"time_embedding.2.bias",
|
||||
"time_embedding.2.weight",
|
||||
"time_projection.1.bias",
|
||||
"time_projection.1.weight"
|
||||
],
|
||||
"block_parameter_suffixes": [
|
||||
"cam_injector_layer1.bias",
|
||||
"cam_injector_layer1.weight",
|
||||
"cam_injector_layer2.bias",
|
||||
"cam_injector_layer2.weight",
|
||||
"cam_scale_layer.bias",
|
||||
"cam_scale_layer.weight",
|
||||
"cam_shift_layer.bias",
|
||||
"cam_shift_layer.weight",
|
||||
"cross_attn.k.bias",
|
||||
"cross_attn.k.weight",
|
||||
"cross_attn.norm_k.weight",
|
||||
"cross_attn.norm_q.weight",
|
||||
"cross_attn.o.bias",
|
||||
"cross_attn.o.weight",
|
||||
"cross_attn.q.bias",
|
||||
"cross_attn.q.weight",
|
||||
"cross_attn.v.bias",
|
||||
"cross_attn.v.weight",
|
||||
"ffn.0.bias",
|
||||
"ffn.0.weight",
|
||||
"ffn.2.bias",
|
||||
"ffn.2.weight",
|
||||
"modulation",
|
||||
"norm3.bias",
|
||||
"norm3.weight",
|
||||
"self_attn.k.bias",
|
||||
"self_attn.k.weight",
|
||||
"self_attn.norm_k.weight",
|
||||
"self_attn.norm_q.weight",
|
||||
"self_attn.o.bias",
|
||||
"self_attn.o.weight",
|
||||
"self_attn.q.bias",
|
||||
"self_attn.q.weight",
|
||||
"self_attn.v.bias",
|
||||
"self_attn.v.weight"
|
||||
]
|
||||
}
|
||||
@@ -1,294 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.lingbot_world.actions import (
|
||||
LINGBOT_CAMERA_ACTION_SCHEMA,
|
||||
LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
|
||||
LingBotCameraControlReducer,
|
||||
integrate_lingbot_camera_actions,
|
||||
parse_lingbot_camera_action_frames,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.session import (
|
||||
ARDiffusionSession,
|
||||
ARDiffusionSessionEvent,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
|
||||
ARDiffusionChunkMetadata,
|
||||
ARDiffusionControlInput,
|
||||
ARDiffusionTickRequest,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _action_control(mode: str, **data) -> ARDiffusionControlInput:
|
||||
return ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema=LINGBOT_CAMERA_ACTION_SCHEMA,
|
||||
data={"mode": mode, **data},
|
||||
)
|
||||
|
||||
|
||||
def _prepared_frames(prepared) -> tuple[tuple[str, ...], ...]:
|
||||
assert len(prepared.controls) == 1
|
||||
return prepared.controls[0].data["frames"]
|
||||
|
||||
|
||||
def test_state_controls_hold_release_and_preserve_short_pulse() -> None:
|
||||
reducer = LingBotCameraControlReducer()
|
||||
pressed = reducer.prepare(
|
||||
current_controls={},
|
||||
events=(
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=1,
|
||||
controls=(
|
||||
_action_control(
|
||||
"state",
|
||||
transitions=[{"actions": ["W"], "client_ts_ms": 10}],
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
chunk_index=0,
|
||||
)
|
||||
assert _prepared_frames(pressed) == (("w",), ("w",), ("w",))
|
||||
assert pressed.controls[0].to_dict()["data"]["frames"] == [["w"], ["w"], ["w"]]
|
||||
reducer.commit(pressed)
|
||||
|
||||
held = reducer.prepare(
|
||||
current_controls={"camera": pressed.controls[0]},
|
||||
events=(),
|
||||
chunk_index=1,
|
||||
)
|
||||
assert _prepared_frames(held) == (("w",), ("w",), ("w",))
|
||||
reducer.commit(held)
|
||||
|
||||
pulse = reducer.prepare(
|
||||
current_controls={"camera": held.controls[0]},
|
||||
events=(
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=2,
|
||||
controls=(
|
||||
_action_control(
|
||||
"state",
|
||||
transitions=[
|
||||
{"actions": ["j"], "client_ts_ms": 20},
|
||||
{"actions": [], "client_ts_ms": 21},
|
||||
],
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
chunk_index=2,
|
||||
)
|
||||
assert _prepared_frames(pulse) == (("j",), (), ())
|
||||
reducer.commit(pulse)
|
||||
|
||||
released = reducer.prepare(
|
||||
current_controls={"camera": pulse.controls[0]},
|
||||
events=(),
|
||||
chunk_index=3,
|
||||
)
|
||||
assert _prepared_frames(released) == ((), (), ())
|
||||
|
||||
|
||||
def test_script_is_transactional_consumed_once_and_neutral_padded() -> None:
|
||||
reducer = LingBotCameraControlReducer()
|
||||
event = ARDiffusionSessionEvent(
|
||||
event_id=1,
|
||||
controls=(
|
||||
_action_control(
|
||||
"script",
|
||||
frames=[["w"], ["d"], ["l"], ["i"]],
|
||||
),
|
||||
),
|
||||
)
|
||||
first_attempt = reducer.prepare(
|
||||
current_controls={},
|
||||
events=(event,),
|
||||
chunk_index=0,
|
||||
)
|
||||
assert _prepared_frames(first_attempt) == (("w",), ("d",), ("l",))
|
||||
|
||||
# An uncommitted failed attempt must not consume the reducer's script.
|
||||
second_attempt = reducer.prepare(
|
||||
current_controls={},
|
||||
events=(event,),
|
||||
chunk_index=0,
|
||||
)
|
||||
assert _prepared_frames(second_attempt) == (("w",), ("d",), ("l",))
|
||||
reducer.commit(second_attempt)
|
||||
|
||||
tail = reducer.prepare(
|
||||
current_controls={"camera": second_attempt.controls[0]},
|
||||
events=(),
|
||||
chunk_index=1,
|
||||
)
|
||||
assert _prepared_frames(tail) == (("i",), (), ())
|
||||
reducer.commit(tail)
|
||||
exhausted = reducer.prepare(
|
||||
current_controls={"camera": tail.controls[0]},
|
||||
events=(),
|
||||
chunk_index=2,
|
||||
)
|
||||
assert _prepared_frames(exhausted) == ((), (), ())
|
||||
|
||||
|
||||
def test_trajectory_control_remains_an_expert_override() -> None:
|
||||
reducer = LingBotCameraControlReducer()
|
||||
trajectory = ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema=LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
|
||||
data={"poses": [[[1.0]]], "intrinsics": [[1.0]]},
|
||||
)
|
||||
prepared = reducer.prepare(
|
||||
current_controls={},
|
||||
events=(
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=1,
|
||||
controls=(trajectory,),
|
||||
),
|
||||
),
|
||||
chunk_index=0,
|
||||
)
|
||||
assert prepared.controls == (trajectory,)
|
||||
|
||||
|
||||
def test_action_integrator_matches_lingbot_motion_and_calibration() -> None:
|
||||
trajectory, pitch = integrate_lingbot_camera_actions(
|
||||
[["w"], ["d"], ["l"]],
|
||||
width=832,
|
||||
height=480,
|
||||
)
|
||||
|
||||
assert pitch == 0.0
|
||||
torch.testing.assert_close(
|
||||
trajectory.poses[0, :3, 3],
|
||||
torch.tensor([0.0, 0.0, 0.05], dtype=torch.float64),
|
||||
atol=1e-6,
|
||||
rtol=0,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
trajectory.poses[1, :3, 3],
|
||||
torch.tensor([0.05, 0.0, 0.05], dtype=torch.float64),
|
||||
atol=1e-6,
|
||||
rtol=0,
|
||||
)
|
||||
expected_yaw = torch.tensor(
|
||||
[
|
||||
[0.9945219, 0.0, 0.10452846],
|
||||
[0.0, 1.0, 0.0],
|
||||
[-0.10452846, 0.0, 0.9945219],
|
||||
],
|
||||
dtype=torch.float64,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
trajectory.poses[2, :3, :3],
|
||||
expected_yaw,
|
||||
atol=1e-6,
|
||||
rtol=0,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
trajectory.intrinsics[0],
|
||||
torch.tensor([500.0, 500.0, 416.0, 240.0]),
|
||||
)
|
||||
|
||||
resized, _ = integrate_lingbot_camera_actions(
|
||||
[[]],
|
||||
width=416,
|
||||
height=240,
|
||||
)
|
||||
scaled = resized.intrinsics[0] * torch.tensor([416 / 832, 240 / 480, 416 / 832, 240 / 480])
|
||||
torch.testing.assert_close(
|
||||
scaled,
|
||||
torch.tensor([500.0, 500.0, 208.0, 120.0]),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("data", "message"),
|
||||
[
|
||||
({"mode": "frames", "frames": [["x"], [], []]}, "W/A/S/D/I/J/K/L"),
|
||||
({"mode": "frames", "frames": [["w"]]}, "exactly one action list"),
|
||||
({"mode": "state", "frames": [[], [], []]}, "mode='frames'"),
|
||||
],
|
||||
)
|
||||
def test_model_tick_action_payload_validation(data, message: str) -> None:
|
||||
with pytest.raises(ValueError, match=message):
|
||||
parse_lingbot_camera_action_frames(data, expected_frames=3)
|
||||
|
||||
|
||||
class _TickConsumer:
|
||||
def __init__(self) -> None:
|
||||
self.ticks: list[ARDiffusionTickRequest] = []
|
||||
|
||||
async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
|
||||
self.ticks.append(tick)
|
||||
return ARDiffusionChunkMetadata.from_tick(tick)
|
||||
|
||||
def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata:
|
||||
return output
|
||||
|
||||
|
||||
class _Lifecycle:
|
||||
async def reset_session(self, session_id: str) -> None:
|
||||
del session_id
|
||||
|
||||
async def close_session(self, session_id: str) -> None:
|
||||
del session_id
|
||||
|
||||
|
||||
def test_session_reducer_keeps_prompt_and_actions_atomic_and_resets() -> None:
|
||||
async def exercise() -> None:
|
||||
consumer = _TickConsumer()
|
||||
reducer = LingBotCameraControlReducer()
|
||||
session = ARDiffusionSession(
|
||||
"world-actions",
|
||||
tick_consumer=consumer,
|
||||
lifecycle=_Lifecycle(),
|
||||
max_pending_events=8,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}",
|
||||
control_reducer=reducer,
|
||||
)
|
||||
await session.accept_event(
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=1,
|
||||
prompt="turn toward the window",
|
||||
controls=(
|
||||
_action_control(
|
||||
"state",
|
||||
transitions=[{"actions": ["d"]}],
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
await session.next_chunk()
|
||||
|
||||
tick = consumer.ticks[0]
|
||||
assert tick.prompt == "turn toward the window"
|
||||
assert tick.applied_event_ids == (1,)
|
||||
assert tick.controls[0].data["frames"] == (("d",), ("d",), ("d",))
|
||||
|
||||
await session.reset()
|
||||
await session.accept_event(
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=2,
|
||||
controls=(
|
||||
_action_control(
|
||||
"state",
|
||||
transitions=[{"actions": []}],
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
await session.next_chunk()
|
||||
assert consumer.ticks[1].chunk_index == 0
|
||||
assert consumer.ticks[1].controls[0].data["frames"] == ((), (), ())
|
||||
|
||||
asyncio.run(exercise())
|
||||
@@ -1,543 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import math
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import ModuleType
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
_MODULE_PATH = Path(__file__).parents[4] / "vllm_omni/diffusion/models/lingbot_world/transformer.py"
|
||||
_MODULE_NAME = "_lingbot_world_attention_under_test"
|
||||
_MISSING_MODULE = object()
|
||||
_STUBBED_MODULE_NAMES = (
|
||||
"vllm",
|
||||
"vllm.distributed",
|
||||
"vllm.model_executor",
|
||||
"vllm.model_executor.layers",
|
||||
"vllm.model_executor.layers.conv",
|
||||
"vllm.model_executor.layers.linear",
|
||||
"vllm.model_executor.model_loader",
|
||||
"vllm.model_executor.model_loader.weight_utils",
|
||||
"vllm.model_executor.utils",
|
||||
"vllm_omni",
|
||||
"vllm_omni.diffusion",
|
||||
"vllm_omni.diffusion.attention",
|
||||
"vllm_omni.diffusion.attention.layer",
|
||||
"vllm_omni.diffusion.layers",
|
||||
"vllm_omni.diffusion.layers.norm",
|
||||
"vllm_omni.diffusion.layers.rope",
|
||||
"vllm_omni.experimental",
|
||||
"vllm_omni.experimental.ar_diffusion",
|
||||
"vllm_omni.experimental.ar_diffusion.kv_cache",
|
||||
"vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _inference_context():
|
||||
with torch.inference_mode():
|
||||
yield
|
||||
|
||||
|
||||
def _install_vllm_stubs() -> None:
|
||||
for name in _STUBBED_MODULE_NAMES:
|
||||
sys.modules[name] = ModuleType(name)
|
||||
|
||||
distributed = sys.modules["vllm.distributed"]
|
||||
distributed.get_tensor_model_parallel_rank = lambda: 0
|
||||
distributed.get_tensor_model_parallel_world_size = lambda: 1
|
||||
distributed.tensor_model_parallel_all_reduce = lambda value: value
|
||||
|
||||
def set_weight_attrs(weight: torch.Tensor, attrs: dict) -> None:
|
||||
for name, value in attrs.items():
|
||||
setattr(weight, name, value)
|
||||
|
||||
sys.modules["vllm.model_executor.utils"].set_weight_attrs = set_weight_attrs
|
||||
sys.modules["vllm.model_executor.model_loader.weight_utils"].default_weight_loader = lambda param, loaded_weight: (
|
||||
param.data.copy_(loaded_weight)
|
||||
)
|
||||
|
||||
class _Linear(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
*,
|
||||
bias: bool = True,
|
||||
return_bias: bool = False,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
del kwargs
|
||||
self.return_bias = return_bias
|
||||
self.weight = nn.Parameter(torch.empty(output_size, input_size))
|
||||
self.bias = nn.Parameter(torch.empty(output_size)) if bias else None
|
||||
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
||||
if self.bias is not None:
|
||||
nn.init.zeros_(self.bias)
|
||||
self.calls = 0
|
||||
|
||||
def forward(self, value: torch.Tensor):
|
||||
self.calls += 1
|
||||
output = F.linear(value, self.weight, self.bias)
|
||||
return (output, self.bias) if self.return_bias else output
|
||||
|
||||
linear = sys.modules["vllm.model_executor.layers.linear"]
|
||||
linear.ColumnParallelLinear = _Linear
|
||||
linear.RowParallelLinear = _Linear
|
||||
|
||||
class _QKVLinear(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
head_size: int,
|
||||
total_num_heads: int,
|
||||
total_num_kv_heads: int | None = None,
|
||||
*,
|
||||
bias: bool = True,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
del kwargs
|
||||
self.num_heads = total_num_heads
|
||||
self.num_kv_heads = total_num_kv_heads or total_num_heads
|
||||
output_size = (self.num_heads + 2 * self.num_kv_heads) * head_size
|
||||
self.weight = nn.Parameter(torch.empty(output_size, hidden_size))
|
||||
self.bias = nn.Parameter(torch.empty(output_size)) if bias else None
|
||||
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
||||
if self.bias is not None:
|
||||
nn.init.zeros_(self.bias)
|
||||
self.calls = 0
|
||||
|
||||
def load_shard(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str) -> None:
|
||||
shard_index = {"q": 0, "k": 1, "v": 2}[shard_id]
|
||||
shard_size = param.shape[0] // 3
|
||||
param.data.narrow(0, shard_index * shard_size, shard_size).copy_(loaded_weight)
|
||||
|
||||
self.weight.weight_loader = load_shard
|
||||
if self.bias is not None:
|
||||
self.bias.weight_loader = load_shard
|
||||
|
||||
def forward(self, value: torch.Tensor):
|
||||
self.calls += 1
|
||||
return F.linear(value, self.weight, self.bias), None
|
||||
|
||||
linear.QKVParallelLinear = _QKVLinear
|
||||
|
||||
class _Conv3dLayer(nn.Conv3d):
|
||||
def __init__(self, in_channels: int, out_channels: int, kernel_size, stride, **kwargs) -> None:
|
||||
super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride, **kwargs)
|
||||
|
||||
sys.modules["vllm.model_executor.layers.conv"].Conv3dLayer = _Conv3dLayer
|
||||
|
||||
class _LayerNorm(nn.LayerNorm):
|
||||
def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True) -> None:
|
||||
super().__init__(dim, eps=eps, elementwise_affine=elementwise_affine)
|
||||
|
||||
sys.modules["vllm_omni.diffusion.layers.norm"].LayerNorm = _LayerNorm
|
||||
|
||||
class _Attention(nn.Module):
|
||||
def __init__(self, *args, softmax_scale: float, **kwargs) -> None:
|
||||
super().__init__()
|
||||
del args
|
||||
self.softmax_scale = softmax_scale
|
||||
self.causal = kwargs["causal"]
|
||||
self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata=None,
|
||||
) -> torch.Tensor:
|
||||
del attn_metadata
|
||||
self.calls.append((query.detach().clone(), key.detach().clone(), value.detach().clone()))
|
||||
scores = torch.einsum("bqhd,bkhd->bhqk", query, key) * self.softmax_scale
|
||||
weights = scores.softmax(dim=-1)
|
||||
return torch.einsum("bhqk,bkhd->bqhd", weights, value)
|
||||
|
||||
sys.modules["vllm_omni.diffusion.attention.layer"].Attention = _Attention
|
||||
|
||||
class _RotaryEmbeddingWan(nn.Module):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__()
|
||||
del args, kwargs
|
||||
self.calls: list[tuple[torch.Tensor, torch.Tensor]] = []
|
||||
|
||||
def forward(self, value: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
|
||||
self.calls.append((cos.detach().clone(), sin.detach().clone()))
|
||||
return value
|
||||
|
||||
sys.modules["vllm_omni.diffusion.layers.rope"].RotaryEmbeddingWan = _RotaryEmbeddingWan
|
||||
|
||||
paged_attention = sys.modules["vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"]
|
||||
|
||||
class _PagedLayerContext:
|
||||
pass
|
||||
|
||||
class _PagedLayerInputs:
|
||||
pass
|
||||
|
||||
paged_attention.ARDiffusionPagedLayerContext = _PagedLayerContext
|
||||
paged_attention.ARDiffusionPagedLayerInputs = _PagedLayerInputs
|
||||
paged_attention.ar_diffusion_paged_attention = lambda *args, **kwargs: (_ for _ in ()).throw(
|
||||
AssertionError("CUDA paged attention was not expected in this CPU test")
|
||||
)
|
||||
paged_attention.paged_write_attn = lambda *args, **kwargs: (_ for _ in ()).throw(
|
||||
AssertionError("paged path was not expected in this direct-cache test")
|
||||
)
|
||||
|
||||
|
||||
def _load_module():
|
||||
assert _MODULE_PATH.exists(), "LingBot attention module has not been implemented"
|
||||
module_names = (*_STUBBED_MODULE_NAMES, _MODULE_NAME)
|
||||
previous = {name: sys.modules.get(name, _MISSING_MODULE) for name in module_names}
|
||||
try:
|
||||
_install_vllm_stubs()
|
||||
spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
finally:
|
||||
for name, original in previous.items():
|
||||
if original is _MISSING_MODULE:
|
||||
sys.modules.pop(name, None)
|
||||
else:
|
||||
sys.modules[name] = original
|
||||
|
||||
|
||||
def test_load_module_restores_process_module_state() -> None:
|
||||
module_names = (*_STUBBED_MODULE_NAMES, _MODULE_NAME)
|
||||
previous = {name: sys.modules.get(name, _MISSING_MODULE) for name in module_names}
|
||||
|
||||
_load_module()
|
||||
|
||||
for name, original in previous.items():
|
||||
assert sys.modules.get(name, _MISSING_MODULE) is original
|
||||
|
||||
|
||||
def test_allocate_lingbot_cache_creates_request_local_layer_storage() -> None:
|
||||
module = _load_module()
|
||||
|
||||
cache = module.allocate_lingbot_cache(
|
||||
batch_size=2,
|
||||
num_layers=3,
|
||||
max_tokens=7,
|
||||
num_local_heads=2,
|
||||
head_dim=4,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
assert len(cache.self_attention) == 3
|
||||
assert len(cache.cross_attention) == 3
|
||||
assert cache.cross_attention == [None, None, None]
|
||||
assert cache.self_attention[0].key.shape == (2, 7, 2, 4)
|
||||
assert cache.self_attention[0].value.shape == (2, 7, 2, 4)
|
||||
assert cache.self_attention[0].key.dtype == torch.float32
|
||||
assert cache.self_attention[0].end == 0
|
||||
assert cache.self_attention[0].key.data_ptr() != cache.self_attention[1].key.data_ptr()
|
||||
|
||||
|
||||
def _allocate_single_layer(module, *, max_tokens: int):
|
||||
return module.allocate_lingbot_cache(
|
||||
batch_size=1,
|
||||
num_layers=1,
|
||||
max_tokens=max_tokens,
|
||||
num_local_heads=1,
|
||||
head_dim=2,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def _set_identity_attention(attention: nn.Module) -> None:
|
||||
qkv = getattr(attention, "qkv", None)
|
||||
if qkv is not None:
|
||||
projections = (qkv, attention.o)
|
||||
else:
|
||||
projections = (attention.q, attention.k, attention.v, attention.o)
|
||||
for projection in projections:
|
||||
with torch.no_grad():
|
||||
identity = torch.eye(2)
|
||||
projection.weight.copy_(torch.cat((identity, identity, identity)) if projection is qkv else identity)
|
||||
if projection.bias is not None:
|
||||
projection.bias.zero_()
|
||||
attention.norm_q = nn.Identity()
|
||||
attention.norm_k = nn.Identity()
|
||||
|
||||
|
||||
def _tokens(*values: float) -> torch.Tensor:
|
||||
return torch.tensor([[[value, 0.0] for value in values]])
|
||||
|
||||
|
||||
def _record_inputs(module: nn.Module) -> list[tuple]:
|
||||
records: list[tuple] = []
|
||||
module.register_forward_pre_hook(lambda _module, args: records.append(args))
|
||||
return records
|
||||
|
||||
|
||||
def _record_outputs(module: nn.Module) -> list[torch.Tensor]:
|
||||
records: list[torch.Tensor] = []
|
||||
|
||||
def record(_module, _args, output) -> None:
|
||||
records.append(output[0] if isinstance(output, tuple) else output)
|
||||
|
||||
module.register_forward_hook(record)
|
||||
return records
|
||||
|
||||
|
||||
def _cache_snapshot(cache) -> tuple[torch.Tensor, torch.Tensor, tuple[int, int, int | None, int]]:
|
||||
return (
|
||||
cache.key.detach().view(torch.uint8).clone(),
|
||||
cache.value.detach().view(torch.uint8).clone(),
|
||||
(cache.end, cache.absolute_end, cache.last_start, cache.sink_end),
|
||||
)
|
||||
|
||||
|
||||
def _assert_cache_unchanged(cache, snapshot) -> None:
|
||||
key, value, metadata = snapshot
|
||||
assert torch.equal(cache.key.detach().view(torch.uint8), key)
|
||||
assert torch.equal(cache.value.detach().view(torch.uint8), value)
|
||||
assert (cache.end, cache.absolute_end, cache.last_start, cache.sink_end) == metadata
|
||||
|
||||
|
||||
def test_self_attention_repeated_offset_overwrites_then_later_offset_appends() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotSelfAttention(dim=2, num_heads=1)
|
||||
_set_identity_attention(attention)
|
||||
cache = _allocate_single_layer(module, max_tokens=6).self_attention[0]
|
||||
|
||||
attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
|
||||
attention(_tokens(10, 20), cache=cache, current_start=0, sink_tokens=0)
|
||||
|
||||
assert cache.end == 2
|
||||
assert cache.absolute_end == 2
|
||||
torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([10.0, 20.0]))
|
||||
|
||||
attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=0)
|
||||
|
||||
assert cache.end == 4
|
||||
assert cache.absolute_end == 4
|
||||
torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([10.0, 20.0, 3.0, 4.0]))
|
||||
|
||||
|
||||
def test_self_attention_uses_one_fused_qkv_projection() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotSelfAttention(dim=2, num_heads=1)
|
||||
cache = _allocate_single_layer(module, max_tokens=2).self_attention[0]
|
||||
|
||||
attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
|
||||
|
||||
assert attention.qkv.calls == 1
|
||||
|
||||
|
||||
def test_self_attention_routes_paged_inputs_through_engine_owned_cache() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotSelfAttention(dim=2, num_heads=1)
|
||||
_set_identity_attention(attention)
|
||||
paged_inputs = module.ARDiffusionPagedLayerInputs()
|
||||
calls = []
|
||||
|
||||
def paged_write(inputs, query, key, value, k_act, v_act, scale):
|
||||
calls.append((inputs, query.clone(), key.clone(), value.clone(), k_act, v_act, scale))
|
||||
return query
|
||||
|
||||
module.paged_write_attn = paged_write
|
||||
output = attention(
|
||||
_tokens(1, 2),
|
||||
cache=paged_inputs,
|
||||
current_start=0,
|
||||
sink_tokens=0,
|
||||
update_cache=False,
|
||||
)
|
||||
|
||||
assert output.shape == (1, 2, 2)
|
||||
assert len(calls) == 1
|
||||
assert calls[0][0] is paged_inputs
|
||||
assert calls[0][4:6] == (None, None)
|
||||
assert attention.attn.calls == []
|
||||
assert not any(hasattr(attention, name) for name in ("q", "k", "v"))
|
||||
|
||||
|
||||
def test_self_attention_is_chunk_causal_without_masking_inside_current_chunk() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotSelfAttention(dim=2, num_heads=1)
|
||||
_set_identity_attention(attention)
|
||||
cache = _allocate_single_layer(module, max_tokens=8).self_attention[0]
|
||||
attention_inputs = _record_inputs(attention.attn)
|
||||
|
||||
attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
|
||||
first_keys = attention_inputs[-1][1].detach().clone()
|
||||
attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=0)
|
||||
second_keys = attention_inputs[-1][1].detach().clone()
|
||||
|
||||
assert attention.attn.causal is False
|
||||
torch.testing.assert_close(first_keys[0, :, 0, 0], torch.tensor([1.0, 2.0]))
|
||||
torch.testing.assert_close(second_keys[0, :, 0, 0], torch.tensor([1.0, 2.0, 3.0, 4.0]))
|
||||
|
||||
|
||||
def test_self_attention_retains_sink_and_latest_local_history_after_eviction() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotSelfAttention(dim=2, num_heads=1)
|
||||
_set_identity_attention(attention)
|
||||
cache = _allocate_single_layer(module, max_tokens=4).self_attention[0]
|
||||
attention_inputs = _record_inputs(attention.attn)
|
||||
|
||||
attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=1)
|
||||
attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=1)
|
||||
attention(_tokens(5, 6), cache=cache, current_start=4, sink_tokens=1)
|
||||
|
||||
visible_keys = attention_inputs[-1][1].detach().clone()
|
||||
assert cache.end == 4
|
||||
assert cache.absolute_end == 6
|
||||
torch.testing.assert_close(visible_keys[0, :, 0, 0], torch.tensor([1.0, 4.0, 5.0, 6.0]))
|
||||
torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([1.0, 4.0, 5.0, 6.0]))
|
||||
|
||||
|
||||
def test_cross_attention_projects_encoder_kv_once_per_request() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotCrossAttention(dim=2, num_heads=1)
|
||||
_set_identity_attention(attention)
|
||||
key_outputs = _record_outputs(attention.k)
|
||||
value_outputs = _record_outputs(attention.v)
|
||||
attention_inputs = _record_inputs(attention.attn)
|
||||
|
||||
output, cache = attention(_tokens(1), _tokens(2, 3), cache=None)
|
||||
second_output, reused_cache = attention(_tokens(4), _tokens(20, 30), cache=cache)
|
||||
|
||||
assert output.shape == second_output.shape == (1, 1, 2)
|
||||
assert reused_cache is cache
|
||||
assert len(key_outputs) == 1
|
||||
assert len(value_outputs) == 1
|
||||
assert cache.end == 2
|
||||
torch.testing.assert_close(cache.key[0, :, 0, 0], torch.tensor([2.0, 3.0]))
|
||||
torch.testing.assert_close(attention_inputs[-1][1], cache.key)
|
||||
|
||||
|
||||
def test_self_attention_cache_is_isolated_between_requests() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotSelfAttention(dim=2, num_heads=1)
|
||||
_set_identity_attention(attention)
|
||||
first = _allocate_single_layer(module, max_tokens=4).self_attention[0]
|
||||
second = _allocate_single_layer(module, max_tokens=4).self_attention[0]
|
||||
|
||||
attention(_tokens(1, 2), cache=first, current_start=0, sink_tokens=0)
|
||||
|
||||
assert first.end == 2
|
||||
assert second.end == 0
|
||||
assert torch.count_nonzero(second.key) == 0
|
||||
attention(_tokens(8), cache=second, current_start=0, sink_tokens=0)
|
||||
torch.testing.assert_close(first.key[0, : first.end, 0, 0], torch.tensor([1.0, 2.0]))
|
||||
torch.testing.assert_close(second.key[0, : second.end, 0, 0], torch.tensor([8.0]))
|
||||
|
||||
|
||||
def test_cross_attention_keeps_checkpoint_qkvo_and_norm_parameter_names() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotCrossAttention(dim=4, num_heads=2)
|
||||
|
||||
names = set(attention.state_dict())
|
||||
|
||||
assert {
|
||||
"q.weight",
|
||||
"k.weight",
|
||||
"v.weight",
|
||||
"o.weight",
|
||||
"norm_q.weight",
|
||||
"norm_k.weight",
|
||||
} <= names
|
||||
assert callable(attention.norm_q.weight.weight_loader)
|
||||
|
||||
|
||||
def test_tp_world_size_one_attention_output_shapes() -> None:
|
||||
module = _load_module()
|
||||
self_attention = module.LingBotSelfAttention(dim=4, num_heads=2)
|
||||
cross_attention = module.LingBotCrossAttention(dim=4, num_heads=2)
|
||||
cache = module.allocate_lingbot_cache(
|
||||
batch_size=2,
|
||||
num_layers=1,
|
||||
max_tokens=5,
|
||||
num_local_heads=2,
|
||||
head_dim=2,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
hidden_states = torch.randn(2, 3, 4)
|
||||
encoder_hidden_states = torch.randn(2, 5, 4)
|
||||
|
||||
self_output = self_attention(hidden_states, cache=cache.self_attention[0], current_start=0, sink_tokens=0)
|
||||
cross_output, cache.cross_attention[0] = cross_attention(
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
cache=cache.cross_attention[0],
|
||||
)
|
||||
|
||||
assert self_output.shape == hidden_states.shape
|
||||
assert cross_output.shape == hidden_states.shape
|
||||
assert cache.cross_attention[0] is not None
|
||||
|
||||
|
||||
def test_self_attention_applies_rotary_embedding_to_current_query_and_key() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotSelfAttention(dim=2, num_heads=1)
|
||||
cache = _allocate_single_layer(module, max_tokens=2).self_attention[0]
|
||||
cos = torch.ones(2, 1)
|
||||
sin = torch.zeros(2, 1)
|
||||
rotary_inputs = _record_inputs(attention.rotary_embedding)
|
||||
|
||||
attention(_tokens(1, 2), cache=cache, current_start=0, rotary_emb=(cos, sin), sink_tokens=0)
|
||||
|
||||
assert len(rotary_inputs) == 2
|
||||
torch.testing.assert_close(rotary_inputs[0][1], cos)
|
||||
torch.testing.assert_close(rotary_inputs[1][2], sin)
|
||||
|
||||
|
||||
def test_attention_block_owns_checkpoint_parent_namespaces() -> None:
|
||||
module = _load_module()
|
||||
block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0")
|
||||
|
||||
names = set(block.state_dict())
|
||||
expected = {
|
||||
*(f"self_attn.{child}.weight" for child in ("qkv", "o", "norm_q", "norm_k")),
|
||||
*(f"cross_attn.{child}.weight" for child in ("q", "k", "v", "o", "norm_q", "norm_k")),
|
||||
}
|
||||
|
||||
assert expected <= names
|
||||
|
||||
|
||||
def test_cross_attention_caches_are_isolated_between_requests() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotCrossAttention(dim=2, num_heads=1)
|
||||
|
||||
_, first = attention(_tokens(1), _tokens(2, 3), cache=None)
|
||||
_, second = attention(_tokens(1), _tokens(8, 9), cache=None)
|
||||
|
||||
assert first.key.untyped_storage().data_ptr() != second.key.untyped_storage().data_ptr()
|
||||
assert first.value.untyped_storage().data_ptr() != second.value.untyped_storage().data_ptr()
|
||||
second_key = second.key.clone()
|
||||
first.key.zero_()
|
||||
torch.testing.assert_close(second.key, second_key)
|
||||
|
||||
|
||||
def test_tp_rmsnorm_weight_loader_selects_rank_shard(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
module = _load_module()
|
||||
monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
|
||||
monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 1)
|
||||
norm = module._LingBotRMSNorm(2, eps=1e-6)
|
||||
|
||||
norm.weight.weight_loader(norm.weight, torch.tensor([10.0, 20.0, 30.0, 40.0]))
|
||||
|
||||
torch.testing.assert_close(norm.weight, torch.tensor([30.0, 40.0]))
|
||||
@@ -1,525 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import math
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
_MODULE_PATH = Path(__file__).parents[4] / "vllm_omni/diffusion/models/lingbot_world/camera.py"
|
||||
_SPEC = importlib.util.spec_from_file_location("_lingbot_world_camera_under_test", _MODULE_PATH)
|
||||
assert _SPEC is not None and _SPEC.loader is not None
|
||||
_CAMERA_MODULE = importlib.util.module_from_spec(_SPEC)
|
||||
sys.modules[_SPEC.name] = _CAMERA_MODULE
|
||||
_SPEC.loader.exec_module(_CAMERA_MODULE)
|
||||
|
||||
CameraTrajectory = _CAMERA_MODULE.CameraTrajectory
|
||||
build_plucker_embedding = _CAMERA_MODULE.build_plucker_embedding
|
||||
interpolate_camera_trajectory = _CAMERA_MODULE.interpolate_camera_trajectory
|
||||
load_camera_trajectory = _CAMERA_MODULE.load_camera_trajectory
|
||||
resolve_trusted_action_directory = _CAMERA_MODULE.resolve_trusted_action_directory
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _identity_poses(num_frames: int) -> np.ndarray:
|
||||
return np.repeat(np.eye(4, dtype=np.float64)[None], num_frames, axis=0)
|
||||
|
||||
|
||||
def _write_trajectory(
|
||||
directory: Path,
|
||||
poses: np.ndarray,
|
||||
intrinsics: np.ndarray,
|
||||
) -> None:
|
||||
np.save(directory / "poses.npy", poses)
|
||||
np.save(directory / "intrinsics.npy", intrinsics)
|
||||
|
||||
|
||||
def _trusted(directory: Path):
|
||||
return resolve_trusted_action_directory(".", directory)
|
||||
|
||||
|
||||
def _trajectory(
|
||||
poses: torch.Tensor,
|
||||
intrinsics: torch.Tensor | None = None,
|
||||
) -> CameraTrajectory:
|
||||
if intrinsics is None:
|
||||
intrinsics = torch.tensor([[832.0 / 3.0, 160.0, 416.0, 240.0]]).repeat(poses.shape[0], 1)
|
||||
return CameraTrajectory(poses=poses, intrinsics=intrinsics)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("present_file", [None, "poses.npy"])
|
||||
def test_load_camera_trajectory_rejects_missing_files(tmp_path: Path, present_file: str | None) -> None:
|
||||
if present_file == "poses.npy":
|
||||
np.save(tmp_path / present_file, _identity_poses(1))
|
||||
|
||||
missing_file = "poses.npy" if present_file is None else "intrinsics.npy"
|
||||
with pytest.raises(FileNotFoundError, match=missing_file):
|
||||
load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("poses", "intrinsics", "message"),
|
||||
[
|
||||
(np.zeros((2, 3, 4)), np.zeros((2, 4)), "poses.npy"),
|
||||
(_identity_poses(2), np.zeros((2, 3)), "intrinsics.npy"),
|
||||
(_identity_poses(0), np.zeros((0, 4)), "1 and 4096 source frames"),
|
||||
(_identity_poses(2), np.zeros((1, 4)), "same number of frames"),
|
||||
],
|
||||
)
|
||||
def test_load_camera_trajectory_rejects_invalid_shapes(
|
||||
tmp_path: Path,
|
||||
poses: np.ndarray,
|
||||
intrinsics: np.ndarray,
|
||||
message: str,
|
||||
) -> None:
|
||||
_write_trajectory(tmp_path, poses, intrinsics)
|
||||
|
||||
with pytest.raises(ValueError, match=message):
|
||||
load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("array_name", ["poses", "intrinsics"])
|
||||
def test_load_camera_trajectory_rejects_non_finite_values(tmp_path: Path, array_name: str) -> None:
|
||||
poses = _identity_poses(1)
|
||||
intrinsics = np.ones((1, 4), dtype=np.float64)
|
||||
if array_name == "poses":
|
||||
poses[0, 0, 0] = np.nan
|
||||
else:
|
||||
intrinsics[0, 0] = np.inf
|
||||
_write_trajectory(tmp_path, poses, intrinsics)
|
||||
|
||||
with pytest.raises(ValueError, match="finite"):
|
||||
load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
|
||||
def test_load_camera_trajectory_preserves_raw_c2w_poses(tmp_path: Path) -> None:
|
||||
first_c2w = np.array(
|
||||
[
|
||||
[0.0, -1.0, 0.0, 2.0],
|
||||
[1.0, 0.0, 0.0, 3.0],
|
||||
[0.0, 0.0, 1.0, 4.0],
|
||||
[0.0, 0.0, 0.0, 1.0],
|
||||
]
|
||||
)
|
||||
second_c2w = np.array(
|
||||
[
|
||||
[0.0, 0.0, 1.0, 5.0],
|
||||
[0.0, 1.0, 0.0, 4.0],
|
||||
[-1.0, 0.0, 0.0, 6.0],
|
||||
[0.0, 0.0, 0.0, 1.0],
|
||||
]
|
||||
)
|
||||
poses = np.stack((first_c2w, second_c2w))
|
||||
intrinsics = np.array([[100.0, 110.0, 120.0, 130.0]] * 2)
|
||||
_write_trajectory(tmp_path, poses, intrinsics)
|
||||
|
||||
trajectory = load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
assert trajectory.poses.dtype == torch.float32
|
||||
assert trajectory.intrinsics.dtype == torch.float32
|
||||
torch.testing.assert_close(trajectory.poses, torch.from_numpy(poses).float())
|
||||
torch.testing.assert_close(trajectory.intrinsics, torch.from_numpy(intrinsics).float())
|
||||
|
||||
|
||||
def test_load_camera_trajectory_rejects_object_dtype_before_materializing(tmp_path: Path) -> None:
|
||||
poses = _identity_poses(1).astype(object)
|
||||
intrinsics = np.ones((1, 4), dtype=np.float32)
|
||||
_write_trajectory(tmp_path, poses, intrinsics)
|
||||
|
||||
with pytest.raises(ValueError, match="poses.npy.*real numeric"):
|
||||
load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
|
||||
def test_load_camera_trajectory_accepts_official_length_and_materializes_request_prefix(tmp_path: Path) -> None:
|
||||
_write_trajectory(
|
||||
tmp_path,
|
||||
_identity_poses(269),
|
||||
np.arange(269 * 4, dtype=np.float32).reshape(269, 4),
|
||||
)
|
||||
|
||||
trajectory = load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
assert trajectory.poses.shape == (117, 4, 4)
|
||||
assert trajectory.intrinsics.shape == (117, 4)
|
||||
np.testing.assert_array_equal(
|
||||
trajectory.intrinsics.numpy(),
|
||||
np.arange(269 * 4, dtype=np.float32).reshape(269, 4)[:117],
|
||||
)
|
||||
|
||||
|
||||
def test_load_camera_trajectory_rejects_unbounded_source_frame_count(tmp_path: Path) -> None:
|
||||
_write_trajectory(
|
||||
tmp_path,
|
||||
_identity_poses(4097),
|
||||
np.ones((4097, 4), dtype=np.float32),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="poses.npy.*bounded camera-array limit|4096 source frames"):
|
||||
load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
|
||||
def test_load_camera_trajectory_rejects_truncated_payload_before_loading(tmp_path: Path) -> None:
|
||||
_write_trajectory(tmp_path, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
|
||||
poses_path = tmp_path / "poses.npy"
|
||||
poses_path.write_bytes(poses_path.read_bytes()[:-1])
|
||||
|
||||
with pytest.raises(ValueError, match="poses.npy byte size"):
|
||||
load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
|
||||
def test_load_camera_trajectory_rejects_oversized_npy_header(tmp_path: Path) -> None:
|
||||
header_size = 20_000
|
||||
(tmp_path / "poses.npy").write_bytes(
|
||||
np.lib.format.magic(2, 0) + header_size.to_bytes(4, "little") + b" " * (header_size - 1) + b"\n"
|
||||
)
|
||||
np.save(tmp_path / "intrinsics.npy", np.ones((1, 4), dtype=np.float32))
|
||||
|
||||
with pytest.raises(ValueError, match="poses.npy.*oversized NPY header"):
|
||||
load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
|
||||
def test_load_camera_trajectory_rejects_symlink_swap_after_resolution(tmp_path: Path) -> None:
|
||||
trusted_root = tmp_path / "trusted"
|
||||
action_directory = trusted_root / "action"
|
||||
outside = tmp_path / "outside"
|
||||
action_directory.mkdir(parents=True)
|
||||
outside.mkdir()
|
||||
_write_trajectory(action_directory, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
|
||||
_write_trajectory(outside, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
|
||||
resolved = resolve_trusted_action_directory("action", trusted_root)
|
||||
action_directory.rename(trusted_root / "original-action")
|
||||
action_directory.symlink_to(outside, target_is_directory=True)
|
||||
|
||||
with pytest.raises(ValueError, match="securely open.*action directory"):
|
||||
load_camera_trajectory(resolved)
|
||||
|
||||
|
||||
def test_load_camera_trajectory_rejects_symlinked_camera_file(tmp_path: Path) -> None:
|
||||
trusted_root = tmp_path / "trusted"
|
||||
action_directory = trusted_root / "action"
|
||||
outside = tmp_path / "outside"
|
||||
action_directory.mkdir(parents=True)
|
||||
outside.mkdir()
|
||||
_write_trajectory(action_directory, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
|
||||
np.save(outside / "poses.npy", _identity_poses(1))
|
||||
resolved = resolve_trusted_action_directory("action", trusted_root)
|
||||
(action_directory / "poses.npy").unlink()
|
||||
(action_directory / "poses.npy").symlink_to(outside / "poses.npy")
|
||||
|
||||
with pytest.raises(ValueError, match="poses.npy.*opened securely"):
|
||||
load_camera_trajectory(resolved)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(__import__("os"), "mkfifo"), reason="requires POSIX FIFO support")
|
||||
def test_load_camera_trajectory_rejects_fifo_without_blocking(tmp_path: Path) -> None:
|
||||
import os
|
||||
|
||||
os.mkfifo(tmp_path / "poses.npy")
|
||||
np.save(tmp_path / "intrinsics.npy", np.ones((1, 4), dtype=np.float32))
|
||||
script = """
|
||||
import importlib.util
|
||||
import sys
|
||||
|
||||
spec = importlib.util.spec_from_file_location("_lingbot_camera_fifo_test", sys.argv[1])
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
action = module.resolve_trusted_action_directory(".", sys.argv[2])
|
||||
try:
|
||||
module.load_camera_trajectory(action)
|
||||
except ValueError as error:
|
||||
print(error)
|
||||
else:
|
||||
raise AssertionError("FIFO camera input was accepted")
|
||||
"""
|
||||
|
||||
completed = subprocess.run(
|
||||
[sys.executable, "-c", script, str(_MODULE_PATH), str(tmp_path)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=3,
|
||||
check=False,
|
||||
)
|
||||
|
||||
assert completed.returncode == 0, completed.stderr
|
||||
assert "poses.npy must be a regular NPY file" in completed.stdout
|
||||
|
||||
|
||||
def test_load_camera_trajectory_decodes_the_validated_file_snapshot(tmp_path: Path, monkeypatch) -> None:
|
||||
original_poses = _identity_poses(1)
|
||||
replacement_poses = original_poses.copy()
|
||||
replacement_poses[0, 0, 3] = 123.0
|
||||
_write_trajectory(tmp_path, original_poses, np.ones((1, 4), dtype=np.float32))
|
||||
poses_path = tmp_path / "poses.npy"
|
||||
original_load = _CAMERA_MODULE.np.load
|
||||
mutation_count = 0
|
||||
|
||||
def mutate_source_then_decode(source, *args, **kwargs):
|
||||
nonlocal mutation_count
|
||||
if mutation_count == 0:
|
||||
mutation_count += 1
|
||||
np.save(poses_path, replacement_poses)
|
||||
return original_load(source, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(_CAMERA_MODULE.np, "load", mutate_source_then_decode)
|
||||
|
||||
trajectory = load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
assert mutation_count == 1
|
||||
torch.testing.assert_close(trajectory.poses, torch.from_numpy(original_poses).float())
|
||||
|
||||
|
||||
def test_load_camera_trajectory_ignores_non_camera_event_files(tmp_path: Path) -> None:
|
||||
poses = _identity_poses(1)
|
||||
intrinsics = np.ones((1, 4), dtype=np.float32)
|
||||
_write_trajectory(tmp_path, poses, intrinsics)
|
||||
np.save(tmp_path / "wasd_events.npy", np.array([{"key": "w"}], dtype=object))
|
||||
|
||||
trajectory = load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
torch.testing.assert_close(trajectory.poses, torch.from_numpy(poses).float())
|
||||
|
||||
|
||||
def test_interpolate_camera_trajectory_uses_linear_and_spherical_interpolation() -> None:
|
||||
poses = torch.eye(4).repeat(2, 1, 1)
|
||||
poses[1, :3, :3] = torch.tensor(
|
||||
[
|
||||
[0.0, -1.0, 0.0],
|
||||
[1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 1.0],
|
||||
]
|
||||
)
|
||||
poses[1, :3, 3] = torch.tensor([2.0, 4.0, 6.0])
|
||||
intrinsics = torch.tensor([[10.0, 20.0, 30.0, 40.0], [20.0, 40.0, 60.0, 80.0]])
|
||||
trajectory = CameraTrajectory(poses=poses, intrinsics=intrinsics)
|
||||
|
||||
result = interpolate_camera_trajectory(trajectory, 3)
|
||||
|
||||
torch.testing.assert_close(result.poses[0], poses[0])
|
||||
torch.testing.assert_close(result.poses[-1], poses[-1])
|
||||
torch.testing.assert_close(result.poses[1, :3, 3], torch.tensor([1.0, 2.0, 3.0]))
|
||||
expected_mid_rotation = torch.tensor(
|
||||
[
|
||||
[math.sqrt(0.5), -math.sqrt(0.5), 0.0],
|
||||
[math.sqrt(0.5), math.sqrt(0.5), 0.0],
|
||||
[0.0, 0.0, 1.0],
|
||||
]
|
||||
)
|
||||
torch.testing.assert_close(result.poses[1, :3, :3], expected_mid_rotation)
|
||||
torch.testing.assert_close(result.intrinsics[1], torch.tensor([15.0, 30.0, 45.0, 60.0]))
|
||||
|
||||
|
||||
def test_build_plucker_embedding_has_forward_center_ray() -> None:
|
||||
trajectory = _trajectory(torch.eye(4).unsqueeze(0))
|
||||
|
||||
result = build_plucker_embedding(
|
||||
trajectory,
|
||||
height=3,
|
||||
width=3,
|
||||
target_height=3,
|
||||
target_width=3,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(result[0, :3, 1, 1], torch.zeros(3))
|
||||
torch.testing.assert_close(result[0, 3:, 1, 1], torch.tensor([0.0, 0.0, 1.0]))
|
||||
|
||||
|
||||
def test_build_plucker_embedding_scales_reference_intrinsics() -> None:
|
||||
trajectory = _trajectory(
|
||||
torch.eye(4).unsqueeze(0),
|
||||
torch.tensor([[416.0, 240.0, 416.0, 240.0]]),
|
||||
)
|
||||
|
||||
result = build_plucker_embedding(
|
||||
trajectory,
|
||||
height=4,
|
||||
width=8,
|
||||
target_height=4,
|
||||
target_width=8,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
expected_corner_direction = torch.tensor([-0.875, -0.75, 1.0])
|
||||
expected_corner_direction /= torch.linalg.vector_norm(expected_corner_direction)
|
||||
torch.testing.assert_close(result[0, 3:, 0, 0], expected_corner_direction)
|
||||
|
||||
|
||||
def test_build_plucker_embedding_samples_target_pixel_centers() -> None:
|
||||
trajectory = _trajectory(
|
||||
torch.eye(4).unsqueeze(0),
|
||||
torch.tensor([[416.0, 240.0, 416.0, 240.0]]),
|
||||
)
|
||||
|
||||
result = build_plucker_embedding(
|
||||
trajectory,
|
||||
height=6,
|
||||
width=10,
|
||||
target_height=2,
|
||||
target_width=5,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
# Target centers map through the source scale: (0.5 * 2, 0.5 * 3).
|
||||
expected_direction = torch.tensor([-0.8, -0.5, 1.0])
|
||||
expected_direction /= torch.linalg.vector_norm(expected_direction)
|
||||
torch.testing.assert_close(result[0, 3:, 0, 0], expected_direction)
|
||||
|
||||
|
||||
def test_build_plucker_embedding_uses_framewise_deltas_for_rotations_and_translations() -> None:
|
||||
poses = torch.eye(4).repeat(3, 1, 1)
|
||||
poses[0, :3, 3] = torch.tensor([5.0, 7.0, 11.0])
|
||||
poses[1, :3, :3] = torch.tensor(
|
||||
[
|
||||
[0.0, -1.0, 0.0],
|
||||
[1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 1.0],
|
||||
]
|
||||
)
|
||||
poses[1, :3, 3] = torch.tensor([5.0, 9.0, 11.0])
|
||||
poses[2, :3, :3] = torch.tensor(
|
||||
[
|
||||
[-1.0, 0.0, 0.0],
|
||||
[0.0, -1.0, 0.0],
|
||||
[0.0, 0.0, 1.0],
|
||||
]
|
||||
)
|
||||
poses[2, :3, 3] = torch.tensor([3.0, 9.0, 11.0])
|
||||
trajectory = _trajectory(
|
||||
poses,
|
||||
torch.tensor([[416.0, 480.0, 208.0, 240.0]]).repeat(3, 1),
|
||||
)
|
||||
|
||||
result = build_plucker_embedding(
|
||||
trajectory,
|
||||
height=1,
|
||||
width=2,
|
||||
target_height=1,
|
||||
target_width=2,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
expected_direction = torch.tensor([0.0, 1.0, 1.0]) / math.sqrt(2.0)
|
||||
torch.testing.assert_close(result[0, :3, 0, 1], torch.zeros(3))
|
||||
torch.testing.assert_close(result[1, :3, 0, 1], torch.tensor([0.0, 1.0, 0.0]))
|
||||
torch.testing.assert_close(result[2, :3, 0, 1], torch.tensor([0.0, 1.0, 0.0]))
|
||||
torch.testing.assert_close(result[1, 3:, 0, 1], expected_direction)
|
||||
torch.testing.assert_close(result[2, 3:, 0, 1], expected_direction)
|
||||
|
||||
|
||||
def test_build_plucker_embedding_normalizes_framewise_translation_by_max_norm() -> None:
|
||||
poses = torch.eye(4).repeat(3, 1, 1)
|
||||
poses[:, :3, 3] = torch.tensor(
|
||||
[
|
||||
[10.0, -2.0, 5.0],
|
||||
[13.0, 2.0, 5.0],
|
||||
[19.0, 10.0, 5.0],
|
||||
]
|
||||
)
|
||||
|
||||
result = build_plucker_embedding(
|
||||
_trajectory(poses),
|
||||
height=1,
|
||||
width=1,
|
||||
target_height=1,
|
||||
target_width=1,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
expected_origins = torch.tensor([[0.0, 0.0, 0.0], [0.3, 0.4, 0.0], [0.6, 0.8, 0.0]])
|
||||
torch.testing.assert_close(result[:, :3, 0, 0], expected_origins)
|
||||
|
||||
|
||||
def test_build_plucker_embedding_interpolates_raw_poses_before_framewise_conversion() -> None:
|
||||
poses = torch.eye(4).repeat(2, 1, 1)
|
||||
poses[0, :3, 3] = torch.tensor([20.0, 4.0, -2.0])
|
||||
poses[1, :3, 3] = torch.tensor([26.0, 12.0, -2.0])
|
||||
raw = _trajectory(poses)
|
||||
|
||||
interpolated = interpolate_camera_trajectory(raw, 3)
|
||||
result = build_plucker_embedding(
|
||||
interpolated,
|
||||
height=1,
|
||||
width=1,
|
||||
target_height=1,
|
||||
target_width=1,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
expected_delta = torch.tensor([0.6, 0.8, 0.0])
|
||||
torch.testing.assert_close(result[0, :3, 0, 0], torch.zeros(3))
|
||||
torch.testing.assert_close(result[1, :3, 0, 0], expected_delta)
|
||||
torch.testing.assert_close(result[2, :3, 0, 0], expected_delta)
|
||||
|
||||
|
||||
def test_build_plucker_embedding_uses_half_pixel_centers_and_origin_direction_channel_order() -> None:
|
||||
poses = torch.eye(4).repeat(2, 1, 1)
|
||||
poses[1, :3, 3] = torch.tensor([2.0, 0.0, 0.0])
|
||||
trajectory = _trajectory(
|
||||
poses,
|
||||
torch.tensor([[832.0, 480.0, 0.0, 0.0]]).repeat(2, 1),
|
||||
)
|
||||
|
||||
result = build_plucker_embedding(
|
||||
trajectory,
|
||||
height=1,
|
||||
width=2,
|
||||
target_height=1,
|
||||
target_width=2,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
expected_direction = torch.tensor([0.25, 0.5, 1.0])
|
||||
expected_direction /= torch.linalg.vector_norm(expected_direction)
|
||||
torch.testing.assert_close(result[1, :3, 0, 0], torch.tensor([1.0, 0.0, 0.0]))
|
||||
torch.testing.assert_close(result[1, 3:, 0, 0], expected_direction)
|
||||
|
||||
|
||||
def test_build_plucker_embedding_honors_shape_dtype_and_device() -> None:
|
||||
trajectory = _trajectory(torch.eye(4).repeat(2, 1, 1))
|
||||
|
||||
result = build_plucker_embedding(
|
||||
trajectory,
|
||||
height=12,
|
||||
width=20,
|
||||
target_height=3,
|
||||
target_width=5,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float64,
|
||||
)
|
||||
|
||||
assert result.shape == (2, 6, 3, 5)
|
||||
assert result.dtype == torch.float64
|
||||
assert result.device == torch.device("cpu")
|
||||
|
||||
|
||||
def test_build_plucker_embedding_is_deterministic() -> None:
|
||||
trajectory = _trajectory(torch.eye(4).repeat(2, 1, 1))
|
||||
kwargs = {
|
||||
"height": 12,
|
||||
"width": 20,
|
||||
"target_height": 3,
|
||||
"target_width": 5,
|
||||
"device": torch.device("cpu"),
|
||||
"dtype": torch.float32,
|
||||
}
|
||||
|
||||
first = build_plucker_embedding(trajectory, **kwargs)
|
||||
second = build_plucker_embedding(trajectory, **kwargs)
|
||||
|
||||
torch.testing.assert_close(first, second, rtol=0.0, atol=0.0)
|
||||
@@ -1,685 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import inspect
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.diffusion.models.lingbot_world import test_lingbot_world_attention as attention_tests
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_weight_index_names.json.fixture"
|
||||
_SHAPE_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_official_shapes.json.fixture"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _inference_context():
|
||||
with torch.inference_mode():
|
||||
yield
|
||||
|
||||
|
||||
def _tiny_model(
|
||||
module,
|
||||
*,
|
||||
num_layers: int = 2,
|
||||
num_frames_per_block: int = 1,
|
||||
sliding_window_num_frames: int = 3,
|
||||
):
|
||||
return module.CausalLingBotWorldTransformer3DModel(
|
||||
patch_size=(1, 2, 2),
|
||||
num_attention_heads=2,
|
||||
attention_head_dim=2,
|
||||
in_channels=36,
|
||||
out_channels=2,
|
||||
text_dim=6,
|
||||
freq_dim=4,
|
||||
ffn_dim=8,
|
||||
num_layers=num_layers,
|
||||
cross_attn_norm=True,
|
||||
eps=1e-6,
|
||||
rope_max_seq_len=16,
|
||||
sink_size=1,
|
||||
num_frames_per_block=num_frames_per_block,
|
||||
sliding_window_num_frames=sliding_window_num_frames,
|
||||
local_attn_size=-1,
|
||||
)
|
||||
|
||||
|
||||
def _cache(module, model, *, max_tokens: int | None = None):
|
||||
if max_tokens is None:
|
||||
window_frames = (
|
||||
model.config.local_attn_size
|
||||
if model.config.local_attn_size != -1
|
||||
else model.config.sliding_window_num_frames
|
||||
)
|
||||
max_tokens = window_frames * 2 * 2
|
||||
return module.allocate_lingbot_cache(
|
||||
batch_size=1,
|
||||
num_layers=len(model.blocks),
|
||||
max_tokens=max_tokens,
|
||||
num_local_heads=2,
|
||||
head_dim=2,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def _checkpoint_parameter_specs(model) -> dict[str, tuple[torch.Size, torch.dtype]]:
|
||||
specs: dict[str, tuple[torch.Size, torch.dtype]] = {}
|
||||
for name, parameter in model.named_parameters():
|
||||
marker = ".self_attn.qkv."
|
||||
if marker not in name:
|
||||
specs[name] = (parameter.shape, parameter.dtype)
|
||||
continue
|
||||
shard_shape = torch.Size((parameter.shape[0] // 3, *parameter.shape[1:]))
|
||||
for projection in ("q", "k", "v"):
|
||||
specs[name.replace(marker, f".self_attn.{projection}.")] = (shard_shape, parameter.dtype)
|
||||
return specs
|
||||
|
||||
|
||||
def _checkpoint_weights_for_model(model) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
|
||||
checkpoint_weights: dict[str, torch.Tensor] = {}
|
||||
expected_parameters: dict[str, torch.Tensor] = {}
|
||||
next_value = 1
|
||||
for name, parameter in model.named_parameters():
|
||||
marker = ".self_attn.qkv."
|
||||
if marker not in name:
|
||||
value = torch.full_like(parameter, next_value)
|
||||
checkpoint_weights[name] = value
|
||||
expected_parameters[name] = value
|
||||
next_value += 1
|
||||
continue
|
||||
shards = []
|
||||
shard_shape = (parameter.shape[0] // 3, *parameter.shape[1:])
|
||||
for projection in ("q", "k", "v"):
|
||||
value = torch.full(shard_shape, next_value, dtype=parameter.dtype, device=parameter.device)
|
||||
checkpoint_weights[name.replace(marker, f".self_attn.{projection}.")] = value
|
||||
shards.append(value)
|
||||
next_value += 1
|
||||
expected_parameters[name] = torch.cat(shards)
|
||||
return checkpoint_weights, expected_parameters
|
||||
|
||||
|
||||
def _self_cache_snapshot(cache) -> list[tuple[torch.Tensor, torch.Tensor, tuple[int, int, int | None, int]]]:
|
||||
return [attention_tests._cache_snapshot(layer_cache) for layer_cache in cache.self_attention]
|
||||
|
||||
|
||||
def _assert_self_cache_unchanged(cache, snapshot) -> None:
|
||||
for layer_cache, layer_snapshot in zip(cache.self_attention, snapshot, strict=True):
|
||||
attention_tests._assert_cache_unchanged(layer_cache, layer_snapshot)
|
||||
|
||||
|
||||
def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera_path() -> None:
|
||||
torch.manual_seed(7)
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module).eval()
|
||||
cache = _cache(module, model)
|
||||
encoder_hidden_states = torch.randn(1, 3, 6)
|
||||
camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
|
||||
cross_key_outputs = [attention_tests._record_outputs(block.cross_attn.k) for block in model.blocks]
|
||||
text_projection_outputs = attention_tests._record_outputs(model.text_embedding)
|
||||
outputs = []
|
||||
|
||||
assert model.patch_embedding.in_channels == 36
|
||||
assert model.patch_embedding_wancamctrl.in_features == 6 * 8 * 8 * 1 * 2 * 2
|
||||
|
||||
for start_frame in range(4):
|
||||
hidden_states = torch.randn(1, 36, 1, 4, 4)
|
||||
timestep = torch.tensor([float(start_frame + 1)])
|
||||
snapshot = _self_cache_snapshot(cache)
|
||||
|
||||
transient_output = model(
|
||||
hidden_states,
|
||||
timestep,
|
||||
encoder_hidden_states,
|
||||
camera_hidden_states,
|
||||
cache=cache,
|
||||
start_frame=start_frame,
|
||||
update_cache=False,
|
||||
)
|
||||
|
||||
assert transient_output.shape == (1, 2, 1, 4, 4)
|
||||
_assert_self_cache_unchanged(cache, snapshot)
|
||||
assert all(layer_cache is not None for layer_cache in cache.cross_attention)
|
||||
|
||||
committed_output = model(
|
||||
hidden_states,
|
||||
timestep,
|
||||
encoder_hidden_states,
|
||||
camera_hidden_states,
|
||||
cache=cache,
|
||||
start_frame=start_frame,
|
||||
update_cache=True,
|
||||
)
|
||||
torch.testing.assert_close(committed_output, transient_output)
|
||||
outputs.append(committed_output)
|
||||
|
||||
expected_token_end = (start_frame + 1) * 2 * 2
|
||||
assert all(layer_cache.absolute_end == expected_token_end for layer_cache in cache.self_attention)
|
||||
|
||||
assert len(outputs) == 4
|
||||
assert all(len(outputs) == 1 for outputs in cross_key_outputs)
|
||||
assert len(text_projection_outputs) == 1
|
||||
|
||||
alternate_cache = _cache(module, model)
|
||||
alternate_output = model(
|
||||
torch.zeros(1, 36, 1, 4, 4),
|
||||
torch.tensor([1.0]),
|
||||
encoder_hidden_states,
|
||||
torch.zeros_like(camera_hidden_states),
|
||||
cache=alternate_cache,
|
||||
start_frame=0,
|
||||
update_cache=False,
|
||||
)
|
||||
camera_output = model(
|
||||
torch.zeros(1, 36, 1, 4, 4),
|
||||
torch.tensor([1.0]),
|
||||
encoder_hidden_states,
|
||||
torch.ones_like(camera_hidden_states),
|
||||
cache=_cache(module, model),
|
||||
start_frame=0,
|
||||
update_cache=False,
|
||||
)
|
||||
assert not torch.equal(camera_output, alternate_output)
|
||||
|
||||
|
||||
def test_partial_cross_attention_cache_projects_text_only_for_missing_layers() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module, num_layers=2).eval()
|
||||
cache = _cache(module, model)
|
||||
hidden_states = torch.randn(1, 36, 1, 4, 4)
|
||||
timestep = torch.tensor([1.0])
|
||||
encoder_hidden_states = torch.randn(1, 3, 6)
|
||||
camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
|
||||
|
||||
model(
|
||||
hidden_states,
|
||||
timestep,
|
||||
encoder_hidden_states,
|
||||
camera_hidden_states,
|
||||
cache=cache,
|
||||
start_frame=0,
|
||||
update_cache=False,
|
||||
)
|
||||
cache.cross_attention[1] = None
|
||||
text_projection_outputs = attention_tests._record_outputs(model.text_embedding)
|
||||
cross_key_outputs = [attention_tests._record_outputs(block.cross_attn.k) for block in model.blocks]
|
||||
|
||||
model(
|
||||
hidden_states,
|
||||
timestep,
|
||||
encoder_hidden_states,
|
||||
camera_hidden_states,
|
||||
cache=cache,
|
||||
start_frame=0,
|
||||
update_cache=False,
|
||||
)
|
||||
|
||||
assert len(text_projection_outputs) == 1
|
||||
assert [len(outputs) for outputs in cross_key_outputs] == [0, 1]
|
||||
assert all(layer_cache is not None for layer_cache in cache.cross_attention)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("frames", [1, 6], ids=("partial", "multiple_blocks"))
|
||||
def test_forward_rejects_chunks_that_do_not_equal_configured_block_size(frames: int) -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(
|
||||
module,
|
||||
num_layers=1,
|
||||
num_frames_per_block=3,
|
||||
sliding_window_num_frames=6,
|
||||
)
|
||||
cache = _cache(module, model)
|
||||
snapshot = _self_cache_snapshot(cache)
|
||||
patch_inputs = attention_tests._record_inputs(model.patch_embedding)
|
||||
|
||||
with pytest.raises(ValueError, match="exactly 3 post-patch frames"):
|
||||
model(
|
||||
torch.randn(1, 36, frames, 4, 4),
|
||||
torch.tensor([1.0]),
|
||||
torch.randn(1, 3, 6),
|
||||
torch.randn(1, 6 * 8 * 8, frames, 4, 4),
|
||||
cache=cache,
|
||||
start_frame=0,
|
||||
update_cache=True,
|
||||
)
|
||||
|
||||
assert patch_inputs == []
|
||||
_assert_self_cache_unchanged(cache, snapshot)
|
||||
assert cache.cross_attention == [None]
|
||||
|
||||
|
||||
def test_forward_accepts_exactly_one_configured_frame_block() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(
|
||||
module,
|
||||
num_layers=1,
|
||||
num_frames_per_block=3,
|
||||
sliding_window_num_frames=6,
|
||||
)
|
||||
|
||||
output = model(
|
||||
torch.randn(1, 36, 3, 4, 4),
|
||||
torch.tensor([1.0]),
|
||||
torch.randn(1, 3, 6),
|
||||
torch.randn(1, 6 * 8 * 8, 3, 4, 4),
|
||||
cache=_cache(module, model),
|
||||
start_frame=0,
|
||||
update_cache=False,
|
||||
)
|
||||
|
||||
assert output.shape == (1, 2, 3, 4, 4)
|
||||
|
||||
|
||||
def test_transformer_allocates_request_cache_from_its_configured_geometry() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(
|
||||
module,
|
||||
num_layers=2,
|
||||
num_frames_per_block=3,
|
||||
sliding_window_num_frames=6,
|
||||
)
|
||||
|
||||
cache = model.allocate_cache(
|
||||
batch_size=2,
|
||||
latent_height=8,
|
||||
latent_width=12,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
assert len(cache.self_attention) == 2
|
||||
assert len(cache.cross_attention) == 2
|
||||
assert cache.cross_attention == [None, None]
|
||||
assert all(layer.key.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention)
|
||||
assert all(layer.value.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention)
|
||||
|
||||
|
||||
def test_video_patch_embedding_uses_temporal_height_width_token_order() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
|
||||
hidden_states = torch.zeros(1, 36, 2, 2, 4)
|
||||
hidden_states[0, 0] = torch.tensor(
|
||||
[
|
||||
[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]],
|
||||
[[9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]],
|
||||
]
|
||||
)
|
||||
with torch.no_grad():
|
||||
model.patch_embedding.weight.zero_()
|
||||
model.patch_embedding.bias.zero_()
|
||||
model.patch_embedding.weight[0, 0, 0, 0, 0] = 1
|
||||
|
||||
tokens = model.patch_embedding(hidden_states).flatten(2).transpose(1, 2)
|
||||
|
||||
torch.testing.assert_close(tokens[0, :, 0], torch.tensor([1.0, 3.0, 9.0, 11.0]))
|
||||
|
||||
|
||||
def test_unpatchify_restores_two_frame_channel_and_spatial_order() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
|
||||
hidden_states = (
|
||||
torch.stack(
|
||||
(
|
||||
torch.arange(0, 8),
|
||||
torch.arange(10, 18),
|
||||
torch.arange(20, 28),
|
||||
torch.arange(30, 38),
|
||||
)
|
||||
)
|
||||
.unsqueeze(0)
|
||||
.float()
|
||||
)
|
||||
expected = torch.tensor(
|
||||
[
|
||||
[
|
||||
[[0, 2, 10, 12], [4, 6, 14, 16]],
|
||||
[[20, 22, 30, 32], [24, 26, 34, 36]],
|
||||
],
|
||||
[
|
||||
[[1, 3, 11, 13], [5, 7, 15, 17]],
|
||||
[[21, 23, 31, 33], [25, 27, 35, 37]],
|
||||
],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
).unsqueeze(0)
|
||||
|
||||
output = model._unpatchify(
|
||||
hidden_states,
|
||||
batch_size=1,
|
||||
frames=2,
|
||||
height=1,
|
||||
width=2,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(output, expected)
|
||||
|
||||
|
||||
def test_head_modulation_broadcasts_distinct_condition_per_frame() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
|
||||
model.head.norm = torch.nn.Identity()
|
||||
with torch.no_grad():
|
||||
model.head.modulation.zero_()
|
||||
model.head.head.weight.zero_()
|
||||
model.head.head.bias.zero_()
|
||||
model.head.head.weight[:, 0] = 1
|
||||
hidden_states = torch.ones(1, 4, model.dim)
|
||||
timestep_embedding = torch.tensor([[[10.0, 0.0, 0.0, 0.0], [20.0, 0.0, 0.0, 0.0]]])
|
||||
expected = torch.tensor([21.0, 21.0, 41.0, 41.0]).view(1, 4, 1).expand(1, 4, 8)
|
||||
|
||||
output = model.head(hidden_states, timestep_embedding)
|
||||
|
||||
torch.testing.assert_close(output, expected)
|
||||
|
||||
|
||||
def test_constructor_defaults_match_official_checkpoint_config() -> None:
|
||||
module = attention_tests._load_module()
|
||||
parameters = inspect.signature(module.CausalLingBotWorldTransformer3DModel.__init__).parameters
|
||||
expected = {
|
||||
"in_channels": 36,
|
||||
"out_channels": 16,
|
||||
"num_layers": 40,
|
||||
"num_attention_heads": 40,
|
||||
"attention_head_dim": 128,
|
||||
"patch_size": (1, 2, 2),
|
||||
"text_dim": 4096,
|
||||
"ffn_dim": 13824,
|
||||
"freq_dim": 256,
|
||||
"rope_max_seq_len": 1024,
|
||||
"eps": 1e-6,
|
||||
"cross_attn_norm": True,
|
||||
"sink_size": 9,
|
||||
"num_frames_per_block": 3,
|
||||
"sliding_window_num_frames": 18,
|
||||
"local_attn_size": -1,
|
||||
}
|
||||
|
||||
assert {name: parameters[name].default for name in expected} == expected
|
||||
|
||||
|
||||
def test_transformer_exposes_parameter_dtype_for_pipeline_runtime() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module)
|
||||
|
||||
assert model.dtype == next(model.parameters()).dtype
|
||||
|
||||
|
||||
def test_transformer_declares_regional_compile_block() -> None:
|
||||
module = attention_tests._load_module()
|
||||
|
||||
assert module.CausalLingBotWorldTransformer3DModel._repeated_blocks == ["LingBotAttentionBlock"]
|
||||
|
||||
|
||||
def test_lingbot_rms_norm_uses_global_tp_square_mean(monkeypatch) -> None:
|
||||
module = attention_tests._load_module()
|
||||
reduced_values: list[torch.Tensor] = []
|
||||
|
||||
monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
|
||||
|
||||
def all_reduce(value: torch.Tensor) -> torch.Tensor:
|
||||
reduced_values.append(value.detach().clone())
|
||||
return value * 2
|
||||
|
||||
monkeypatch.setattr(module, "tensor_model_parallel_all_reduce", all_reduce)
|
||||
norm = module._LingBotRMSNorm(hidden_size=2, eps=0.0)
|
||||
value = torch.tensor([[3.0, 4.0]])
|
||||
|
||||
output = norm(value)
|
||||
|
||||
torch.testing.assert_close(output, value / math.sqrt(12.5))
|
||||
assert len(reduced_values) == 1
|
||||
torch.testing.assert_close(reduced_values[0], torch.tensor([[25.0]]))
|
||||
|
||||
|
||||
def test_attention_rejects_heads_not_divisible_by_tp_size(monkeypatch) -> None:
|
||||
module = attention_tests._load_module()
|
||||
monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 3)
|
||||
|
||||
with pytest.raises(ValueError, match="num_heads=2.*tp_size=3"):
|
||||
module.LingBotSelfAttention(dim=4, num_heads=2)
|
||||
|
||||
|
||||
def test_constructor_rejects_unsupported_qk_norm_with_config_error() -> None:
|
||||
module = attention_tests._load_module()
|
||||
|
||||
with pytest.raises(ValueError, match="qk_norm.*rms_norm_across_heads"):
|
||||
module.CausalLingBotWorldTransformer3DModel(qk_norm="layer_norm")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["image_dim", "added_kv_proj_dim", "pos_embed_seq_len"])
|
||||
def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None:
|
||||
module = attention_tests._load_module()
|
||||
|
||||
with pytest.raises(ValueError, match=field):
|
||||
module.CausalLingBotWorldTransformer3DModel(**{field: 4})
|
||||
|
||||
|
||||
def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> None:
|
||||
module = attention_tests._load_module()
|
||||
|
||||
with pytest.raises(RuntimeError, match="quant_config.*not supported"):
|
||||
module.CausalLingBotWorldTransformer3DModel(quant_config=object())
|
||||
|
||||
|
||||
def test_from_config_accepts_diffusers_metadata_and_normalizes_patch_size() -> None:
|
||||
module = attention_tests._load_module()
|
||||
|
||||
class ConfigProbe(module.CausalLingBotWorldTransformer3DModel):
|
||||
def __init__(self, **kwargs) -> None:
|
||||
self.received_kwargs = kwargs
|
||||
|
||||
config = {
|
||||
"_class_name": "CausalLingBotWorldTransformer3DModel",
|
||||
"_diffusers_version": "0.35.0.dev0",
|
||||
"patch_size": [1, 2, 2],
|
||||
"num_attention_heads": 40,
|
||||
"attention_head_dim": 128,
|
||||
"in_channels": 36,
|
||||
"out_channels": 16,
|
||||
"text_dim": 4096,
|
||||
"freq_dim": 256,
|
||||
"ffn_dim": 13824,
|
||||
"num_layers": 40,
|
||||
"cross_attn_norm": True,
|
||||
"eps": 1e-6,
|
||||
"image_dim": None,
|
||||
"added_kv_proj_dim": None,
|
||||
"rope_max_seq_len": 1024,
|
||||
"pos_embed_seq_len": None,
|
||||
"qk_norm": "rms_norm_across_heads",
|
||||
"sink_size": 9,
|
||||
"num_frames_per_block": 3,
|
||||
"sliding_window_num_frames": 18,
|
||||
"local_attn_size": -1,
|
||||
}
|
||||
|
||||
model = ConfigProbe.from_config(config, prefix="transformer")
|
||||
|
||||
assert model.received_kwargs["patch_size"] == (1, 2, 2)
|
||||
assert model.received_kwargs["num_layers"] == 40
|
||||
assert model.received_kwargs["prefix"] == "transformer"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value"),
|
||||
[
|
||||
("in_channels", 35),
|
||||
("num_layers", 39),
|
||||
("num_attention_heads", 32),
|
||||
("attention_head_dim", 64),
|
||||
("ffn_dim", 10240),
|
||||
("num_frames_per_block", 4),
|
||||
("sliding_window_num_frames", 24),
|
||||
("sink_size", 0),
|
||||
],
|
||||
)
|
||||
def test_from_config_rejects_checkpoint_topology_drift(field: str, value: object) -> None:
|
||||
module = attention_tests._load_module()
|
||||
config = {
|
||||
"_class_name": "CausalLingBotWorldTransformer3DModel",
|
||||
"_diffusers_version": "0.35.0.dev0",
|
||||
"patch_size": [1, 2, 2],
|
||||
"num_attention_heads": 40,
|
||||
"attention_head_dim": 128,
|
||||
"in_channels": 36,
|
||||
"out_channels": 16,
|
||||
"text_dim": 4096,
|
||||
"freq_dim": 256,
|
||||
"ffn_dim": 13824,
|
||||
"num_layers": 40,
|
||||
"cross_attn_norm": True,
|
||||
"eps": 1e-6,
|
||||
"image_dim": None,
|
||||
"added_kv_proj_dim": None,
|
||||
"rope_max_seq_len": 1024,
|
||||
"pos_embed_seq_len": None,
|
||||
"qk_norm": "rms_norm_across_heads",
|
||||
"sink_size": 9,
|
||||
"num_frames_per_block": 3,
|
||||
"sliding_window_num_frames": 18,
|
||||
"local_attn_size": -1,
|
||||
}
|
||||
config[field] = value
|
||||
|
||||
with pytest.raises(ValueError, match=field):
|
||||
module.CausalLingBotWorldTransformer3DModel.from_config(config)
|
||||
|
||||
|
||||
def test_from_config_ignores_non_semantic_checkpoint_metadata() -> None:
|
||||
module = attention_tests._load_module()
|
||||
config = {
|
||||
"_class_name": "CausalLingBotWorldTransformer3DModel",
|
||||
"architectures": ["CausalLingBotWorldTransformer3DModel"],
|
||||
"future_diffusers_metadata": {"version": 1},
|
||||
}
|
||||
|
||||
with torch.device("meta"):
|
||||
model = module.CausalLingBotWorldTransformer3DModel.from_config(config)
|
||||
|
||||
assert model.config.in_channels == 36
|
||||
assert model.config.num_layers == 40
|
||||
|
||||
|
||||
def test_load_weights_uses_parameter_loaders_and_rejects_unknown_model_keys() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module, num_layers=1)
|
||||
checkpoint_weights, expected_parameters = _checkpoint_weights_for_model(model)
|
||||
loader_calls: list[str] = []
|
||||
qkv_weight = model.blocks[0].self_attn.qkv.weight
|
||||
|
||||
def record_loader(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str) -> None:
|
||||
loader_calls.append(shard_id)
|
||||
shard_index = {"q": 0, "k": 1, "v": 2}[shard_id]
|
||||
shard_size = param.shape[0] // 3
|
||||
param.data.narrow(0, shard_index * shard_size, shard_size).copy_(loaded_weight)
|
||||
|
||||
qkv_weight.weight_loader = record_loader
|
||||
loaded = model.load_weights(iter(checkpoint_weights.items()))
|
||||
|
||||
assert set(checkpoint_weights) <= loaded
|
||||
assert set(dict(model.named_parameters())) <= loaded
|
||||
assert loader_calls == ["q", "k", "v"]
|
||||
for name, param in model.named_parameters():
|
||||
torch.testing.assert_close(param, expected_parameters[name])
|
||||
|
||||
first_name = next(iter(checkpoint_weights))
|
||||
partial = _tiny_model(module, num_layers=1)
|
||||
partial_loaded = partial.load_weights([(first_name, checkpoint_weights[first_name])])
|
||||
expected_partial = {first_name}
|
||||
for projection in ("q", "k", "v"):
|
||||
marker = f".self_attn.{projection}."
|
||||
if marker in first_name:
|
||||
expected_partial.add(first_name.replace(marker, ".self_attn.qkv."))
|
||||
break
|
||||
assert partial_loaded == expected_partial
|
||||
|
||||
with pytest.raises(KeyError, match="unexpected_model.weight"):
|
||||
model.load_weights([("unexpected_model.weight", torch.ones(1))])
|
||||
|
||||
|
||||
def test_load_weights_consumes_checkpoint_iterator_incrementally() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module, num_layers=1)
|
||||
first_name, first_param = next(iter(model.named_parameters()))
|
||||
loaded_weight = torch.full_like(first_param, 17)
|
||||
|
||||
def weights():
|
||||
yield first_name, loaded_weight
|
||||
torch.testing.assert_close(first_param, loaded_weight)
|
||||
|
||||
assert model.load_weights(weights()) == {first_name}
|
||||
|
||||
|
||||
def test_real_auto_weights_loader_covers_fused_model_parameter_namespace() -> None:
|
||||
utils = pytest.importorskip("vllm.model_executor.models.utils")
|
||||
module = attention_tests._load_module()
|
||||
transformer = _tiny_model(module, num_layers=1)
|
||||
checkpoint_weights, _ = _checkpoint_weights_for_model(transformer)
|
||||
|
||||
class Pipeline(torch.nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.transformer = transformer
|
||||
|
||||
pipeline = Pipeline()
|
||||
loaded = utils.AutoWeightsLoader(pipeline).load_weights(
|
||||
(f"transformer.{name}", value) for name, value in checkpoint_weights.items()
|
||||
)
|
||||
|
||||
assert set(dict(pipeline.named_parameters())) <= loaded
|
||||
|
||||
|
||||
def test_checkpoint_weight_index_fixture_matches_model_namespaces() -> None:
|
||||
module = attention_tests._load_module()
|
||||
fixture = json.loads(_FIXTURE_PATH.read_text())
|
||||
model = _tiny_model(module, num_layers=1)
|
||||
parameter_names = set(_checkpoint_parameter_specs(model))
|
||||
|
||||
assert {name.split(".", 1)[0] for name in parameter_names} == set(fixture["top_level_modules"])
|
||||
assert {name for name in parameter_names if not name.startswith("blocks.")} == set(fixture["non_block_parameters"])
|
||||
assert {name.removeprefix("blocks.0.") for name in parameter_names if name.startswith("blocks.0.")} == set(
|
||||
fixture["block_parameter_suffixes"]
|
||||
)
|
||||
expanded_names = set(fixture["non_block_parameters"])
|
||||
expanded_names.update(
|
||||
f"blocks.{layer}.{suffix}" for layer in range(40) for suffix in fixture["block_parameter_suffixes"]
|
||||
)
|
||||
canonical_names = "\n".join(sorted(expanded_names))
|
||||
|
||||
assert len(expanded_names) == fixture["index_parameter_count"] == 1421
|
||||
assert hashlib.sha256(canonical_names.encode()).hexdigest() == fixture["canonical_name_set_sha256"]
|
||||
assert fixture["canonical_name_to_shard_sha256"] == (
|
||||
"e1addc27d2b1fad6f10226a23a265ee3e3c61e74dce4d5859727bdf180a67992"
|
||||
)
|
||||
assert list(fixture["shards"]) == [
|
||||
f"diffusion_pytorch_model-{index:05}-of-00008.safetensors" for index in range(1, 9)
|
||||
]
|
||||
assert sum(shard["parameter_count"] for shard in fixture["shards"].values()) == 1421
|
||||
assert all(len(shard["parameter_names_sha256"]) == 64 for shard in fixture["shards"].values())
|
||||
|
||||
|
||||
def test_official_default_shapes_match_public_safetensors_header_fixture() -> None:
|
||||
module = attention_tests._load_module()
|
||||
fixture = json.loads(_SHAPE_FIXTURE_PATH.read_text())
|
||||
|
||||
with torch.device("meta"):
|
||||
model = module.CausalLingBotWorldTransformer3DModel()
|
||||
parameter_specs = _checkpoint_parameter_specs(model)
|
||||
expected_dtype = {"F32": torch.float32}[fixture["dtype"]]
|
||||
|
||||
assert len(fixture["shard_representatives"]) == 8
|
||||
assert set(fixture["shard_representatives"].values()) <= set(fixture["parameter_shapes"])
|
||||
|
||||
for name, expected_shape in fixture["parameter_shapes"].items():
|
||||
shape, dtype = parameter_specs[name]
|
||||
assert tuple(shape) == tuple(expected_shape), name
|
||||
assert dtype == expected_dtype
|
||||
@@ -1,39 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
_ROOT = Path(__file__).parents[4]
|
||||
_PRODUCTION_FILES = (
|
||||
_ROOT / "vllm_omni/diffusion/models/lingbot_world/camera.py",
|
||||
_ROOT / "vllm_omni/diffusion/models/lingbot_world/transformer.py",
|
||||
_ROOT / "vllm_omni/diffusion/models/lingbot_world/pipeline.py",
|
||||
)
|
||||
|
||||
|
||||
def test_all_lingbot_production_functions_have_complete_annotations() -> None:
|
||||
missing: list[str] = []
|
||||
for path in _PRODUCTION_FILES:
|
||||
tree = ast.parse(path.read_text())
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
continue
|
||||
arguments = (*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs)
|
||||
for argument in arguments:
|
||||
if argument.arg not in {"self", "cls"} and argument.annotation is None:
|
||||
missing.append(f"{path.name}:{node.lineno}:{node.name}:{argument.arg}")
|
||||
if node.args.vararg is not None and node.args.vararg.annotation is None:
|
||||
missing.append(f"{path.name}:{node.lineno}:{node.name}:*{node.args.vararg.arg}")
|
||||
if node.args.kwarg is not None and node.args.kwarg.annotation is None:
|
||||
missing.append(f"{path.name}:{node.lineno}:{node.name}:**{node.args.kwarg.arg}")
|
||||
if node.returns is None:
|
||||
missing.append(f"{path.name}:{node.lineno}:{node.name}:return")
|
||||
|
||||
assert missing == []
|
||||
File diff suppressed because it is too large.
Load diff
@@ -1,2 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -1,366 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""π0 LeRobot parity (in-process): the bit-for-bit correctness oracle.
|
||||
|
||||
Verifies that vllm-omni's ``Pi0ForActionPrediction`` produces bit-for-bit
|
||||
matching action chunks with LeRobot's ``PI0Policy`` when fed:
|
||||
* the same weights (``lerobot/pi0_base``)
|
||||
* the same pre-processed inputs (images, masks, tokens, state)
|
||||
* the same initial noise tensor
|
||||
|
||||
π0 is a flow-matching model (Euler-integrated ODE from t=1 → t=0 with a fixed
|
||||
``num_steps``), so the output is deterministic once the noise is fixed and
|
||||
``torch.allclose`` on the final action chunk is a valid oracle. This is the
|
||||
authoritative correctness oracle for the vllm-omni π0 port (max|Δ| < 1e-4).
|
||||
|
||||
The OpenPI websocket online-serving e2e lives separately in
|
||||
``tests/e2e/online_serving/test_pi0_expansion.py``.
|
||||
|
||||
Run in a SEPARATE ``lerobot[pi]`` venv (avoids dep conflict with the vllm-omni
|
||||
env), with the vllm-omni ``pi0`` package importable::
|
||||
|
||||
python -m pytest tests/diffusion/models/pi0/test_pi0_parity.py -v -s
|
||||
|
||||
Skipped automatically when LeRobot is not installed (e.g. the vllm-omni env).
|
||||
The run is CPU/float32 with fixed defaults; the only override is
|
||||
``PI0_PARITY_MODEL_PATH`` (a local pi0_base dir in LeRobot format, to skip the
|
||||
HF download; defaults to ``lerobot/pi0_base``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import importlib.util
|
||||
import os
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
# local_model: needs real weights + a lerobot venv (transformers 5.3.0), so it
|
||||
# runs locally rather than in the ready-CI (see docs/contributing/ci/CI_5levels.md).
|
||||
# Additionally gated on lerobot being importable.
|
||||
_HAS_LEROBOT = importlib.util.find_spec("lerobot") is not None
|
||||
|
||||
pytestmark = [pytest.mark.local_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
# ─── Config (fixed; matches LeRobot defaults for ``lerobot/pi0_base``) ──
|
||||
DEVICE = "cpu"
|
||||
DTYPE_STR = "float32"
|
||||
ATOL = 1e-4
|
||||
NUM_STEPS = 10
|
||||
BATCH_SIZE = 2
|
||||
ACTION_DIM = 32
|
||||
STATE_DIM = 32
|
||||
ACTION_HORIZON = 50
|
||||
MAX_TOKEN_LEN = 48
|
||||
|
||||
# The only knob: point at a local pi0_base dir (LeRobot format) to skip the HF
|
||||
# download; defaults to the HF repo id.
|
||||
MODEL_PATH = os.environ.get("PI0_PARITY_MODEL_PATH", "lerobot/pi0_base")
|
||||
|
||||
|
||||
def _resolve_checkpoint_dir() -> str:
|
||||
"""Return a local dir containing the pi0_base checkpoint (download if needed)."""
|
||||
if os.path.isdir(MODEL_PATH):
|
||||
return MODEL_PATH
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
return snapshot_download(repo_id=MODEL_PATH, repo_type="model")
|
||||
|
||||
|
||||
# ─── Dummy dataset stats (identity) ──────────────────────────────────
|
||||
def _dummy_dataset_stats() -> dict:
|
||||
return {
|
||||
"observation.state": {
|
||||
"mean": torch.zeros(STATE_DIM),
|
||||
"std": torch.ones(STATE_DIM),
|
||||
"q01": torch.zeros(STATE_DIM),
|
||||
"q99": torch.ones(STATE_DIM),
|
||||
},
|
||||
"action": {
|
||||
"mean": torch.zeros(ACTION_DIM),
|
||||
"std": torch.ones(ACTION_DIM),
|
||||
"q01": torch.zeros(ACTION_DIM),
|
||||
"q99": torch.ones(ACTION_DIM),
|
||||
},
|
||||
"images": {
|
||||
cam: {
|
||||
"mean": torch.zeros(3, 224, 224),
|
||||
"std": torch.ones(3, 224, 224),
|
||||
"q01": torch.zeros(3, 224, 224),
|
||||
"q99": torch.ones(3, 224, 224),
|
||||
}
|
||||
for cam in ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb")
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _create_dummy_batch(batch_size: int = BATCH_SIZE, device: str = DEVICE) -> dict:
|
||||
"""Reproducible dummy inputs — identical across both implementations."""
|
||||
g = torch.Generator(device="cpu").manual_seed(0)
|
||||
prompt = "Pick up the red block and place it in the bin"
|
||||
return {
|
||||
"observation.state": torch.randn(batch_size, STATE_DIM, generator=g, dtype=torch.float32).to(device),
|
||||
"action": torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device),
|
||||
"observation.images.base_0_rgb": torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to(
|
||||
device
|
||||
),
|
||||
"observation.images.left_wrist_0_rgb": torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to(
|
||||
device
|
||||
),
|
||||
"observation.images.right_wrist_0_rgb": torch.rand(
|
||||
batch_size, 3, 224, 224, generator=g, dtype=torch.float32
|
||||
).to(device),
|
||||
"task": [prompt for _ in range(batch_size)],
|
||||
}
|
||||
|
||||
|
||||
# ─── LeRobot instantiation ────────────────────────────────────────────
|
||||
def _instantiate_lerobot():
|
||||
from lerobot.policies.pi0 import PI0Policy
|
||||
from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors
|
||||
|
||||
policy = PI0Policy.from_pretrained(MODEL_PATH, strict=True)
|
||||
policy.to(DEVICE)
|
||||
policy.config.device = DEVICE
|
||||
policy.eval()
|
||||
|
||||
pre, post = make_pi0_pre_post_processors(config=policy.config, dataset_stats=_dummy_dataset_stats())
|
||||
return policy, pre, post
|
||||
|
||||
|
||||
# ─── vllm-omni instantiation ──────────────────────────────────────────
|
||||
def _instantiate_vllm_omni():
|
||||
"""Build the vllm-omni π0 model in isolation (no pipeline, no engine)."""
|
||||
from vllm_omni.diffusion.models.pi0 import Pi0Config, Pi0ForActionPrediction
|
||||
|
||||
cfg = Pi0Config(
|
||||
max_action_dim=ACTION_DIM,
|
||||
max_state_dim=STATE_DIM,
|
||||
chunk_size=ACTION_HORIZON,
|
||||
num_inference_steps=NUM_STEPS,
|
||||
dtype=DTYPE_STR,
|
||||
)
|
||||
model = Pi0ForActionPrediction(cfg)
|
||||
model.to(DEVICE).eval()
|
||||
_load_lerobot_weights(model)
|
||||
return model
|
||||
|
||||
|
||||
def _load_lerobot_weights(model):
|
||||
"""Feed the ``lerobot/pi0_base`` safetensors into the vllm-omni model.
|
||||
|
||||
``Pi0ForActionPrediction.load_weights`` handles the leading ``model.``
|
||||
prefix and the flat→nested / lm_head→embed_tokens remaps, so we can pass
|
||||
the raw checkpoint dict.
|
||||
"""
|
||||
import safetensors.torch
|
||||
|
||||
cache_dir = _resolve_checkpoint_dir()
|
||||
path = os.path.join(cache_dir, "model.safetensors")
|
||||
state = safetensors.torch.load_file(path)
|
||||
model.load_weights(list(state.items()))
|
||||
|
||||
|
||||
# ─── Helpers to extract LeRobot's pre-processed inputs ────────────────
|
||||
def _extract_lerobot_model_inputs(lerobot_policy, processed_batch):
|
||||
"""Mimic what ``PI0Policy.predict_action_chunk`` feeds into
|
||||
``self.model.sample_actions``. We use these *exact* tensors for vllm-omni
|
||||
so any divergence must come from the core model, not preprocessing.
|
||||
"""
|
||||
images, img_masks = lerobot_policy._preprocess_images(processed_batch)
|
||||
from lerobot.utils.constants import (
|
||||
OBS_LANGUAGE_ATTENTION_MASK,
|
||||
OBS_LANGUAGE_TOKENS,
|
||||
)
|
||||
|
||||
lang_tokens = processed_batch[OBS_LANGUAGE_TOKENS]
|
||||
lang_masks = processed_batch[OBS_LANGUAGE_ATTENTION_MASK]
|
||||
state = lerobot_policy.prepare_state(processed_batch)
|
||||
return images, img_masks, lang_tokens, lang_masks, state
|
||||
|
||||
|
||||
# ─── Shared fixed-noise sampler ───────────────────────────────────────
|
||||
def _make_fixed_noise(batch_size: int, device: str) -> torch.Tensor:
|
||||
g = torch.Generator(device="cpu").manual_seed(42)
|
||||
return torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device)
|
||||
|
||||
|
||||
# ─── Main test ────────────────────────────────────────────────────────
|
||||
@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
|
||||
def test_pi0_vllm_omni_vs_lerobot():
|
||||
print("\n[parity] Instantiating LeRobot…")
|
||||
lerobot_policy, lerobot_pre, _ = _instantiate_lerobot()
|
||||
|
||||
print("[parity] Instantiating vllm-omni…")
|
||||
omni_model = _instantiate_vllm_omni()
|
||||
|
||||
print("[parity] Preparing shared inputs…")
|
||||
raw_batch = _create_dummy_batch()
|
||||
processed_batch = lerobot_pre(copy.deepcopy(raw_batch))
|
||||
images, img_masks, lang_tokens, lang_masks, state = _extract_lerobot_model_inputs(lerobot_policy, processed_batch)
|
||||
noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE)
|
||||
|
||||
print(f"[parity] state.shape={state.shape} lang_tokens.shape={lang_tokens.shape}")
|
||||
print(f"[parity] images[0].shape={images[0].shape} (num_cams={len(images)})")
|
||||
print(f"[parity] noise.shape={noise.shape} dtype={noise.dtype}")
|
||||
|
||||
# ── LeRobot forward ──
|
||||
print("[parity] Running LeRobot sample_actions…")
|
||||
with torch.no_grad():
|
||||
lerobot_actions = lerobot_policy.model.sample_actions(
|
||||
images,
|
||||
img_masks,
|
||||
lang_tokens,
|
||||
lang_masks,
|
||||
state,
|
||||
noise=noise,
|
||||
num_steps=NUM_STEPS,
|
||||
)
|
||||
print(
|
||||
f"[parity] LeRobot actions: shape={lerobot_actions.shape} "
|
||||
f"mean={lerobot_actions.mean().item():.6f} std={lerobot_actions.std().item():.6f}"
|
||||
)
|
||||
|
||||
# ── vllm-omni forward ──
|
||||
print("[parity] Running vllm-omni sample_actions…")
|
||||
with torch.no_grad():
|
||||
omni_actions = omni_model.sample_actions(
|
||||
images=images,
|
||||
image_masks=img_masks,
|
||||
lang_tokens=lang_tokens,
|
||||
lang_masks=lang_masks,
|
||||
state=state,
|
||||
noise=noise,
|
||||
num_steps=NUM_STEPS,
|
||||
)
|
||||
print(
|
||||
f"[parity] vllm-omni actions: shape={omni_actions.shape} "
|
||||
f"mean={omni_actions.mean().item():.6f} std={omni_actions.std().item():.6f}"
|
||||
)
|
||||
|
||||
# ── Compare ──
|
||||
diff = (lerobot_actions.float() - omni_actions.float()).abs()
|
||||
max_diff = diff.max().item()
|
||||
mean_diff = diff.mean().item()
|
||||
print(f"[parity] |Δ| max={max_diff:.2e} mean={mean_diff:.2e} atol={ATOL:.1e}")
|
||||
close = torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=ATOL)
|
||||
print(f"[parity] torch.allclose(atol={ATOL}): {close}")
|
||||
|
||||
if not close:
|
||||
print("\n[parity] ⚠️ Outputs diverge — running per-stage diagnostics…")
|
||||
_diagnose_divergence(
|
||||
lerobot_policy.model,
|
||||
omni_model,
|
||||
images,
|
||||
img_masks,
|
||||
lang_tokens,
|
||||
lang_masks,
|
||||
state,
|
||||
noise,
|
||||
)
|
||||
|
||||
assert close, (
|
||||
f"vllm-omni vs LeRobot actions differ beyond atol={ATOL}. max_diff={max_diff:.2e} mean_diff={mean_diff:.2e}"
|
||||
)
|
||||
|
||||
|
||||
# ─── Per-stage divergence diagnostics ─────────────────────────────────
|
||||
@torch.no_grad()
|
||||
def _diagnose_divergence(
|
||||
lerobot_flow_model,
|
||||
omni_model,
|
||||
images,
|
||||
img_masks,
|
||||
lang_tokens,
|
||||
lang_masks,
|
||||
state,
|
||||
noise,
|
||||
):
|
||||
"""Localize a numerical mismatch to a specific pipeline stage:
|
||||
1. Prefix embeddings (SigLIP / embed_tokens / projector) — image vs lang.
|
||||
2. Prefix KV cache layer 0 — PaliGemma LM attention.
|
||||
3. A single denoise_step velocity at t=1.0 — action expert forward.
|
||||
"""
|
||||
from vllm_omni.diffusion.models.pi0.modeling_pi0 import (
|
||||
make_att_2d_masks,
|
||||
prepare_attention_masks_4d,
|
||||
)
|
||||
|
||||
# ── Stage 1: prefix embeddings ──
|
||||
lr_prefix_embs, lr_prefix_pad, lr_prefix_att = lerobot_flow_model.embed_prefix(
|
||||
images, img_masks, lang_tokens, lang_masks
|
||||
)
|
||||
sg_prefix_embs, sg_prefix_pad, sg_prefix_att = omni_model.embed_prefix(images, img_masks, lang_tokens, lang_masks)
|
||||
total_diff = (lr_prefix_embs.float() - sg_prefix_embs.float()).abs().max().item()
|
||||
print(f"[diag] prefix_embs max |Δ| = {total_diff:.2e} (shape={tuple(sg_prefix_embs.shape)})")
|
||||
print(f"[diag] prefix_pad_masks equal: {torch.equal(lr_prefix_pad, sg_prefix_pad)}")
|
||||
print(f"[diag] prefix_att_masks equal: {torch.equal(lr_prefix_att.bool(), sg_prefix_att.bool())}")
|
||||
|
||||
num_cams = len(images)
|
||||
img_len = 256 * num_cams
|
||||
lang_len = lr_prefix_embs.shape[1] - img_len
|
||||
img_diff = (lr_prefix_embs[:, :img_len].float() - sg_prefix_embs[:, :img_len].float()).abs().max().item()
|
||||
lang_diff = (lr_prefix_embs[:, img_len:].float() - sg_prefix_embs[:, img_len:].float()).abs().max().item()
|
||||
print(f"[diag] image slice [:{img_len}] max|Δ| = {img_diff:.2e} (num_cams={num_cams})")
|
||||
print(f"[diag] lang slice [{img_len}:] max|Δ| = {lang_diff:.2e} (len={lang_len})")
|
||||
|
||||
def _stats(name, t):
|
||||
t = t.float()
|
||||
print(
|
||||
f"[diag] {name} prefix_embs: mean={t.mean().item():+.4f} "
|
||||
f"std={t.std().item():.4f} min={t.min().item():+.2f} "
|
||||
f"max={t.max().item():+.2f}"
|
||||
)
|
||||
|
||||
_stats("LeRobot ", lr_prefix_embs)
|
||||
_stats("vllm-omni", sg_prefix_embs)
|
||||
|
||||
# ── Stage 2: prefix KV cache ──
|
||||
prefix_att_2d = make_att_2d_masks(sg_prefix_pad, sg_prefix_att)
|
||||
prefix_pos = torch.cumsum(sg_prefix_pad, dim=1) - 1
|
||||
prefix_att_4d = prepare_attention_masks_4d(prefix_att_2d)
|
||||
|
||||
lerobot_flow_model.paligemma_with_expert.paligemma.language_model.config._attn_implementation = "eager"
|
||||
_, lr_kv = lerobot_flow_model.paligemma_with_expert.forward(
|
||||
attention_mask=prefix_att_4d,
|
||||
position_ids=prefix_pos,
|
||||
past_key_values=None,
|
||||
inputs_embeds=[lr_prefix_embs, None],
|
||||
use_cache=True,
|
||||
)
|
||||
_, sg_kv = omni_model.paligemma_with_expert.forward(
|
||||
attention_mask=prefix_att_4d,
|
||||
position_ids=prefix_pos,
|
||||
past_key_values=None,
|
||||
inputs_embeds=[sg_prefix_embs, None],
|
||||
use_cache=True,
|
||||
)
|
||||
|
||||
def _layer0_kv(cache):
|
||||
if isinstance(cache, list):
|
||||
return cache[0]
|
||||
layer0 = cache.layers[0]
|
||||
return layer0.keys, layer0.values
|
||||
|
||||
try:
|
||||
lr_k, lr_v = _layer0_kv(lr_kv)
|
||||
sg_k, sg_v = _layer0_kv(sg_kv)
|
||||
dk = (lr_k.float() - sg_k.float()).abs().max().item()
|
||||
dv = (lr_v.float() - sg_v.float()).abs().max().item()
|
||||
print(f"[diag] prefix KV layer0 K max|Δ|={dk:.2e} V max|Δ|={dv:.2e}")
|
||||
except Exception as e: # noqa: BLE001
|
||||
print(f"[diag] could not extract prefix KV for comparison: {e}")
|
||||
|
||||
# ── Stage 3: a single denoise_step at t=1.0 ──
|
||||
bsize = state.shape[0]
|
||||
t = torch.ones(bsize, dtype=torch.float32, device=state.device)
|
||||
lr_vt = lerobot_flow_model.denoise_step(state, sg_prefix_pad, lr_kv, noise, t)
|
||||
sg_vt = omni_model.denoise_step(state, sg_prefix_pad, sg_kv, noise, t)
|
||||
print(f"[diag] denoise_step(t=1) v_t max|Δ| = {(lr_vt.float() - sg_vt.float()).abs().max().item():.2e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_pi0_vllm_omni_vs_lerobot()
|
||||
@@ -1,585 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""CPU unit tests for the π0 (Pi-Zero) VLA model.
|
||||
|
||||
Deliberately lightweight: pure utility code (config plumbing, mask building,
|
||||
sinusoidal embeddings, normalization, image preprocessing, weight-load remap)
|
||||
that runs on CPU with no weights and no lerobot. The full LeRobot parity check
|
||||
lives in ``tests/diffusion/models/pi0/test_pi0_parity.py`` and requires a model download.
|
||||
|
||||
pytest tests/diffusion/models/pi0/test_pi0_units.py -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import math
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.pi0.config import Pi0Config
|
||||
from vllm_omni.diffusion.models.pi0.modeling_pi0 import (
|
||||
OPENPI_ATTENTION_MASK_VALUE,
|
||||
Pi0ForActionPrediction,
|
||||
_apply_norm,
|
||||
_build_norm_buffers,
|
||||
create_sinusoidal_pos_embedding,
|
||||
get_gemma_config,
|
||||
make_att_2d_masks,
|
||||
prepare_attention_masks_4d,
|
||||
)
|
||||
from vllm_omni.diffusion.models.pi0.processor_pi0 import (
|
||||
Pi0ImageProcessor,
|
||||
build_model_inputs,
|
||||
pil_image_to_tensor,
|
||||
resize_with_pad,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Pi0Config dataclass + checkpoint resolver
|
||||
# ----------------------------------------------------------------------------
|
||||
_LEROBOT_CFG = {
|
||||
"type": "pi0",
|
||||
"paligemma_variant": "gemma_2b",
|
||||
"action_expert_variant": "gemma_300m",
|
||||
"chunk_size": 50,
|
||||
"max_action_dim": 32,
|
||||
"max_state_dim": 32,
|
||||
"num_inference_steps": 10,
|
||||
"image_resolution": [224, 224],
|
||||
"tokenizer_max_length": 48,
|
||||
"dtype": "float32",
|
||||
# Training-only keys that must be filtered out of the dataclass.
|
||||
"optimizer_lr": 2.5e-5,
|
||||
"scheduler_warmup_steps": 1000,
|
||||
"input_features": {
|
||||
"observation.images.base_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
|
||||
"observation.images.left_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
|
||||
"observation.images.right_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
|
||||
"observation.state": {"type": "STATE", "shape": [32]},
|
||||
},
|
||||
}
|
||||
|
||||
_EXPECTED_CAMERA_ORDER = [
|
||||
"observation.images.base_0_rgb",
|
||||
"observation.images.left_wrist_0_rgb",
|
||||
"observation.images.right_wrist_0_rgb",
|
||||
]
|
||||
|
||||
|
||||
def test_config_camera_order_from_input_features():
|
||||
c = Pi0Config.from_model_config(_LEROBOT_CFG)
|
||||
assert c.image_feature_keys == _EXPECTED_CAMERA_ORDER
|
||||
|
||||
|
||||
def test_config_filters_training_only_keys():
|
||||
c = Pi0Config.from_model_config(_LEROBOT_CFG)
|
||||
assert not hasattr(c, "optimizer_lr")
|
||||
assert not hasattr(c, "scheduler_warmup_steps")
|
||||
|
||||
|
||||
def test_config_coerces_image_resolution_to_tuple():
|
||||
c = Pi0Config.from_model_config(_LEROBOT_CFG)
|
||||
assert c.image_resolution == (224, 224)
|
||||
assert isinstance(c.image_resolution, tuple)
|
||||
|
||||
|
||||
def test_config_non_square_resolution_raises():
|
||||
with pytest.raises(ValueError):
|
||||
Pi0Config(image_resolution=(224, 256))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [(224,), (224, 224, 3), 224])
|
||||
def test_config_malformed_resolution_raises(bad):
|
||||
with pytest.raises(ValueError):
|
||||
Pi0Config(image_resolution=bad)
|
||||
|
||||
|
||||
def test_config_defaults_and_empty():
|
||||
c = Pi0Config.from_model_config(None)
|
||||
assert c.chunk_size == 50
|
||||
assert c.max_action_dim == 32
|
||||
assert c.max_state_dim == 32
|
||||
assert c.num_inference_steps == 10
|
||||
assert c.image_resolution == (224, 224)
|
||||
assert c.tokenizer_max_length == 48
|
||||
assert c.max_cameras == 3
|
||||
assert c.norm_stats is None
|
||||
assert c.image_key_map == {}
|
||||
|
||||
|
||||
def test_config_deploy_yaml_override():
|
||||
"""Deploy-yaml model_config is authoritative (e.g. force dtype/variant)."""
|
||||
c = Pi0Config.from_model_config({**_LEROBOT_CFG, "dtype": "bfloat16", "num_inference_steps": 4})
|
||||
assert c.dtype == "bfloat16"
|
||||
assert c.num_inference_steps == 4
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Kernel utilities: sinusoidal timestep embedding
|
||||
# ----------------------------------------------------------------------------
|
||||
def test_sinusoidal_embedding_shape():
|
||||
emb = create_sinusoidal_pos_embedding(torch.tensor([0.0, 0.5, 1.0]), dimension=64)
|
||||
assert emb.shape == (3, 64)
|
||||
|
||||
|
||||
def test_sinusoidal_embedding_values_at_zero():
|
||||
emb = create_sinusoidal_pos_embedding(torch.tensor([0.0]), dimension=4)
|
||||
assert torch.allclose(emb[0, :2].float(), torch.zeros(2), atol=1e-6)
|
||||
assert torch.allclose(emb[0, 2:].float(), torch.ones(2), atol=1e-6)
|
||||
|
||||
|
||||
def test_sinusoidal_embedding_odd_dim_raises():
|
||||
with pytest.raises(ValueError):
|
||||
create_sinusoidal_pos_embedding(torch.tensor([1.0]), dimension=3)
|
||||
|
||||
|
||||
def test_sinusoidal_embedding_non_1d_raises():
|
||||
with pytest.raises(ValueError):
|
||||
create_sinusoidal_pos_embedding(torch.tensor([[0.0, 0.5]]), dimension=4)
|
||||
|
||||
|
||||
def test_sinusoidal_embedding_distinct_timesteps_differ():
|
||||
emb = create_sinusoidal_pos_embedding(torch.tensor([0.1, 0.2, 0.5, 0.9]), dimension=16)
|
||||
for i in range(emb.shape[0]):
|
||||
for j in range(i + 1, emb.shape[0]):
|
||||
assert not torch.allclose(emb[i], emb[j])
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Kernel utilities: attention mask construction
|
||||
# ----------------------------------------------------------------------------
|
||||
def test_make_att_2d_masks_bidirectional():
|
||||
result = make_att_2d_masks(torch.ones(1, 4, dtype=torch.bool), torch.zeros(1, 4, dtype=torch.long))
|
||||
assert result.all()
|
||||
|
||||
|
||||
def test_make_att_2d_masks_causal():
|
||||
result = make_att_2d_masks(torch.ones(1, 4, dtype=torch.bool), torch.ones(1, 4, dtype=torch.long))
|
||||
expected = torch.tril(torch.ones(4, 4, dtype=torch.bool))
|
||||
assert torch.equal(result[0], expected)
|
||||
|
||||
|
||||
def test_make_att_2d_masks_prefix_lm():
|
||||
result = make_att_2d_masks(torch.ones(1, 5, dtype=torch.bool), torch.tensor([[0, 0, 0, 1, 1]], dtype=torch.long))
|
||||
assert result[0, 0, 0]
|
||||
assert not result[0, 0, 3]
|
||||
assert result[0, 3, 0]
|
||||
assert not result[0, 3, 4]
|
||||
|
||||
|
||||
def test_make_att_2d_masks_respects_padding():
|
||||
pad = torch.tensor([[1, 1, 1, 0]], dtype=torch.bool) # last token is padding
|
||||
att = torch.zeros(1, 4, dtype=torch.long) # full bidirectional
|
||||
result = make_att_2d_masks(pad, att)
|
||||
assert not result[0, :, 3].any(), "Nothing should attend to padded key"
|
||||
assert not result[0, 3, :].any(), "Padded query should attend to nothing"
|
||||
|
||||
|
||||
def test_prepare_attention_masks_4d():
|
||||
mask_2d = torch.tensor([[[True, False], [True, True]]])
|
||||
mask_4d = prepare_attention_masks_4d(mask_2d)
|
||||
assert mask_4d.shape == (1, 1, 2, 2)
|
||||
assert mask_4d[0, 0, 0, 0].item() == 0.0
|
||||
assert mask_4d[0, 0, 1, 1].item() == 0.0
|
||||
expected = torch.tensor(OPENPI_ATTENTION_MASK_VALUE, dtype=mask_4d.dtype)
|
||||
assert torch.equal(mask_4d[0, 0, 0, 1], expected)
|
||||
|
||||
|
||||
def test_attention_mask_value_finite_and_negative():
|
||||
assert math.isfinite(OPENPI_ATTENTION_MASK_VALUE)
|
||||
assert OPENPI_ATTENTION_MASK_VALUE < -1e30
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Normalization round-trips
|
||||
# ----------------------------------------------------------------------------
|
||||
def test_build_norm_buffers_missing():
|
||||
assert _build_norm_buffers(None, "state") is None
|
||||
assert _build_norm_buffers({}, "state") is None
|
||||
assert _build_norm_buffers({"state": None}, "state") is None
|
||||
|
||||
|
||||
def test_build_norm_buffers_mean_std():
|
||||
stats = {"state": {"mode": "mean_std", "mean": [1.0, 2.0], "std": [0.5, 0.5]}}
|
||||
out = _build_norm_buffers(stats, "state")
|
||||
assert out["mode"] == "mean_std"
|
||||
assert torch.equal(out["mean"], torch.tensor([1.0, 2.0]))
|
||||
assert torch.equal(out["std"], torch.tensor([0.5, 0.5]))
|
||||
|
||||
|
||||
def test_build_norm_buffers_min_max():
|
||||
stats = {"action": {"mode": "min_max", "min": [-1.0, 0.0], "max": [1.0, 2.0]}}
|
||||
out = _build_norm_buffers(stats, "action")
|
||||
assert out["mode"] == "min_max"
|
||||
assert torch.equal(out["min"], torch.tensor([-1.0, 0.0]))
|
||||
assert torch.equal(out["max"], torch.tensor([1.0, 2.0]))
|
||||
|
||||
|
||||
def test_apply_norm_mean_std_forward_inverse():
|
||||
stats = _build_norm_buffers({"s": {"mode": "mean_std", "mean": [0.5, -0.2], "std": [2.0, 0.1]}}, "s")
|
||||
x = torch.tensor([[1.0, 0.3]])
|
||||
z = _apply_norm(x, stats, inverse=False)
|
||||
assert torch.allclose(z, torch.tensor([[0.25, 5.0]]), atol=1e-5)
|
||||
x_back = _apply_norm(z, stats, inverse=True)
|
||||
assert torch.allclose(x_back, x, atol=1e-5)
|
||||
|
||||
|
||||
def test_apply_norm_min_max_forward_inverse():
|
||||
stats = _build_norm_buffers({"s": {"mode": "min_max", "min": [0.0, -1.0], "max": [1.0, 1.0]}}, "s")
|
||||
x = torch.tensor([[0.25, 0.5]])
|
||||
z = _apply_norm(x, stats, inverse=False)
|
||||
assert torch.allclose(z, torch.tensor([[-0.5, 0.5]]), atol=1e-5)
|
||||
x_back = _apply_norm(z, stats, inverse=True)
|
||||
assert torch.allclose(x_back, x, atol=1e-5)
|
||||
|
||||
|
||||
def test_apply_norm_preserves_padded_tail():
|
||||
stats = _build_norm_buffers({"s": {"mode": "mean_std", "mean": [0.0, 0.0], "std": [1.0, 1.0]}}, "s")
|
||||
x = torch.tensor([[1.0, 2.0, 99.0, -99.0]])
|
||||
z = _apply_norm(x, stats, inverse=False)
|
||||
assert torch.allclose(z[:, :2], x[:, :2])
|
||||
assert torch.allclose(z[:, 2:], x[:, 2:])
|
||||
|
||||
|
||||
def test_apply_norm_noop_when_stats_none():
|
||||
x = torch.randn(2, 5)
|
||||
assert torch.equal(_apply_norm(x, None, inverse=False), x)
|
||||
assert torch.equal(_apply_norm(x, None, inverse=True), x)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Gemma variant configs
|
||||
# ----------------------------------------------------------------------------
|
||||
def test_gemma_2b_dims():
|
||||
c = get_gemma_config("gemma_2b")
|
||||
assert c.width == 2048
|
||||
assert c.depth == 18
|
||||
|
||||
|
||||
def test_gemma_300m_dims():
|
||||
c = get_gemma_config("gemma_300m")
|
||||
assert c.width == 1024
|
||||
assert c.depth == 18
|
||||
|
||||
|
||||
def test_gemma_unknown_variant_raises():
|
||||
with pytest.raises(ValueError):
|
||||
get_gemma_config("gemma_7b")
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# load_weights remap rules (tiny model; needs torch + transformers)
|
||||
# ----------------------------------------------------------------------------
|
||||
def _tiny_pi0_model():
|
||||
"""Construct a tiny π0 model for remap tests (full Gemma depth but small
|
||||
action/state dims — the PaliGemma backbone is always full-size)."""
|
||||
cfg = Pi0Config(max_action_dim=8, max_state_dim=8, chunk_size=2)
|
||||
return Pi0ForActionPrediction(cfg)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _silence_pi0_loader(caplog):
|
||||
# The remap tests feed a single synthetic weight; the loader's
|
||||
# "missing params" warning is expected noise here.
|
||||
logging.getLogger("vllm_omni.diffusion.models.pi0.modeling_pi0").setLevel(logging.ERROR)
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_remap_paligemma_submodules_to_nested_layout():
|
||||
model = _tiny_pi0_model()
|
||||
params = dict(model.named_parameters())
|
||||
target = next(k for k in params if k.startswith("paligemma_with_expert.paligemma.model.vision_tower."))
|
||||
flat_key = "model." + target.replace(
|
||||
"paligemma_with_expert.paligemma.model.vision_tower.",
|
||||
"paligemma_with_expert.paligemma.vision_tower.",
|
||||
)
|
||||
payload = torch.full(params[target].shape, 0.12345)
|
||||
before = params[target].detach().clone()
|
||||
model.load_weights([(flat_key, payload)])
|
||||
after = dict(model.named_parameters())[target]
|
||||
assert not torch.equal(before, after), "flat→nested remap broken"
|
||||
assert torch.allclose(after, payload)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_lm_head_remap_to_embed_tokens():
|
||||
model = _tiny_pi0_model()
|
||||
params = dict(model.named_parameters())
|
||||
target = "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
|
||||
assert target in params
|
||||
payload = torch.full(params[target].shape, 0.5)
|
||||
before = params[target].detach().clone()
|
||||
model.load_weights([("model.paligemma_with_expert.paligemma.lm_head.weight", payload)])
|
||||
after = dict(model.named_parameters())[target]
|
||||
assert not torch.equal(before, after), "lm_head→embed_tokens remap broken"
|
||||
assert torch.allclose(after, payload)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_vision_tower_version_robust_remap():
|
||||
"""transformers >=5.4 flattens SigLIP from vision_tower.vision_model.* to
|
||||
vision_tower.*; the lerobot/pi0_base checkpoint uses the vision_model infix.
|
||||
load_weights must adapt so vision weights aren't silently dropped (which
|
||||
would leave the vision encoder at random init)."""
|
||||
model = _tiny_pi0_model()
|
||||
params = dict(model.named_parameters())
|
||||
# A representative vision-tower param as it exists on the model.
|
||||
target = next(
|
||||
k
|
||||
for k in params
|
||||
if k.startswith("paligemma_with_expert.paligemma.model.vision_tower.") and "patch_embedding.weight" in k
|
||||
)
|
||||
vt = "paligemma_with_expert.paligemma.model.vision_tower."
|
||||
rest = target[len(vt) :]
|
||||
# Build the checkpoint key in whichever layout differs from the model's.
|
||||
if rest.startswith("vision_model."):
|
||||
ckpt_key = vt + rest[len("vision_model.") :]
|
||||
else:
|
||||
ckpt_key = vt + "vision_model." + rest
|
||||
payload = torch.full(params[target].shape, 0.073)
|
||||
before = params[target].detach().clone()
|
||||
model.load_weights([("model." + ckpt_key, payload)])
|
||||
after = dict(model.named_parameters())[target]
|
||||
assert not torch.equal(before, after), "vision_tower version-robust remap broken"
|
||||
assert torch.allclose(after, payload)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Image processor
|
||||
# ----------------------------------------------------------------------------
|
||||
def test_resize_with_pad_square():
|
||||
out = resize_with_pad(torch.ones(1, 3, 224, 224), 224, 224)
|
||||
assert out.shape == (1, 3, 224, 224)
|
||||
|
||||
|
||||
def test_resize_with_pad_wide():
|
||||
"""Wide source (H < W): padding applied to top/bottom (rows)."""
|
||||
out = resize_with_pad(torch.zeros(1, 3, 100, 200), 224, 224)
|
||||
assert out.shape == (1, 3, 224, 224)
|
||||
assert out[0, 0, 0, 112].item() == -1.0 # top row is padded
|
||||
|
||||
|
||||
def test_resize_with_pad_tall():
|
||||
"""Tall source (H > W): padding applied to left/right (cols)."""
|
||||
out = resize_with_pad(torch.zeros(1, 3, 400, 100), 224, 224)
|
||||
assert out.shape == (1, 3, 224, 224)
|
||||
assert out[0, 0, 112, 0].item() == -1.0 # left-most column padded
|
||||
|
||||
|
||||
def test_resize_with_pad_rejects_non_4d():
|
||||
with pytest.raises(ValueError):
|
||||
resize_with_pad(torch.zeros(3, 100, 200), 224, 224)
|
||||
|
||||
|
||||
def test_pil_to_tensor():
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
img = Image.fromarray(np.full((10, 10, 3), 128, dtype=np.uint8))
|
||||
t = pil_image_to_tensor(img)
|
||||
assert t.shape == (1, 3, 10, 10)
|
||||
expected = 128.0 / 255.0 * 2.0 - 1.0
|
||||
assert abs(t[0, 0, 5, 5].item() - expected) < 1e-4
|
||||
|
||||
|
||||
def test_pil_to_tensor_rgba_to_rgb():
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
rgba = np.full((8, 8, 4), 200, dtype=np.uint8)
|
||||
img = Image.fromarray(rgba, mode="RGBA")
|
||||
t = pil_image_to_tensor(img)
|
||||
assert t.shape == (1, 3, 8, 8)
|
||||
|
||||
|
||||
def test_make_empty_image():
|
||||
proc = Pi0ImageProcessor(image_size=224)
|
||||
empty = proc.make_empty_image()
|
||||
assert empty.shape == (1, 3, 224, 224)
|
||||
assert (empty == -1.0).all()
|
||||
|
||||
|
||||
def test_preprocess_single_resizes_ndarray():
|
||||
import numpy as np
|
||||
|
||||
proc = Pi0ImageProcessor(image_size=224)
|
||||
arr = np.full((120, 120, 3), 100, dtype=np.uint8)
|
||||
t = proc.preprocess_single(arr)
|
||||
assert t.shape == (1, 3, 224, 224)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# build_model_inputs: camera ordering, padding, masks, state
|
||||
# ----------------------------------------------------------------------------
|
||||
class _FakeTokenizer:
|
||||
"""Minimal stand-in for the PaliGemma tokenizer (no download)."""
|
||||
|
||||
def __call__(
|
||||
self, text, padding=None, max_length=None, truncation=None, add_special_tokens=None, return_tensors=None
|
||||
):
|
||||
ids = [2] + [ord(c) % 100 for c in text][: max_length - 1]
|
||||
ids = ids[:max_length]
|
||||
attn = [1] * len(ids)
|
||||
while len(ids) < max_length:
|
||||
ids.append(0)
|
||||
attn.append(0)
|
||||
return {"input_ids": ids, "attention_mask": attn}
|
||||
|
||||
|
||||
def test_build_model_inputs_camera_order_and_padding():
|
||||
import numpy as np
|
||||
|
||||
config = Pi0Config.from_model_config(_LEROBOT_CFG) # 3-camera config
|
||||
device = torch.device("cpu")
|
||||
# Provide only 2 of the 3 declared cameras.
|
||||
robot_obs = {
|
||||
"images": {
|
||||
"observation.images.base_0_rgb": np.full((224, 224, 3), 50, dtype=np.uint8),
|
||||
"observation.images.left_wrist_0_rgb": np.full((224, 224, 3), 80, dtype=np.uint8),
|
||||
},
|
||||
"state": [0.1] * 14, # raw state shorter than max_state_dim=32
|
||||
"prompt": "pick up the block",
|
||||
}
|
||||
images, image_masks, lang_tokens, lang_masks, state = build_model_inputs(
|
||||
robot_obs, config, _FakeTokenizer(), device
|
||||
)
|
||||
assert len(images) == 3
|
||||
assert all(img.shape == (1, 3, 224, 224) for img in images)
|
||||
masks = [bool(m.item()) for m in image_masks]
|
||||
assert masks == [True, True, False]
|
||||
# The padded (3rd) camera must be all -1.
|
||||
assert (images[2] == -1.0).all()
|
||||
assert lang_tokens.shape == (1, 48)
|
||||
assert lang_masks.shape == (1, 48)
|
||||
assert state.shape == (1, 32)
|
||||
# State padded with zeros beyond the 14 real entries.
|
||||
assert torch.allclose(state[0, 14:], torch.zeros(18))
|
||||
|
||||
|
||||
def test_build_model_inputs_state_truncated():
|
||||
config = Pi0Config.from_model_config(_LEROBOT_CFG)
|
||||
robot_obs = {"images": {}, "state": [0.5] * 40, "prompt": "x"}
|
||||
_, _, _, _, state = build_model_inputs(robot_obs, config, _FakeTokenizer(), torch.device("cpu"))
|
||||
assert state.shape == (1, 32)
|
||||
|
||||
|
||||
def test_build_model_inputs_image_key_map():
|
||||
import numpy as np
|
||||
|
||||
cfg = Pi0Config.from_model_config({**_LEROBOT_CFG, "image_key_map": {"cam_high": "observation.images.base_0_rgb"}})
|
||||
robot_obs = {
|
||||
"images": {"cam_high": np.full((224, 224, 3), 60, dtype=np.uint8)},
|
||||
"prompt": "go",
|
||||
}
|
||||
_, image_masks, _, _, _ = build_model_inputs(robot_obs, cfg, _FakeTokenizer(), torch.device("cpu"))
|
||||
masks = [bool(m.item()) for m in image_masks]
|
||||
assert masks == [True, False, False] # only base_0_rgb (mapped) is present
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Version-stability / self-consistency (gated on a real pi0_base checkpoint)
|
||||
#
|
||||
# The LeRobot parity gate (test_pi0_parity.py) only validates correctness at
|
||||
# transformers 5.3.0 (where LeRobot loads pi0_base cleanly). Production vllm-omni
|
||||
# ships transformers 5.6.2, and LeRobot is not a usable reference there (it
|
||||
# incorrectly loads the SigLIP vision tower at >=5.4). So we pin the kernel's deterministic
|
||||
# output to golden values established at 5.3.0 and re-check them here: running on
|
||||
# the production transformers proves the port is version-stable and guards the
|
||||
# vision_tower / embed_scale version-robust paths against regressions.
|
||||
#
|
||||
# PI0_SELFCONSIST_CKPT=/path/to/pi0_base \
|
||||
# pytest tests/diffusion/models/pi0/test_pi0_units.py -k version_stable -v -s
|
||||
# ----------------------------------------------------------------------------
|
||||
import os # noqa: E402
|
||||
|
||||
_SELFCONSIST_CKPT = os.environ.get("PI0_SELFCONSIST_CKPT")
|
||||
|
||||
# Golden values established at transformers 5.3.0 (where parity = 7.15e-07).
|
||||
_GOLDEN_SUM = -67.808113
|
||||
_GOLDEN_STD = 0.252789
|
||||
_GOLDEN_MEAN = -0.021190
|
||||
_GOLDEN_EMBED_IMAGE_STD = 2.17603
|
||||
|
||||
|
||||
def _selfconsist_inputs(device):
|
||||
# NB: draw order is load-bearing — the golden values were established with
|
||||
# state drawn first, then the 3 images, then lang (shared RNG, seed 0).
|
||||
g = torch.Generator(device="cpu").manual_seed(0)
|
||||
b = 2
|
||||
state = torch.randn(b, 32, generator=g)
|
||||
images = [torch.rand(b, 3, 224, 224, generator=g) for _ in range(3)]
|
||||
masks = [torch.ones(b, dtype=torch.bool) for _ in range(3)]
|
||||
lang = torch.randint(0, 1000, (b, 48), generator=g)
|
||||
lmask = torch.ones(b, 48, dtype=torch.bool)
|
||||
ng = torch.Generator(device="cpu").manual_seed(42)
|
||||
noise = torch.randn(b, 50, 32, generator=ng)
|
||||
return (
|
||||
[im.to(device) for im in images],
|
||||
[m.to(device) for m in masks],
|
||||
lang.to(device),
|
||||
lmask.to(device),
|
||||
state.to(device),
|
||||
noise.to(device),
|
||||
)
|
||||
|
||||
|
||||
def _selfconsist_model(device):
|
||||
import safetensors.torch
|
||||
|
||||
model = Pi0ForActionPrediction(
|
||||
Pi0Config(max_action_dim=32, max_state_dim=32, chunk_size=50, num_inference_steps=10, dtype="float32")
|
||||
)
|
||||
state = safetensors.torch.load_file(os.path.join(_SELFCONSIST_CKPT, "model.safetensors"))
|
||||
filled = model.load_weights(list(state.items()))
|
||||
expected = {n for n, _ in model.named_parameters() if "rotary_emb" not in n and not n.endswith(".inv_freq")}
|
||||
missing = expected - (filled or set())
|
||||
assert not missing, f"{len(missing)} params received no weight (e.g. {sorted(missing)[:3]})"
|
||||
return model.to(device).eval()
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.skipif(not _SELFCONSIST_CKPT, reason="Set PI0_SELFCONSIST_CKPT to a pi0_base dir.")
|
||||
@torch.no_grad()
|
||||
def test_version_stable_actions():
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
model = _selfconsist_model(device)
|
||||
images, masks, lang, lmask, state, noise = _selfconsist_inputs(device)
|
||||
actions = (
|
||||
model.sample_actions(
|
||||
images=images,
|
||||
image_masks=masks,
|
||||
lang_tokens=lang,
|
||||
lang_masks=lmask,
|
||||
state=state,
|
||||
noise=noise,
|
||||
num_steps=10,
|
||||
)
|
||||
.float()
|
||||
.cpu()
|
||||
)
|
||||
s, m, std = actions.sum().item(), actions.mean().item(), actions.std().item()
|
||||
assert abs(s - _GOLDEN_SUM) < 0.5, f"action sum drifted: {s} vs {_GOLDEN_SUM}"
|
||||
assert abs(std - _GOLDEN_STD) < 0.01, f"action std drifted: {std} vs {_GOLDEN_STD}"
|
||||
assert abs(m - _GOLDEN_MEAN) < 0.01, f"action mean drifted: {m} vs {_GOLDEN_MEAN}"
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.skipif(not _SELFCONSIST_CKPT, reason="Set PI0_SELFCONSIST_CKPT to a pi0_base dir.")
|
||||
@torch.no_grad()
|
||||
def test_version_stable_embed_image():
|
||||
"""Directly guards the SigLIP vision-tower load across the >=5.4 flattening."""
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
model = _selfconsist_model(device)
|
||||
g = torch.Generator(device="cpu").manual_seed(7)
|
||||
img = torch.rand(1, 3, 224, 224, generator=g, dtype=torch.float32).to(device)
|
||||
std = model.paligemma_with_expert.embed_image(img).float().std().item()
|
||||
assert abs(std - _GOLDEN_EMBED_IMAGE_STD) < 0.05, (
|
||||
f"embed_image std drifted: {std} vs {_GOLDEN_EMBED_IMAGE_STD} — "
|
||||
"likely a vision_tower weight-load regression under this transformers version."
|
||||
)
|
||||
@@ -418,77 +418,3 @@ class TestShutdownCleansUpFutures:
|
||||
|
||||
assert len(executor._rpc_futures) == 0
|
||||
assert len(executor._output_futures) == 0
|
||||
|
||||
|
||||
class _RacyFuture(concurrent.futures.Future):
|
||||
# done() always lies and reports False, to deterministically force the
|
||||
# pump's check-then-act race window without needing real thread timing.
|
||||
def done(self) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def _racy_cancelled_future() -> concurrent.futures.Future:
|
||||
fut = _RacyFuture()
|
||||
fut.cancel()
|
||||
assert fut.cancelled()
|
||||
return fut
|
||||
|
||||
|
||||
class TestResultPumpCancelledFutureRace:
|
||||
"""Regression for #5793: a future cancelled concurrently with the pump's
|
||||
resolve call (e.g. an asyncio.wait_for timeout, or a request abort) must
|
||||
be dropped, not raise InvalidStateError and kill the pump thread.
|
||||
"""
|
||||
|
||||
def test_compute_done_racing_cancel_does_not_crash_pump(self):
|
||||
executor = _make_executor()
|
||||
fut = _racy_cancelled_future()
|
||||
with executor._futures_lock:
|
||||
executor._rpc_futures["1"] = fut
|
||||
|
||||
msg = AsyncDiffusionOutput(kind=AsyncOutputKind.COMPUTE_DONE, rpc_id="1", async_output_id="abc")
|
||||
_feed_one_msg_to_pump(executor, msg)
|
||||
|
||||
assert fut.cancelled()
|
||||
|
||||
def test_output_ready_racing_cancel_does_not_crash_pump(self):
|
||||
executor = _make_executor()
|
||||
fut = _racy_cancelled_future()
|
||||
with executor._futures_lock:
|
||||
executor._output_futures["abc123"] = fut
|
||||
|
||||
msg = AsyncDiffusionOutput(
|
||||
kind=AsyncOutputKind.OUTPUT_READY,
|
||||
async_output_id="abc123",
|
||||
output=DiffusionOutput(output="late"),
|
||||
)
|
||||
_feed_one_msg_to_pump(executor, msg)
|
||||
|
||||
assert fut.cancelled()
|
||||
|
||||
def test_batch_split_racing_cancel_does_not_crash_pump(self):
|
||||
executor = _make_executor()
|
||||
cancelled_fut = _racy_cancelled_future()
|
||||
healthy_fut = concurrent.futures.Future()
|
||||
with executor._futures_lock:
|
||||
executor._output_futures["batch-1/r-aborted"] = cancelled_fut
|
||||
executor._output_futures["batch-1/r-healthy"] = healthy_fut
|
||||
executor._batch_split_map["batch-1"] = {
|
||||
"batch-1/r-aborted": "r-aborted",
|
||||
"batch-1/r-healthy": "r-healthy",
|
||||
}
|
||||
|
||||
outputs = {
|
||||
"r-aborted": DiffusionOutput(output="late"),
|
||||
"r-healthy": DiffusionOutput(output="ok"),
|
||||
}
|
||||
msg = AsyncDiffusionOutput(
|
||||
kind=AsyncOutputKind.OUTPUT_READY,
|
||||
async_output_id="batch-1",
|
||||
output=_FakeBatchOutput(outputs),
|
||||
)
|
||||
_feed_one_msg_to_pump(executor, msg)
|
||||
|
||||
assert cancelled_fut.cancelled()
|
||||
assert healthy_fut.done()
|
||||
assert healthy_fut.result(timeout=1.0) is outputs["r-healthy"]
|
||||
@@ -1728,112 +1728,3 @@ def test_purge_is_noop_on_empty_deques(build_adapter):
|
||||
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests={})
|
||||
assert running_queue == []
|
||||
assert waiting_queue == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Chunk-wait deadline (RFC #4855 R1.1, issue #3833)
|
||||
#
|
||||
# Before this, the async-chunk path had no deadline of any kind, so a dropped
|
||||
# terminal chunk or an upstream stage that died mid-stream parked the request
|
||||
# in WAITING_FOR_CHUNK forever. None of these scenarios were covered.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def _park_in_chunk_wait(adapter, request, *, waiting=True):
|
||||
"""Drive a request through one process_pending_chunks round into the wait."""
|
||||
queue = DummyWaitingQueue([request]) if waiting else [request]
|
||||
if waiting:
|
||||
adapter.process_pending_chunks(queue, [])
|
||||
else:
|
||||
adapter.process_pending_chunks(DummyWaitingQueue(), queue)
|
||||
return queue
|
||||
|
||||
|
||||
def test_chunk_wait_clock_starts_when_a_request_parks(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
|
||||
request = _req("r1", RequestStatus.WAITING)
|
||||
|
||||
_park_in_chunk_wait(adapter, request)
|
||||
|
||||
assert request.status == RequestStatus.WAITING_FOR_CHUNK
|
||||
assert "r1" in adapter._waiting_since
|
||||
|
||||
|
||||
def test_chunk_wait_does_not_expire_before_the_deadline(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
|
||||
_park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
|
||||
|
||||
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
|
||||
assert "r1" in adapter._waiting_since
|
||||
|
||||
|
||||
def test_dropped_terminal_chunk_expires_the_request(build_adapter):
|
||||
"""The #3833 scenario: upstream stops sending and never marks the stream done."""
|
||||
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
|
||||
_park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
|
||||
|
||||
adapter._waiting_since["r1"] -= 601.0
|
||||
|
||||
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"r1"}
|
||||
# Cleared, so a second sweep does not re-report the same request.
|
||||
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
|
||||
|
||||
|
||||
def test_arriving_chunk_resets_the_clock(build_adapter):
|
||||
"""A slow but healthy stream must never expire: the deadline measures stall
|
||||
time between chunks, not the lifetime of the stream."""
|
||||
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
|
||||
request = _req("r1", RequestStatus.WAITING)
|
||||
queue = _park_in_chunk_wait(adapter, request)
|
||||
adapter.restore_queues(queue, [], scheduler_requests={"r1": request})
|
||||
|
||||
# Age the wait to just under the deadline, then deliver a chunk.
|
||||
adapter._waiting_since["r1"] -= 599.0
|
||||
adapter._finished_load_reqs.add("r1")
|
||||
adapter.process_pending_chunks(queue, [])
|
||||
|
||||
assert "r1" not in adapter._waiting_since
|
||||
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
|
||||
|
||||
|
||||
def test_a_disabled_deadline_never_expires(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
|
||||
_park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
|
||||
adapter._waiting_since["r1"] -= 10_000.0
|
||||
|
||||
assert adapter.collect_timed_out_request_ids(timeout_s=0.0) == set()
|
||||
assert adapter.collect_timed_out_request_ids(timeout_s=-1.0) == set()
|
||||
|
||||
|
||||
def test_finished_request_leaves_no_stale_timestamp(build_adapter):
|
||||
"""Otherwise a completed request would be reported as timed out later."""
|
||||
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
|
||||
request = _req("r1", RequestStatus.WAITING)
|
||||
_park_in_chunk_wait(adapter, request)
|
||||
|
||||
adapter.finish_requests(["r1"], RequestStatus.FINISHED_STOPPED, {"r1": request})
|
||||
|
||||
assert "r1" not in adapter._waiting_since
|
||||
assert adapter.collect_timed_out_request_ids(timeout_s=0.001) == set()
|
||||
|
||||
|
||||
def test_aborted_request_leaves_no_stale_timestamp(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
|
||||
_park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
|
||||
|
||||
adapter.cleanup_receiver("r1")
|
||||
|
||||
assert "r1" not in adapter._waiting_since
|
||||
|
||||
|
||||
def test_expiry_is_per_request(build_adapter):
|
||||
"""One stalled stream must not take down its healthy neighbours."""
|
||||
adapter, _ = build_adapter(stage_id=1, model_mode="ar", max_num_seqs=4)
|
||||
stalled = _req("stalled", RequestStatus.WAITING)
|
||||
healthy = _req("healthy", RequestStatus.WAITING)
|
||||
adapter.process_pending_chunks(DummyWaitingQueue([stalled, healthy]), [])
|
||||
|
||||
adapter._waiting_since["stalled"] -= 601.0
|
||||
|
||||
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"stalled"}
|
||||
assert "healthy" in adapter._waiting_since
|
||||
@@ -463,7 +463,7 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.full_model
|
||||
@pytest.mark.benchmark
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
@@ -591,7 +591,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
|
||||
assert audio_metrics["relative_l2"] <= LTX25_AUDIO_RELATIVE_L2_THRESHOLD
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.full_model
|
||||
@pytest.mark.benchmark
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
|
||||
@@ -544,7 +544,7 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", CASES, ids=lambda case: case.name)
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.full_model
|
||||
@pytest.mark.benchmark
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
|
||||
@@ -17,7 +17,7 @@ from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.media import get_asset_path
|
||||
from tests.helpers.runtime import OmniServer
|
||||
|
||||
pytestmark = [pytest.mark.diffusion]
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
SINGLE_MODEL = "Qwen/Qwen-Image-Edit"
|
||||
@@ -253,7 +253,6 @@ def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.benchmark
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_qwen_image_edit_single_matches_diffusers(
|
||||
@@ -279,7 +278,6 @@ def test_qwen_image_edit_single_matches_diffusers(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.benchmark
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_qwen_image_edit_2511_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None:
|
||||
|
||||
@@ -17,7 +17,7 @@ from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServer
|
||||
|
||||
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
MODEL_ID = "Qwen/Qwen-Image-Layered"
|
||||
|
||||
@@ -11,7 +11,7 @@ from tests.e2e.accuracy.helpers import assert_video_similarity_metrics, probe_vi
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL = "BestWishYsh/Helios-Distilled"
|
||||
PROMPT = "A serene lakeside sunrise with mist over the water."
|
||||
|
||||
@@ -138,7 +138,7 @@ def _make_file_lora_request(adapter_dir: Path) -> LoRARequest:
|
||||
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.slow,
|
||||
pytest.mark.full_model,
|
||||
pytest.mark.diffusion,
|
||||
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
|
||||
]
|
||||
|
||||
@@ -109,7 +109,7 @@ def _diffusion_sampling_params(
|
||||
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.slow,
|
||||
pytest.mark.full_model,
|
||||
pytest.mark.diffusion,
|
||||
]
|
||||
|
||||
|
||||
@@ -3,18 +3,19 @@
|
||||
"""Smoke test for HunyuanImage-3.0 Image-to-Text (I2T) pipeline."""
|
||||
|
||||
from collections.abc import Generator
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt
|
||||
|
||||
MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
|
||||
DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image3_ar.yaml")
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image3_i2t.yaml"
|
||||
|
||||
# First 20 generated token IDs from the HF greedy reference on this input.
|
||||
# vllm-omni AR output matches this prefix bitwise; the two implementations
|
||||
@@ -44,14 +45,14 @@ EXPECTED_PREFIX_TOKEN_IDS: list[int] = [
|
||||
# Decoded form, kept only for human-readable assertion messages.
|
||||
EXPECTED_PREFIX_TEXT = "The image is a solid, uniform green color with no variations, objects, or details present. It"
|
||||
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion, pytest.mark.cuda]
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def omni() -> Generator[Omni, None, None]:
|
||||
with OmniRunner(
|
||||
MODEL_NAME,
|
||||
deploy_config=DEPLOY_CONFIG_PATH,
|
||||
deploy_config=str(STAGE_CONFIG_PATH),
|
||||
trust_remote_code=True,
|
||||
) as runner:
|
||||
yield runner.omni
|
||||
|
||||
@@ -1,84 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""End-to-end smoke test for LingBot-World v2 generation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniRunnerHandler
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
MODEL = os.environ.get(
|
||||
"VLLM_OMNI_LINGBOT_WORLD_V2_CHECKPOINT_PATH",
|
||||
"robbyant/lingbot-world-v2-14b-causal-fast-diffusers",
|
||||
)
|
||||
_IMAGE_PATH = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH")
|
||||
_ACTION_DIR_VALUE = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_ACTION_DIR")
|
||||
_ACTION_DIR = Path(_ACTION_DIR_VALUE).expanduser().resolve() if _ACTION_DIR_VALUE else None
|
||||
_HAS_INPUT_ASSETS = _IMAGE_PATH is not None and _ACTION_DIR is not None
|
||||
|
||||
_OMNI_RUNNER_PARAM = (
|
||||
MODEL,
|
||||
None,
|
||||
{
|
||||
"flow_shift": 5.0,
|
||||
"parallel_config": DiffusionParallelConfig(tensor_parallel_size=1),
|
||||
"model_config": {
|
||||
"lingbot_action_root": str(_ACTION_DIR.parent) if _ACTION_DIR is not None else "",
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.slow,
|
||||
pytest.mark.diffusion,
|
||||
pytest.mark.skipif(
|
||||
not _HAS_INPUT_ASSETS,
|
||||
reason=("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH and VLLM_OMNI_LINGBOT_WORLD_V2_ACTION_DIR are required"),
|
||||
),
|
||||
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
|
||||
]
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_lingbot_world_v2_tp1_one_block(omni_runner_handler: OmniRunnerHandler) -> None:
|
||||
"""Load the real checkpoint and generate one causal block on one GPU."""
|
||||
|
||||
assert _IMAGE_PATH is not None
|
||||
assert _ACTION_DIR is not None
|
||||
response = omni_runner_handler.send_diffusion_request(
|
||||
{
|
||||
"model": omni_runner_handler.runner.model_name,
|
||||
"prompt": "The camera moves slowly forward through the scene.",
|
||||
"multi_modal_data": {"image": str(Path(_IMAGE_PATH).expanduser().resolve())},
|
||||
"sampling_params": OmniDiffusionSamplingParams(
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=9,
|
||||
num_inference_steps=4,
|
||||
max_sequence_length=512,
|
||||
seed=42,
|
||||
fps=16,
|
||||
extra_args={
|
||||
"action_path": _ACTION_DIR.name,
|
||||
"flow_shift": 5.0,
|
||||
},
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
assert response.success
|
||||
assert response.images is not None
|
||||
assert len(response.images) == 1
|
||||
video = np.asarray(response.images[0])
|
||||
assert video.shape == (1, 9, 480, 832, 3)
|
||||
assert np.isfinite(video).all()
|
||||
assert float(np.ptp(video)) > 1e-6
|
||||
assert float(np.mean(np.abs(np.diff(video, axis=1)))) > 1e-6
|
||||
@@ -96,7 +96,7 @@ def _get_test_config() -> str:
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"),
|
||||
pytest.mark.slow,
|
||||
pytest.mark.full_model,
|
||||
pytest.mark.tts,
|
||||
pytest.mark.parametrize(
|
||||
"omni_runner",
|
||||
@@ -271,6 +271,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) -> None:
|
||||
"""MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
|
||||
@@ -282,6 +283,7 @@ def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str)
|
||||
assert not torch.all(audio == 0), "Audio is silence"
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_moss_tts_realtime_chinese(omni_runner: OmniRunner, ref_audio_path: str) -> None:
|
||||
"""MossTTSRealtime: Chinese input produces non-empty audio."""
|
||||
|
||||
@@ -22,7 +22,7 @@ from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
|
||||
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
PROMPT = "A futuristic city skyline at twilight, cyberpunk style, ultra-detailed, high resolution."
|
||||
NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark"
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
"""Weekly e2e coverage for BAGEL diffusion multi-replica serving.
|
||||
"""Nightly-only e2e coverage for BAGEL diffusion multi-replica serving.
|
||||
|
||||
This test needs 4 H100 GPUs; keep it out of generic test_bagel_* jobs.
|
||||
It is covered by the weekly Diffusion · H100 · Multi-GPU lane
|
||||
(``.buildkite/cuda/test-weekly.yml``, ``mirror_hardwares: h100_4``).
|
||||
"""
|
||||
|
||||
import os
|
||||
@@ -19,7 +17,7 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
|
||||
MULTI_REPLICA_DEPLOY = get_deploy_config_path("ci/bagel_multi_replicas_4gpu.yaml")
|
||||
ROUTE_STRESS_REQUESTS = 6
|
||||
|
||||
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
|
||||
|
||||
test_params = [
|
||||
OmniServerParams(
|
||||
|
||||
@@ -16,7 +16,7 @@ parity vs. the Diffusers baseline is covered separately (see the support plan, s
|
||||
From ``tests/``::
|
||||
|
||||
pytest -s -v e2e/online_serving/test_boogu_image.py \
|
||||
-m "slow and diffusion and H100"
|
||||
-m "full_model and diffusion and H100" --run-level=full_model
|
||||
"""
|
||||
|
||||
import os
|
||||
@@ -28,7 +28,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL = "Boogu/Boogu-Image-0.1-Base"
|
||||
T2I_PROMPT = "A mountain lake at sunset, photorealistic, cinematic lighting"
|
||||
|
||||
@@ -19,7 +19,7 @@ reference image is supported.
|
||||
From ``tests/``::
|
||||
|
||||
pytest -s -v e2e/online_serving/test_boogu_image_edit.py \
|
||||
-m "slow and diffusion and H100"
|
||||
-m "full_model and diffusion and H100" --run-level=full_model
|
||||
"""
|
||||
|
||||
import base64
|
||||
@@ -35,7 +35,7 @@ from tests.helpers.runtime import (
|
||||
dummy_messages_from_mix_data,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
EDIT_PROMPT = "Change the style to a colored pencil drawing."
|
||||
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||
|
||||
@@ -25,7 +25,7 @@ model and are intentionally omitted.
|
||||
From ``tests/``::
|
||||
|
||||
pytest -s -v e2e/online_serving/test_boogu_image_edit_expansion.py \
|
||||
-m "slow and diffusion and H100"
|
||||
-m "full_model and diffusion and H100" --run-level=full_model
|
||||
"""
|
||||
|
||||
import os
|
||||
@@ -38,7 +38,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL = "Boogu/Boogu-Image-0.1-Edit"
|
||||
EDIT_PROMPT = "Change the style to a colored pencil drawing."
|
||||
|
||||
@@ -30,7 +30,7 @@ omitted (add rows here when those features land).
|
||||
From ``tests/``::
|
||||
|
||||
pytest -s -v e2e/online_serving/test_boogu_image_expansion.py \
|
||||
-m "slow and diffusion and H100"
|
||||
-m "full_model and diffusion and H100" --run-level=full_model
|
||||
"""
|
||||
|
||||
import os
|
||||
@@ -42,7 +42,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL = "Boogu/Boogu-Image-0.1-Base"
|
||||
T2I_PROMPT = "A mountain lake at sunset, photorealistic, cinematic lighting"
|
||||
|
||||
@@ -30,7 +30,7 @@ from tests.helpers.runtime import (
|
||||
)
|
||||
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL = os.environ.get("GLM_IMAGE_MODEL_PATH", "zai-org/GLM-Image")
|
||||
|
||||
|
||||
@@ -63,7 +63,7 @@ def test_text_to_image_moe(omni_server: OmniServer, openai_client: OpenAIClientH
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.full_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize("omni_server", _get_server_cases(MODEL), indirect=True)
|
||||
def test_video_generation_modes_moe(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
|
||||
@@ -10,7 +10,7 @@ from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.media import generate_synthetic_image
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
DISTILLED_MODEL = os.getenv("VLLM_TEST_LTX23_DISTILLED_MODEL", "diffusers/LTX-2.3-Distilled-Diffusers")
|
||||
BASE_MODEL = os.getenv("VLLM_TEST_LTX23_MODEL", "diffusers/LTX-2.3-Diffusers")
|
||||
|
||||
@@ -20,7 +20,7 @@ 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.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
pytestmark = [pytest.mark.omni, pytest.mark.slow]
|
||||
pytestmark = [pytest.mark.omni, pytest.mark.full_model]
|
||||
|
||||
_SKIP_NEED_4_H100_NOT_CI = pytest.mark.skip(
|
||||
reason="Requires 4x H100 GPUs; skipped in CI for now.",
|
||||
|
||||
Loaded 100 of 191 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user