Compare commits

...
Author SHA1 Message Date
hsliuustc0106 b0942b9c23 [Core] Validate HSDP shard size before initialization
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
2026-08-29 06:00:36 +08:00
hsliuustc0106 5c8e0e85d8 [Core] Identify HSDP replica ranks in worker titles
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
2026-08-28 18:27:56 +08:00
hsliuustc0106 65ac6794f9 [Core] Restore topology-aware diffusion worker titles
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
2026-08-28 18:19:10 +08:00
Hongsheng Liu d5aced9296 Revert "[Core] Make diffusion worker titles topology-aware" (#6717) 2026-08-28 17:30:03 +08:00
Wallbreazzz 4ad455e61e [Feature][TTS] adaptive chunk ramp (Phase 2 buffer-feedback controller) (#6001)
Signed-off-by: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com>
2026-08-28 02:11:05 -07:00
amy-why-3459 da94c3ebe1 [Bugfix][MiniCPM-o] Bind omni-duplex video frames to the unit they close (#6404)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-28 16:58:20 +08:00
prettygirlisnotme 13173d6bbb [Core] Make diffusion worker titles topology-aware (#6307)
Signed-off-by: Dragonborn <1078757180@qq.com>
2026-08-28 16:00:30 +08:00
Sy03 009b992c5b [Bugfix][MiniCPM-o] Isolate duplex handoff failures and barge-in (#6170)
Signed-off-by: Sy03 <1370724210@qq.com>
2026-08-28 15:46:41 +08:00
Mu GuanLinandHongsheng Liu 04b97a5922 [Perf] MiniMax-H3 VAE Decoder Ops (#6607)
Signed-off-by: mglyn <1203789601@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-28 15:14:33 +08:00
TJianandroot 15c6ab8fa7 [AMD][CI] Fix AMD CI (Partial) (#6704)
Signed-off-by: root <root@tw020.mia.tensorwave.lan>
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Co-authored-by: root <root@tw020.mia.tensorwave.lan>
2026-08-28 14:34:11 +08:00
zzh c201fd77c6 Add HYvideo1.5 benchmark (#6349)
Signed-off-by: zzh <943967662@qq.com>
2026-08-28 14:31:03 +08:00
b6a4baa659 [Feature] Support text encoder Disaggregation for H3 (#5885)
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Signed-off-by: yuanwu <yuan.wu@intel.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: yuanwu <yuan.wu@intel.com>
2026-08-28 14:27:05 +08:00
WeiQing Chen ed837e350e [Test] Stabilize MiniMax-H3 reference accuracy compile (#6688)
Signed-off-by: david6666666 <530634352@qq.com>
2026-08-28 14:18:28 +08:00
BruceLoveDecimalandamy-why-3459 ab8d45d3cc [Bugfix][MiniCPM-o] Retire CFM DiT CUDA graphs a generation at a time (#6587)
Signed-off-by: BruceLoveDecimal <liuqihao970610@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-28 14:07:57 +08:00
Eugene Zhao f8f8aa8eb7 [Bugfix] Honor stage selection in online profiler endpoints (#6609)
Signed-off-by: qujing226 <qujing226@gmail.com>
2026-08-28 11:41:47 +08:00
Hongsheng Liu 6ec5ab11bd [Bugfix] Make HWR filesystem lifecycle transitions fail-closed (#6692)
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-08-28 09:30:17 +08:00
Hongsheng Liuandhsliuustc0106 4245fd6955 [Fix] Restore MammothModa2 support (#6694)
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
2026-08-28 08:32:32 +08:00
Mu GuanLinandHongsheng Liu 5e4daaf197 [Diffusion] Add LTX-2.5 Diffusion VAE decoder support (#6189)
Signed-off-by: mglyn <1203789601@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-27 23:54:40 +08:00
NATUREandamy-why-3459 262ff14693 [Bugfix][MiniCPM-o] Bound native auto-response continuation (#6630)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-27 19:51:47 +08:00
Deep Shahandamy-why-3459 069e6c5fbc [CI/Build][MiniCPM-o] Derive duplex admission-probe limit from the deploy config (#6678)
Signed-off-by: Deep Shah <deep@socratic.co>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-27 19:20:01 +08:00
NATURE 54c1fc64b7 [Bugfix] Gate async-chunk segment resume by session mode (#6680)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
2026-08-27 18:52:09 +08:00
Zhanshenzhuiandamy-why-3459 c6df39e9f7 [Benchmark][MiniCPM-o] Add Omni-DuplexEval support (#6634)
Signed-off-by: Zhanshenzhui <99891002+zyforsure@users.noreply.github.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-27 16:31:50 +08:00
Zheng Wengangandamy-why-3459 94b1546ce3 [Model] Qwen3-Omni: register thinker-only pipeline for Instruct serve (#6284)
Signed-off-by: ZhengWG <zwg0606@gmail.com>
Signed-off-by: Zheng Wengang <zwg0606@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-27 08:21:03 +00:00
19da238246 [Bugfix][NPU] Eager-import mindiesd for every diffusion attention backend (#6054)
Signed-off-by: brandneway <gyuan4892@gmail.com>
Co-authored-by: changetheway <guotaoyuan1@h-partners.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-08-27 14:41:36 +08:00
211 changed files with 19198 additions and 990 deletions

No files matched your search

+89 -20
View File
@@ -1,4 +1,7 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# vllm-omni customized version
# Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/bootstrap-amd.sh
# Last synced: 2025-12-15
@@ -6,6 +9,9 @@
set -euo pipefail
# The bootstrap runs from the repository root; ShellCheck does not follow
# sourced files unless invoked with -x.
# shellcheck disable=SC1091
source .buildkite/common/scripts/resolve_skip_ci.sh
if [[ -z "${RUN_ALL:-}" ]]; then
@@ -63,6 +69,8 @@ upload_pipeline() {
# Install minijinja
ls .buildkite || buildkite-agent annotate --style error 'Please merge upstream main branch for buildkite CI'
curl -sSfL https://github.com/mitsuhiko/minijinja/releases/download/2.3.1/minijinja-cli-installer.sh | sh
# Installed by the minijinja bootstrap above and only present on the CI agent.
# shellcheck disable=SC1091
source /var/lib/buildkite-agent/.cargo/env
if [[ $BUILDKITE_PIPELINE_SLUG == "fastcheck" ]]; then
@@ -90,22 +98,82 @@ upload_pipeline() {
cd .buildkite/amd
# Select test definition file: merge suite for main, ready suite for PRs.
# For debugging, DEBUG_TEST_YAML can override the selection — accepts
# "merge" or "ready" (case-insensitive).
# For debugging, DEBUG_TEST_YAML accepts a comma-separated list containing
# "merge" and/or "ready" (case-insensitive). Multiple suites are combined
# before rendering so they share one amd-build step.
if [[ -n "${DEBUG_TEST_YAML:-}" ]]; then
case "${DEBUG_TEST_YAML,,}" in
merge)
TEST_YAML="test-amd-merge.yml"
;;
ready)
TEST_YAML="test-amd-ready.yml"
;;
*)
echo "ERROR: DEBUG_TEST_YAML must be 'merge' or 'ready', got '$DEBUG_TEST_YAML'" >&2
declare -a DEBUG_TEST_SPECS=()
declare -A SEEN_DEBUG_TESTS=()
IFS=',' read -ra REQUESTED_DEBUG_TESTS <<< "${DEBUG_TEST_YAML,,}"
for requested_test in "${REQUESTED_DEBUG_TESTS[@]}"; do
# Permit readable values such as "ready, merge," and ignore the
# empty item produced by a trailing comma.
requested_test="${requested_test//[[:space:]]/}"
[[ -z "$requested_test" ]] && continue
if [[ -n "${SEEN_DEBUG_TESTS[$requested_test]:-}" ]]; then
echo "ERROR: duplicate DEBUG_TEST_YAML suite '$requested_test'" >&2
exit 1
;;
esac
echo "DEBUG_TEST_YAML override: using $TEST_YAML"
fi
SEEN_DEBUG_TESTS[$requested_test]=1
case "$requested_test" in
ready)
DEBUG_TEST_SPECS+=("READY_TESTS:test-amd-ready.yml")
;;
merge)
DEBUG_TEST_SPECS+=("MERGE_TESTS:test-amd-merge.yml")
;;
*)
echo "ERROR: DEBUG_TEST_YAML entries must be 'merge' or 'ready', got '$requested_test'" >&2
exit 1
;;
esac
done
if [[ ${#DEBUG_TEST_SPECS[@]} -eq 0 ]]; then
echo "ERROR: DEBUG_TEST_YAML did not contain a test suite" >&2
exit 1
elif [[ ${#DEBUG_TEST_SPECS[@]} -eq 1 ]]; then
TEST_YAML="${DEBUG_TEST_SPECS[0]#*:}"
else
TEST_YAML=$(mktemp "${TMPDIR:-/tmp}/amd-debug-tests.XXXXXX.yml")
python - "$TEST_YAML" "${DEBUG_TEST_SPECS[@]}" <<'PY'
import sys
import yaml
output_path, *suite_specs = sys.argv[1:]
combined = {"env": {}, "steps": []}
for suite_spec in suite_specs:
group_name, input_path = suite_spec.split(":", 1)
with open(input_path, encoding="utf-8") as test_file:
suite = yaml.safe_load(test_file)
for name, value in (suite.get("env") or {}).items():
previous = combined["env"].get(name, value)
if previous != value:
raise ValueError(
f"Conflicting environment value for {name}: {previous!r} != {value!r}"
)
combined["env"][name] = value
suite_steps = []
for entry in suite.get("steps") or []:
if "group" in entry:
suite_steps.extend(entry.get("steps") or [])
else:
suite_steps.append(entry)
combined["steps"].append({"group": group_name, "steps": suite_steps})
with open(output_path, "w", encoding="utf-8") as output_file:
yaml.safe_dump(combined, output_file, sort_keys=False)
PY
fi
echo "DEBUG_TEST_YAML override: using ${DEBUG_TEST_SPECS[*]}"
elif [[ $BUILDKITE_BRANCH == "main" ]]; then
TEST_YAML="test-amd-merge.yml"
else
@@ -130,19 +198,20 @@ upload_pipeline() {
> pipeline.yaml
)
cat pipeline.yaml
if [[ "$TEST_YAML" == "${TMPDIR:-/tmp}/amd-debug-tests."*.yml ]]; then
rm -f -- "$TEST_YAML"
fi
buildkite-agent artifact upload pipeline.yaml
buildkite-agent pipeline upload pipeline.yaml
exit 0
}
get_diff() {
$(git add .)
echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
git diff --name-only --diff-filter=ACMDR "$(git merge-base origin/main HEAD)"
}
get_diff_main() {
$(git add .)
echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
git diff --name-only --diff-filter=ACMDR HEAD~1
}
file_diff=$(get_diff)
@@ -181,7 +250,7 @@ for file in $file_diff; do
# First check if file matches any pattern
matches_pattern=0
for pattern in "${patterns[@]}"; do
if [[ $file == $pattern* ]] || [[ $file == $pattern ]]; then
if [[ $file == "$pattern"* ]] || [[ $file == "$pattern" ]]; then
matches_pattern=1
break
fi
@@ -191,7 +260,7 @@ for file in $file_diff; do
if [[ $matches_pattern -eq 1 ]]; then
matches_ignore=0
for ignore in "${ignore_patterns[@]}"; do
if [[ $file == $ignore* ]] || [[ $file == $ignore ]]; then
if [[ $file == "$ignore"* ]] || [[ $file == "$ignore" ]]; then
matches_ignore=1
break
fi
+44 -16
View File
@@ -21,7 +21,7 @@ steps:
- label: "Simple · Diffusion Test · Shard %N/%t"
agent_pool: mi300_1
parallelism: 2
parallelism: 4
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -71,26 +71,30 @@ steps:
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
bash -c '
set +e
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
- pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
- pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs) · Shard 1/2"
agent_pool: mi300_4
depends_on: amd-build
mirror_hardwares: [amdproduction]
timeout_in_minutes: 40
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
- timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
- pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs) · Shard 2/2"
agent_pool: mi300_4
depends_on: amd-build
mirror_hardwares: [amdproduction]
timeout_in_minutes: 40
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
# merge-only tests
- label: "Diffusion Tensor Parallelism Test"
@@ -175,17 +179,41 @@ steps:
- label: "Qwen3-TTS CustomVoice E2E Test · Shard %N/%t"
agent_pool: mi300_1
parallelism: 2
parallelism: 3
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
timeout_in_minutes: 80
commands:
- |
bash -c '
timeout --signal=TERM --kill-after=2m 80m 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_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
TESTS=(
"tests/e2e/online_serving/test_qwen3_tts_customvoice.py::test_default_cuda_graph_startup[async_chunk]"
"tests/e2e/online_serving/test_qwen3_tts_customvoice.py::test_text_to_audio_002[async_chunk]"
"tests/e2e/offline_inference/test_qwen3_tts_customvoice.py::test_text_to_audio_001[no_cuda_graph]"
)
TEST="$${TESTS[$$BUILDKITE_PARALLEL_JOB]}"
pytest -s -v "$$TEST" -m "advanced_model" --run-level "advanced_model"
'
# TODO: Restore this case to the blocking sharded step once its long-running
# Whisper validation on AMD CI is resolved.
- label: "Qwen3-TTS CustomVoice async_chunk Test (quarantined)"
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: NonBlocking
timeout_in_minutes: 80
commands:
- |
timeout --signal=TERM --kill-after=2m 75m 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_qwen3_tts_customvoice.py::test_text_to_audio_001[async_chunk]" -m "advanced_model" --run-level "advanced_model"
'
- label: "Qwen3-TTS Base E2E Test"
+1 -4
View File
@@ -191,10 +191,7 @@ steps:
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 45m bash -c '
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"
'
- pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
# Need to debug before reneable numerical changes across large PRs
+13
View File
@@ -166,6 +166,7 @@ steps:
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
- tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
- vllm_omni/deploy/minicpmo_4_5.yaml
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
@@ -216,3 +217,15 @@ steps:
pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
"
mirror_hardwares: h100_1
- label: "Diffusion · HunyuanVideo-1.5 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_hunyuan_video_15_expansion.py
- vllm_omni/diffusion/models/hunyuan_video/
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
"
mirror_hardwares: h100_1
+23
View File
@@ -272,6 +272,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 · HunyuanVideo-1.5 Function Test"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · Doc Test"
timeout_in_minutes: 60
commands:
@@ -322,6 +328,23 @@ steps:
exit $$EXIT2
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · Perf Test · HunyuanVideo-1.5"
key: nightly-diffusion-x2v-performance-hunyuanvideo15
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json
EXIT1=$$?
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json
EXIT2=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
if [ $$EXIT1 -ne 0 ]; then exit $$EXIT1; fi
exit $$EXIT2
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · Perf Test · MiniMax-H3"
key: nightly-diffusion-x2v-performance-minimax-h3
timeout_in_minutes: 180
+13
View File
@@ -205,6 +205,7 @@ steps:
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
- tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
- vllm_omni/deploy/minicpmo_4_5.yaml
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
@@ -337,3 +338,15 @@ steps:
pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: h100_1
- label: "Diffusion · HunyuanVideo-1.5 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_hunyuan_video_15_expansion.py
- vllm_omni/diffusion/models/hunyuan_video/
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m 'core_model and diffusion' --run-level 'core_model'
"
mirror_hardwares: h100_1
+4
View File
@@ -114,6 +114,10 @@ Configuration classes.
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
- [vllm_omni.transformers_utils.configs.cosyvoice3.CosyVoice3Config][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
- [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][]
- [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][]
+1 -1
View File
@@ -59,7 +59,7 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze
| `edges` | list | optional | `null` | Explicit edge list for the KV transfer graph. Auto-derived from stage inputs if omitted. |
| `stages` | list | optional | `[]` | Per-stage runtime overrides matched by `stage_id`. Pipeline stages are still created from `PipelineConfig` when this list is empty. |
| `platforms` | dict | optional | `null` | Keyed by `npu` / `rocm` / `xpu`, each contains a `stages:` list with per-platform overrides applied on top of the CUDA defaults. |
| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only`). |
| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only` / `qwen3_omni_moe_thinker_only`). |
| `trust_remote_code` | bool \| null | optional | `null` | **Pipeline-wide.** Trust HF remote code on model load; applies to every stage when specified. |
| `distributed_executor_backend` | str \| null | optional | `null` | **Pipeline-wide.** Distributed executor backend forwarded to vLLM (`"mp"`, `"ray"`, `"external_launcher"`). If omitted, vLLM auto-selects backend from runtime topology. |
| `dtype` | str \| null | optional | `null` | **Pipeline-wide.** Model dtype for every stage. |
+47
View File
@@ -257,6 +257,53 @@ Lookup checks the marker both before validation and after constructing the
lease so a weaker overlapping lookup cannot return a lease after a stronger
validation has denied the artifact.
If publishing the external deny marker fails, lookup returns a retryable
storage failure and marks the artifact directory with a persistent invalidation
mode before releasing the shared lock. Cooperating readers check that mode both
while opening and immediately before returning a lease. Lookup then makes a
bounded attempt to take the exclusive artifact lock and quarantine the entry;
the quarantine transition removes the internal mode marker. If a cooperating
builder replaces the marked inode while the locks are exchanged, containment
preserves that new publication instead of quarantining it. If neither the
external marker nor the directory-mode marker can be persisted, cross-process
exclusion cannot be guaranteed on an unavailable filesystem and the domain
must remain failed rather than treating the artifact as an ordinary invalid
cache entry.
New publications carry the internal invalidation mode through the staging
rename. Hardening clears it and is the logical authority transition for a
process crash; the temporary- and artifact-parent syncs then make the rename
durable, with an artifact-inode sync first making the hardened mode durable. A
process that dies before hardening leaves a persistently denied
entry, while one that dies after successful hardening leaves an authoritative
ready entry. Any synchronous failure after rename is marked invalid and
contained in quarantine before `PUBLICATION_FAILED` is returned, so that
nonretryable result cannot later resolve as a ready hit. If the filesystem
refuses both the inode marker and the quarantine move, the store instead returns
a retryable storage failure with `outcome_uncertain` details; as with deny-marker
double failure, cross-process exclusion cannot be guaranteed until storage
recovers. Once the artifact-parent sync succeeds, a later artifact- or
build-lock release error is logged and lookup resolves the committed entry
instead of reporting a contradictory publication failure. The same authority
rule applies when a competing publisher already installed byte-identical
semantic content; staging-cleanup errors are retried without downgrading that
validated ready entry to publication failure.
Before widening a hardened artifact for a lifecycle move, the store durably
sets the same invalidation mode. A process crash before rename therefore leaves
the writable source denied. After rename, the store re-hardens the destination
but retains the marker until both parent-directory syncs make the move durable;
only then does it durably clear the marker. An interrupted move therefore leaves
a marked transition that the next exact-key lifecycle operation re-hardens.
Explicit cleanup uses the locked move and retries exact-key
`.cleanup.*` tombstones left by an interrupted removal; a later build performs
the same tombstone reconciliation before producing replacement content.
An operational failure while inspecting a noncooperative competing publication
is a retryable storage failure. The competing entry remains authoritative and
is not mislabeled as corrupt or quarantined merely because its manifest could
not be read transiently.
`inspect_domain()` and `inspect_artifact()` return typed, read-only snapshots.
Inspection validation does not create a deny marker. Reported lock activity is
point-in-time evidence only; kernel locks are not a persistent owner registry.
+2
View File
@@ -69,6 +69,8 @@ th {
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
| `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | — |
| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | — |
+1
View File
@@ -129,6 +129,7 @@ The following tables show which models support each feature:
| **Krea 2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **OmniGen2** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
@@ -28,6 +28,18 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/deploy_config_file
```
To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct
weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key
`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
```
Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select
the same single-stage pipeline without `--deploy-config`.
### Launch individual stages (stage-based CLI)
Adopt the stage-based CLI architecture to independently instantiate execution processes per functional stage.
@@ -0,0 +1,117 @@
"""
Offline inference example: MammothModa2 image summarization (single AR stage).
Example:
uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
--model path/to/MammothModa2-Preview \
--deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \
--image /path/to/input.jpg \
--question "Please summarize the content of this image."
"""
from __future__ import annotations
import argparse
import os
from PIL import Image
from vllm import SamplingParams
from vllm.multimodal.image import convert_image_mode
from vllm_omni import Omni
DEFAULT_SYSTEM = "You are a helpful assistant."
DEFAULT_QUESTION = "Please summarize the content of this image."
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.")
parser.add_argument("--image", type=str, required=True, help="Path to input image.")
parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
parser.add_argument(
"--max-tokens",
type=int,
default=512,
help="Max new tokens to generate.",
)
parser.add_argument("--temperature", type=float, default=0.2)
parser.add_argument("--top-p", type=float, default=0.9)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--trust-remote-code", action="store_true")
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
return parser.parse_args()
def build_prompt(system: str, question: str) -> str:
return (
f"<|im_start|>system\n{system}<|im_end|>\n"
"<|im_start|>user\n"
"<|vision_start|><|image_pad|><|vision_end|>"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
def main() -> None:
args = parse_args()
if not os.path.exists(args.image):
raise FileNotFoundError(f"Image file not found: {args.image}")
pil_image = Image.open(args.image)
image_data = convert_image_mode(pil_image, "RGB")
prompt = build_prompt(args.system, args.question)
omni = Omni(
model=args.model,
deploy_config=args.deploy_config,
trust_remote_code=args.trust_remote_code,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
)
try:
sp = SamplingParams(
temperature=float(args.temperature),
top_p=float(args.top_p),
top_k=-1,
max_tokens=int(args.max_tokens),
seed=int(args.seed),
detokenize=True,
)
# NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None].
# Consume it inside the try block so the worker isn't closed early.
outputs = list(
omni.generate(
[
{
"prompt": prompt,
"multi_modal_data": {"image": image_data},
"additional_information": {"omni_task": ["chat"]},
}
],
[sp],
)
)
finally:
omni.close()
lines: list[str] = []
for stage_outputs in outputs:
ro = stage_outputs
text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
lines.append("answer:\n")
lines.append(text.strip() + "\n")
lines.append("\n")
print("\n".join(lines))
if __name__ == "__main__":
main()
@@ -599,6 +599,17 @@ def main():
if args.seed is not None and hasattr(params, "seed"):
params.seed = args.seed
# MammothModa2's AR stage emits one visual token per grid cell,
# one EOL token per row, and one final look-ahead token whose hidden
# state is unavailable. Size the first stage from the prompt metadata
# instead of SamplingParams' default of 16 tokens.
prompt_info = prompt_dict.get("additional_information", {})
if idx == 0 and prompt_info.get("omni_task") == ["t2i"]:
ar_width = int(prompt_info.get("ar_width", [0])[0])
ar_height = int(prompt_info.get("ar_height", [0])[0])
if ar_width > 0 and ar_height > 0:
params.max_tokens = ar_height * (ar_width + 1) + 1
if not diffusion_replaced and len(sampling_params_list) == 1:
sampling_params_list = [diffusion_params]
+31 -3
View File
@@ -10,10 +10,12 @@ Generate text from text or image inputs with vLLM-Omni's shared offline entrypoi
| --- | --- | --- | --- |
| `ByteDance-Seed/BAGEL-7B-MoT` | Yes | Yes | `vllm_omni/deploy/bagel.yaml` |
| `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | `vllm_omni/deploy/hunyuan_image3_ar.yaml` |
| `bytedance-research/MammothModa2-Preview` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` |
| `bytedance-research/MammothModa2-Dev` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` |
The script recognizes these model families from `config.json`, applies the
model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3.
BAGEL uses its registered default deploy.
model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3
and MammothModa2. BAGEL uses its registered default deploy.
## Text-To-Text
@@ -33,6 +35,22 @@ python examples/offline_inference/x_to_text/x_to_text.py \
--prompt "Explain multimodal inference in three concise sentences."
```
### MammothModa2
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model bytedance-research/MammothModa2-Preview \
--prompt "Explain multimodal inference in three concise sentences."
```
The same command supports the Dev checkpoint:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model bytedance-research/MammothModa2-Dev \
--prompt "Explain multimodal inference in three concise sentences."
```
## Image-To-Text
Pass one image with `--image`. The default question can be replaced with any
@@ -54,6 +72,16 @@ python examples/offline_inference/x_to_text/x_to_text.py \
--prompt "Describe this image in detail."
```
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model bytedance-research/MammothModa2-Preview \
--image image.png \
--prompt "Describe this image in detail."
```
Use `bytedance-research/MammothModa2-Dev` in the same command to run I2T with
the Dev checkpoint.
## Key Arguments
| Argument | Default | Description |
@@ -74,4 +102,4 @@ python examples/offline_inference/x_to_text/x_to_text.py \
The committed HunyuanImage-3 AR-only default uses four GPUs with TP=4. To use a
different topology, pass a custom AR-only YAML through `--deploy-config`.
BAGEL uses its repository deploy default unless overridden.
BAGEL and MammothModa2 use their repository deploy defaults unless overridden.
@@ -15,6 +15,7 @@ from vllm_omni.model_extras import build_x_to_text_prompt, get_x_to_text_model_f
_REPO_ROOT = Path(__file__).resolve().parents[3]
_HUNYUAN_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml"
_MAMMOTH_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "mammoth_moda2_ar.yaml"
def parse_args() -> argparse.Namespace:
@@ -61,7 +62,7 @@ def main() -> None:
omni_kwargs: dict[str, Any] = {
"model": args.model,
"mode": "image-to-text" if image is not None else "text-to-text",
"trust_remote_code": args.trust_remote_code or family == "hunyuan_image3",
"trust_remote_code": args.trust_remote_code or family in {"hunyuan_image3", "mammoth_moda2"},
"log_stats": args.log_stats,
}
if args.enforce_eager is not None:
@@ -70,6 +71,8 @@ def main() -> None:
omni_kwargs["deploy_config"] = args.deploy_config
elif family == "hunyuan_image3":
omni_kwargs["deploy_config"] = str(_HUNYUAN_AR_DEPLOY_CONFIG)
elif family == "mammoth_moda2":
omni_kwargs["deploy_config"] = str(_MAMMOTH_AR_DEPLOY_CONFIG)
omni = Omni(**omni_kwargs)
try:
@@ -159,6 +159,33 @@ python examples/online_serving/minicpmo/realtime_duplex_demo.py \
--output-dir /tmp/minicpmo_realtime_duplex_demo
```
Video input uses the same session. PyAV demuxes JPEG frames at `--video-fps`
(default 1.0) and vLLM `load_audio` extracts a 16 kHz mono WAV unless
`--input-wav` overrides the soundtrack. Frames keep their capture size
(`--frame-max-side 0`); the server `process_image` normalizes at 448.
`--stack-frames N` raises the visual refresh rate the way official duplex does:
each 1 s unit also samples the `N-1` sub-frames captured inside it and tiles
them into a single composite image sent next to that unit's base frame. The
audio cadence never changes — a unit is always one second — and the wire always
carries 2 images per unit however large `N` is, because the sub-frames share one
composite. Official uses 5 for high refresh rate mode:
```bash
python examples/online_serving/minicpmo/realtime_duplex_demo.py \
--url ws://localhost:8099/v1/realtime?duplex=1 \
--model openbmb/MiniCPM-o-4_5 \
--input-video /path/to/clip.mp4 \
--stack-frames 5 \
--ref-audio /path/to/MiniCPM-o-Demo/assets/ref_audio/ref_minicpm_signature.wav \
--output-dir /tmp/minicpmo_realtime_duplex_video_demo
```
Detail inside a composite is capped by `scale_resolution=448` at
`max_slice_nums=1`: official suggests HD slicing (`max_slice_nums=[2, 1]`) for
stacked frames, which the duplex adapter does not implement yet. Reading small
text or digits out of a wide scene is limited by that, not by frame timing.
## Open the experimental browser client
The browser UI serves the page and proxies the same-origin Realtime WebSocket to
@@ -192,6 +219,64 @@ stronger `response-required` mode is diagnostic: it requires a purpose-built
two-response WAV, its `--input-sha256`, and an
`--expect-second-response-substring` value.
## Run Omni-DuplexEval
Omni-DuplexEval generation uses the vLLM-Omni native MiniCPM-o duplex
endpoint. Evaluation uses a separate OpenAI-compatible multimodal judge
served from the same vLLM-Omni environment. This validates the vLLM-Omni
duplex client and local judge integration; it does not reproduce the paper's
original MiniCPM-o inference implementation.
Start the duplex generation server:
```bash
vllm serve openbmb/MiniCPM-o-4_5 --omni --trust-remote-code \
--deploy-config vllm_omni/deploy/minicpmo_4_5_duplex.yaml \
--served-model-name openbmb/MiniCPM-o-4_5 \
--host 0.0.0.0 --port 8099
```
Start a separate multimodal judge. The allowed path must contain any local
videos passed with `--judge-video-mode video_url`:
```bash
vllm serve Qwen/Qwen2.5-VL-7B-Instruct \
--served-model-name Qwen/Qwen2.5-VL-7B-Instruct \
--allowed-local-media-path /data/omni-duplex-eval \
--host 0.0.0.0 --port 8000
```
Generate, evaluate, and summarize:
```bash
vllm bench omni-duplex-eval --omni generate \
--url ws://127.0.0.1:8099/v1/realtime?duplex=1 \
--model openbmb/MiniCPM-o-4_5 \
--ref-audio /data/ref.wav \
--dataset Hothan/Omni-DuplexEval \
--concurrency 2 \
--response-root /data/omni-duplex-eval/responses
vllm bench omni-duplex-eval --omni evaluate \
--dataset Hothan/Omni-DuplexEval \
--response-root /data/omni-duplex-eval/responses \
--score-root /data/omni-duplex-eval/scores \
--judge-base-url http://127.0.0.1:8000 \
--judge-model Qwen/Qwen2.5-VL-7B-Instruct \
--judge-video-mode video_url \
--eval-workers 4
vllm bench omni-duplex-eval --omni summarize \
--score-root /data/omni-duplex-eval/scores
```
The dataset names such as `RTD_OCR` and `PR_correction` are Hugging Face
splits, not dataset configurations. Pass one with `--split`, or omit it to
run all nine splits. `--limit 1` is useful for a smoke test. Realtime
generation records `clock=media`; artifacts generated with
`--pace as-fast-as-possible` record `clock=invalid` and evaluation rejects
them unless `--allow-invalid-clock` is explicit.
## Related examples
- [Offline MiniCPM-o inference](../../offline_inference/minicpmo/)
@@ -12,7 +12,9 @@ from __future__ import annotations
import argparse
import asyncio
import base64
import io
import json
import math
import sys
import time
from pathlib import Path
@@ -22,8 +24,8 @@ if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
PCM16_BYTES_PER_SAMPLE, # noqa: F401 - compatibility export
PCM16_SAMPLE_RATE, # noqa: F401 - compatibility export
PCM16_BYTES_PER_SAMPLE,
PCM16_SAMPLE_RATE,
RealtimeDuplexClient,
RealtimeEventCollector,
build_realtime_url,
@@ -38,6 +40,10 @@ from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
reference_audio_data_url as _ref_audio_data_url,
)
from vllm_omni.experimental.fullduplex.video_stacking import ( # noqa: E402
concat_frames_b64,
unit_subframe_offsets,
)
class _StreamingOutputWriter:
@@ -159,14 +165,258 @@ def _event_count_after(
return sum(event.get("type") == event_type for event in events[index + 1 :])
async def run_demo(args: argparse.Namespace) -> dict[str, object]:
input_pcm16 = read_pcm16_wav(Path(args.input_wav))
if not input_pcm16:
raise ValueError("input WAV has no audio")
def _probe_video_frames_and_fps(video_path: Path) -> tuple[int, float]:
"""Frame count and average FPS via PyAV (same approach as Daily-Omni)."""
import av
with av.open(str(video_path)) as container:
stream = container.streams.video[0]
avg_fps = float(stream.average_rate) if stream.average_rate else 0.0
num_frames = stream.frames or 0
if num_frames <= 0:
duration: float | None = None
if stream.duration is not None and stream.time_base is not None:
duration = float(stream.duration * stream.time_base)
elif container.duration is not None:
duration = container.duration / av.time_base
if duration and avg_fps:
num_frames = int(math.floor(duration * avg_fps))
if num_frames <= 0:
stream.thread_type = "AUTO"
num_frames = sum(1 for _ in container.decode(stream))
if num_frames <= 0:
raise ValueError(f"No decodable video frames in {video_path}")
return num_frames, avg_fps or 1.0
def _upright(image, rotation: int):
"""Apply a container display-matrix rotation, in degrees, to ``image``.
PyAV hands back the stored raster and leaves the display matrix for the
caller, unlike the ffmpeg CLI which autorotates. Phone clips routinely
carry +/-90, so skipping this feeds the model a sideways frame: arrows and
digits stop being readable and the model mirrors the motion it reports.
"""
from PIL import Image
quarter_turns = {
90: Image.Transpose.ROTATE_90,
180: Image.Transpose.ROTATE_180,
270: Image.Transpose.ROTATE_270,
}
turn = quarter_turns.get(int(rotation) % 360)
return image.transpose(turn) if turn is not None else image
def _decode_video_frames_rgb(video_path: Path, frame_idx: list[int]) -> list:
"""Decode presentation-order ``frame_idx`` into upright RGB ``PIL.Image``s."""
import av
wanted = sorted({int(i) for i in frame_idx})
if not wanted:
return []
decoded: dict[int, object] = {}
with av.open(str(video_path)) as container:
stream = container.streams.video[0]
stream.thread_type = "AUTO"
cursor = 0
for position, frame in enumerate(container.decode(stream)):
if wanted[cursor] > position:
continue
image = _upright(frame.to_image(), getattr(frame, "rotation", 0) or 0)
while cursor < len(wanted) and wanted[cursor] <= position:
decoded[wanted[cursor]] = image
cursor += 1
if cursor >= len(wanted):
break
if not decoded:
raise ValueError(f"Decoded no frames from {video_path} for indices {wanted[:8]}")
last_image = decoded[max(decoded)]
return [decoded.get(int(i), last_image) for i in frame_idx]
def _resize_max_side(image, max_side: int):
width, height = image.size
longest = max(width, height)
if longest <= max_side or max_side <= 0:
return image
scale = max_side / float(longest)
new_size = (max(1, int(round(width * scale))), max(1, int(round(height * scale))))
return image.resize(new_size)
def _jpeg_bytes(image, *, quality: int = 95) -> bytes:
buf = io.BytesIO()
image.convert("RGB").save(buf, format="JPEG", quality=quality)
return buf.getvalue()
def _to_pcm16_bytes(audio_np) -> bytes:
import numpy as np
arr = np.asarray(audio_np).reshape(-1)
if arr.dtype != np.int16:
arr = (np.clip(arr, -1.0, 1.0) * 32767.0).astype(np.int16)
return arr.tobytes()
def _pcm16_wav_duration_s(path: Path) -> float:
pcm16 = read_pcm16_wav(path)
return len(pcm16) / float(PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
def _extract_video_frames(
video_path: Path,
work_dir: Path,
*,
fps: float,
max_side: int,
duration_s: float,
stack_frames: int = 1,
) -> tuple[list[str], list[str | None]]:
"""Demux JPEG frames only. Frame count follows ``duration_s``, not soundtrack.
Returns the base frames (one per model unit) and the parallel stacked track.
With ``stack_frames=N`` each unit also samples its ``N-1`` interior
sub-frames and tiles them into one composite, which is how official duplex
raises the visual refresh rate while keeping one second of audio per unit.
A unit whose interior falls past the end of the clip stacks nothing.
"""
num_frames, avg_fps = _probe_video_frames_and_fps(video_path)
sample_fps = max(float(fps), 1e-6)
unit_s = 1.0 / sample_fps
timeline_s = max(float(duration_s), unit_s)
num_samples = max(1, int(math.ceil(timeline_s * sample_fps - 1e-9)))
offsets = unit_subframe_offsets(stack_frames, unit_s=unit_s)
def frame_at(seconds: float) -> int:
# Official get_video_frame_audio_segments uses truncating int(), not
# round(), so 29.97 fps lands on the same presentation index.
return min(int(seconds * avg_fps), num_frames - 1)
base_idx = [frame_at(unit * unit_s) for unit in range(num_samples)]
sub_times = [
[unit * unit_s + offset for offset in offsets if unit * unit_s + offset < timeline_s]
for unit in range(num_samples)
]
sub_idx = [frame_at(seconds) for times in sub_times for seconds in times]
images = _decode_video_frames_rgb(video_path, base_idx + sub_idx)
sub_images = iter(images[len(base_idx) :])
frame_dir = work_dir / "frames"
frame_dir.mkdir(parents=True, exist_ok=True)
for stale in frame_dir.glob("*.jpg"):
stale.unlink()
frames_b64: list[str] = []
stacked_b64: list[str | None] = []
for index, (image, times) in enumerate(zip(images[: len(base_idx)], sub_times), start=1):
jpeg = _jpeg_bytes(_resize_max_side(image, max_side))
(frame_dir / f"frame_{index:05d}.jpg").write_bytes(jpeg)
frames_b64.append(base64.b64encode(jpeg).decode("ascii"))
interior = [_resize_max_side(next(sub_images), max_side) for _ in times]
if not interior:
stacked_b64.append(None)
continue
composite = concat_frames_b64(interior)
(frame_dir / f"stack_{index:05d}.jpg").write_bytes(base64.b64decode(composite))
stacked_b64.append(composite)
return frames_b64, stacked_b64
def _extract_video_soundtrack(video_path: Path, work_dir: Path) -> Path:
"""Write a 16 kHz mono WAV from the video soundtrack."""
try:
from vllm.multimodal.media.audio import load_audio
except ImportError as exc: # pragma: no cover
raise SystemExit("--input-video needs vLLM audio helpers (load_audio)") from exc
audio_np, sr = load_audio(str(video_path), sr=PCM16_SAMPLE_RATE, mono=True)
work_dir.mkdir(parents=True, exist_ok=True)
audio_path = work_dir / "input_16k.wav"
write_pcm16_wav(audio_path, _to_pcm16_bytes(audio_np), sample_rate_hz=int(sr))
return audio_path
def _extract_video_input(
video_path: Path,
work_dir: Path,
*,
fps: float,
max_side: int,
stack_frames: int,
) -> tuple[Path, list[str], list[str | None]]:
"""Demux a video into duplex inputs: 16 kHz mono WAV plus JPEG frames.
Use this only when the video soundtrack is the audio source. When the
caller already has ``--input-wav``, extract frames against that WAV
duration instead so a silent clip cannot fail ``load_audio``.
"""
audio_path = _extract_video_soundtrack(video_path, work_dir)
frames_b64, stacked_b64 = _extract_video_frames(
video_path,
work_dir,
fps=fps,
max_side=max_side,
duration_s=_pcm16_wav_duration_s(audio_path),
stack_frames=stack_frames,
)
return audio_path, frames_b64, stacked_b64
def _resolve_duplex_av_inputs(
*,
input_wav: str | None,
input_video: str | None,
work_dir: Path,
fps: float,
max_side: int,
stack_frames: int = 1,
) -> tuple[str, list[str], list[str | None]]:
"""Return ``(wav_path, frames_b64, stacked_b64)``. External WAV wins over video audio."""
if input_video:
video_path = Path(input_video)
if input_wav:
frames, stacked = _extract_video_frames(
video_path,
work_dir,
fps=fps,
max_side=max_side,
duration_s=_pcm16_wav_duration_s(Path(input_wav)),
stack_frames=stack_frames,
)
return input_wav, frames, stacked
extracted_wav, frames, stacked = _extract_video_input(
video_path,
work_dir,
fps=fps,
max_side=max_side,
stack_frames=stack_frames,
)
return str(extracted_wav), frames, stacked
if not input_wav:
raise SystemExit("provide --input-wav, --input-video, or both")
return input_wav, [], []
async def run_demo(args: argparse.Namespace) -> dict[str, object]:
output_dir = Path(args.output_dir)
stream_writer = _StreamingOutputWriter(output_dir)
input_wav, video_frames, stacked_video_frames = _resolve_duplex_av_inputs(
input_wav=args.input_wav,
input_video=args.input_video,
work_dir=output_dir / "video_input",
fps=args.video_fps,
max_side=args.frame_max_side,
stack_frames=args.stack_frames,
)
input_pcm16 = read_pcm16_wav(Path(input_wav))
if not input_pcm16:
raise ValueError("input WAV has no audio")
url = build_realtime_url(
args.url,
args.model,
@@ -184,10 +434,12 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
timeout_s=args.timeout_s,
)
stream_event_cursor = len(client.events.events)
await client.stream_pcm16(
frames_sent = await client.stream_pcm16(
input_pcm16,
chunk_ms=args.chunk_ms,
realtime=not args.no_realtime_pacing,
video_frames=video_frames,
stacked_video_frames=stacked_video_frames,
)
commit_event_cursor = len(client.events.events)
stream_decision = _latest_model_decision(client.events.events, stream_event_cursor)
@@ -310,6 +562,10 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
committed_index,
),
},
"input_video": args.input_video,
"video_frames_extracted": len(video_frames),
"video_frames_sent": frames_sent,
"video_stacked_frames_extracted": sum(frame is not None for frame in stacked_video_frames),
"audio_bytes": len(audio),
"audio_chunk_count": len(stream_writer.audio_chunk_paths),
"audio_chunk_files": [str(path) for path in stream_writer.audio_chunk_paths],
@@ -356,7 +612,48 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1")
parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
parser.add_argument("--session-id")
parser.add_argument("--input-wav", required=True)
parser.add_argument(
"--input-wav",
help="Mono 16 kHz PCM16 clip to stream. Optional when --input-video carries the audio track.",
)
parser.add_argument(
"--input-video",
help=(
"Video to stream as omni duplex input. PyAV demuxes frames and "
"vLLM load_audio extracts a 16 kHz mono WAV (unless --input-wav overrides the audio)."
),
)
parser.add_argument(
"--video-fps",
type=float,
default=1.0,
help="Frame extraction rate. The wire contract carries ~1 frame per second of audio.",
)
parser.add_argument(
"--frame-max-side",
type=int,
default=0,
help=(
"Optional client-side downscale of the longest side, in pixels. "
"0 (default) keeps the capture resolution so the server "
"process_image can normalize at scale_resolution=448 from a "
"higher-quality source. The web demo never resizes on the client."
),
)
parser.add_argument(
"--stack-frames",
type=int,
default=1,
help=(
"Frames sampled per 1 s audio unit. 1 (default) sends the unit's "
"frame alone. N > 1 also samples the N-1 sub-frames inside the unit "
"and tiles them into one composite sent next to the base frame, so "
"motion (arrows, doors, digits changing) is visible at N fps while "
"the audio stays 1 s per unit. Official uses 5 for high refresh "
"rate mode; the wire carries 2 images per unit at any N because the "
"sub-frames share one composite."
),
)
parser.add_argument(
"--ref-audio",
required=True,
@@ -372,7 +669,7 @@ def parse_args() -> argparse.Namespace:
"--temperature",
type=float,
default=None,
help="Stage0 sampling temperature; response-required soft-interrupt uses 0.0.",
help="Stage 0 sampling temperature; omit to preserve the model default.",
)
parser.add_argument("--no-realtime-pacing", action="store_true")
parser.add_argument("--require-audio", action="store_true")
@@ -24,6 +24,18 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/your_deploy_config.yaml
```
To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct
weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key
`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
```
Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select
the same single-stage pipeline without `--deploy-config`.
For a 3x-GPU multi-replica layout (talker/code2wav scale-out on cuda:1,2),
use `--stage-overrides` on top of the default config:
+5 -1
View File
@@ -71,11 +71,11 @@ dev = [
"voxcpm==2.0.3", # for VoxCPM2 path-locating in tests
"torchdiffeq==0.2.5", # for covo_audio Test
"FlagEmbedding==1.4.0", # for hunyuan_image3 tests
"kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
"s3tokenizer==0.3.0", # for tts test
# Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo,
# whose top-level stepaudio2 package conflicts with this distribution.
"step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml)
"silero-vad==6.1", # Realtime server-VAD duplex E2E
]
demo = [
@@ -85,6 +85,10 @@ demo = [
"cn2an>=0.5.22", # Chinese number normalization for TTS (e.g. "1次" -> "一次")
]
server-vad = [
"silero-vad==6.1",
]
forced-aligner = [
"qwen-asr>=0.0.6", # faithful Japanese/Korean word segmentation for the TTS forced aligner
]
+47 -5
View File
@@ -10,14 +10,22 @@ checkpoint for T2V and first-frame I2V. Every output includes synchronized
48 kHz stereo audio.
The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5)
repository is not directly loadable with `--model`. It supplies the official
upsampler and LoRA sidecars used by the Full/SFT two-stage pipeline, so accept
both model licenses and authenticate before first use.
repository is not directly loadable with `--model`. It supplies the canonical
Diffusion VAE checkpoint and the official upsampler and LoRA sidecars used by
the Full/SFT two-stage pipeline, so accept both model licenses and authenticate
before first use.
The vLLM-Omni integration code is provided under the project's
[Apache-2.0 license](../../LICENSE). The LTX-2.5 model weights are separate
artifacts covered by the
[LTX-2.x Community License](https://github.com/Lightricks/LTX-2/blob/main/LICENSE.md),
not the vLLM-Omni license. Users are responsible for reviewing and complying
with the model license before downloading or using the weights.
## Choose a pipeline
| Pipeline | Mode | Default size | Schedule |
|---|---|---:|---:|
| --- | --- | ---: | ---: |
| `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps |
| `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 steps |
| `LTX2DistilledOneStagePipeline` | Distilled one-stage | 960x544 | 8 steps |
@@ -28,6 +36,40 @@ latent upsampler, and run a three-step refinement stage. Select the class with
`--model-class-name`; no `--task-type` flag is required. Supplying one initial
image selects I2V, while omitting it selects T2V.
## Video decoder
LTX-2.5 uses the canonical Native Diffusion VAE decoder by default. To opt in
to the legacy convolutional VAE, set the startup-only stage-0 model extra
`ltx2_use_conv_vae: true`. The choice applies to all requests and all four
pipeline classes above. DiffVAE loads its NATTEN kernel from Hugging Face Hub
during startup. `decode_timestep` and `decode_noise_scale` condition only the
legacy ConvVAE and have no effect when DiffVAE is selected.
For offline Python usage:
```python
omni = Omni(
model="Lightricks/LTX-2.5-Diffusers",
model_class_name="LTX2Pipeline",
stage_overrides='{"0":{"extras":{"ltx2_use_conv_vae":true}}}',
)
```
For online serving:
```bash
vllm serve Lightricks/LTX-2.5-Diffusers \
--omni \
--model-class-name LTX2Pipeline \
--stage-overrides '{"0":{"extras":{"ltx2_use_conv_vae":true}}}'
```
Both decoders are untiled by default. Set `vae_use_tiling` for memory-saving
serial tiling; DiffVAE tiles only above 80 frames or 768 pixels in either spatial
dimension. DiffVAE also supports distributed VAE decode; see the
[VAE Parallelism Guide](../../docs/user_guide/diffusion/parallelism/vae_parallelism.md).
DiffVAE is decoder-only, so I2V still uses the convolutional VAE encoder.
## Prerequisites
```bash
@@ -42,7 +84,7 @@ Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and
## Hardware
| GPU | Status | Recommended scope |
|---|---|---|
| --- | --- | --- |
| NVIDIA B300 | Verified | All four canonical pipelines |
| NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
| NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
+212
View File
@@ -0,0 +1,212 @@
# MammothModa2
> MammothModa2-Preview and MammothModa2-Dev unified understanding and generation
## Summary
- Vendor: ByteDance Research
- Models: `bytedance-research/MammothModa2-Preview`, `bytedance-research/MammothModa2-Dev`
- Tasks: Preview and Dev text-to-image (AR → DiT); Dev text/image understanding
- Mode: Offline inference
- Maintainer: Community
## When to use this recipe
Use this recipe to run MammothModa2-Preview through shared task-oriented
offline examples. Text-to-image uses the shared image example (`text_to_image.py`) instead of a model-specific script.
The generic example formats the AR prompt, drives the AR → DiT stage pipeline,
and forwards MammothModa2-specific generation parameters through the
pipeline-declared `extra_body` contract.
MammothModa2's DiT stage consumes its inputs through the multi-stage kwargs
interface (not `OmniDiffusionRequest`), so its generation knobs
(`text_guidance_scale`, `cfg_range`, `num_inference_steps`) are passed via
`--extra-body` rather than the standard `--num-inference-steps` / `--cfg-scale`
flags. Image size uses the standard `--height` / `--width` flags.
## References
- Upstream model:
[`bytedance-research/MammothModa2-Preview`](https://huggingface.co/bytedance-research/MammothModa2-Preview)
- Dev model:
[`bytedance-research/MammothModa2-Dev`](https://huggingface.co/bytedance-research/MammothModa2-Dev)
- Related offline example:
[`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
- Related T2T/I2T example:
[`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py)
- Declared parameters:
[`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)
- Deploy config:
[`vllm_omni/deploy/mammoth_moda2.yaml`](../../vllm_omni/deploy/mammoth_moda2.yaml)
## Hardware Support
The default deploy config runs both the AR and DiT stages on a single GPU
(`devices: "0"`). The committed `gpu_memory_utilization` split (stage-0 AR `0.5`,
stage-1 DiT `0.3`) is sized for an ~80 GB GPU. The model also fits on a 48 GB GPU
after rebalancing the split so the AR weights (~23 GB) leave room for the KV
cache — see the note under *1x L40S 48GB*.
## GPU
### 1x L40S 48GB
> **48 GB config adjustment:** the committed
> `vllm_omni/deploy/mammoth_moda2.yaml` uses
> `gpu_memory_utilization` 0.5 / 0.3 (sized for ~80 GB). To fit on a 48 GB L40S,
> set the stage-0 (AR) value to `0.8` and the stage-1 (DiT) value to `0.16`
> before running. (On an ~80 GB GPU, leave the defaults unchanged.)
#### Environment
- OS: Linux
- Python: Match the repository requirements for your checkout
- Driver / runtime: NVIDIA CUDA environment with one L40S 48 GB (verified) or an
~80 GB GPU for the default config
- vLLM version: Match the repository requirements for your checkout
- vLLM-Omni version or commit: Use the commit you are deploying from
#### Offline Commands
Download the model:
```bash
hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
```
Run text-to-image with the shared offline example from the repository root. The
deploy config sets `trust_remote_code`, so no extra flag is needed. Forward the
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 \
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
--height 1024 \
--width 1024 \
--extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
--output mammoth_t2i.png
```
The `--extra-body` JSON forwards MammothModa2-specific parameters into
`OmniDiffusionSamplingParams.extra_args`. Keys are filtered against the model's
declared `extra_body_params` (see
[`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)),
so unknown keys for MammothModa2 are silently dropped:
- `text_guidance_scale` — classifier-free guidance scale for the DiT stage
(default `9.0`; CFG is active only when `> 1.0`).
- `cfg_range` — relative step range `[start, end]` over which CFG is applied
(default `[0.0, 1.0]`).
- `num_inference_steps` — number of DiT denoising steps (default `50`).
`--height` and `--width` must be multiples of 16.
Run text-to-text through the shared understanding example. It recognizes the
MammothModa2 checkpoint and automatically selects `mammoth_moda2_ar.yaml`:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model ./MammothModa2-Preview \
--prompt "Explain multimodal generation in three sentences."
```
Add an image for image-to-text or image summarization. The shared example
uses MammothModa2's chat and vision-token template:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model ./MammothModa2-Preview \
--image /path/to/input.jpg \
--prompt "Please summarize the content of this image."
```
#### Verification
The example writes the generated image to the `--output` path. Confirm the file
exists and is a valid image:
```bash
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
Qwen2.5-VL. vLLM-Omni selects the matching implementation from the nested
`llm_config.model_type`; no checkpoint edits or `trust_remote_code` flag are
required.
Text-to-text and image-to-text use the AR-only deploy. Text-to-image loads the
Qwen3 generation experts (`gen_mlp`), extra visual vocabulary and image head,
then sends the generated visual tokens and hidden states to the DiT stage.
Download the checkpoint:
```bash
hf download bytedance-research/MammothModa2-Dev --local-dir ./MammothModa2-Dev
```
Run text-to-text through the shared understanding example. It recognizes the
Dev checkpoint as MammothModa2 and automatically selects
`mammoth_moda2_ar.yaml`:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model ./MammothModa2-Dev \
--prompt "Explain multimodal generation in three sentences."
```
Add an image for image-to-text or image summarization:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model ./MammothModa2-Dev \
--image ./image.png \
--prompt "Please summarize the content of this image."
```
The Dev checkpoint is approximately 47.55 GiB on disk. In the verified AR-only
run, loaded model weights used approximately 16.97 GiB of GPU memory before KV
and encoder caches. Allow additional GPU memory for those caches and the input
image.
@@ -0,0 +1,82 @@
# MiniMax H3 disaggregated text encoder
This opt-in topology runs the Qwen3-VL text encoder as a vLLM stage and sends
its hidden states and token-role metadata to an encoder-free diffusion stage.
The standard MiniMax H3 recipes remain single-stage and continue to load the
text encoder inside the diffusion pipeline.
## Start the server
Choose the topology explicitly and load its deployment defaults:
```bash
vllm-omni serve MiniMaxAI/MiniMax-H3 \
--omni \
--deploy-config vllm_omni/deploy/minimax_h3_disaggregated.yaml
```
The default deployment assigns stage 0 to GPUs 0-1 with tensor parallel size
2 and `max_num_seqs: 1`. Stage 1 uses GPUs 2-5 with diffusion tensor parallel
size 1, Ulysses degree 4, and VAE patch parallel size 4. Adjust the
`devices`, `tensor_parallel_size`, and stage 1 `parallel_config` values in a
deployment override for the available hardware. Diffusion quantization,
layerwise offload, distributed layerwise offload, VAE parallelism, and USP
settings use the same stage 1 options documented in [MiniMax-H3.md](MiniMax-H3.md).
For example, this five-GPU topology assigns one GPU to the encoder and four
to the diffusion stage. `--stage-overrides` keeps placement and parallelism
scoped to the owning stage rather than broadcasting an override to both:
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3,4 \
vllm-omni serve MiniMaxAI/MiniMax-H3 \
--omni \
--deploy-config vllm_omni/deploy/minimax_h3_disaggregated.yaml \
--stage-overrides '{"0":{"devices":"0","tensor_parallel_size":1},"1":{"devices":"1,2,3,4","tensor_parallel_size":1,"ulysses_degree":4,"vae_patch_parallel_size":4}}'
```
For memory-constrained deployments, start from the CPU-offload or distributed
layerwise-offload profiles in [MiniMax-H3.md](MiniMax-H3.md). Apply memory and
quantization options only to Stage 1 with `--stage-overrides`; retain the
encoder's BF16 configuration and the video/audio VAEs' FP32 precision. Select
one offload strategy per deployment:
```bash
# Stage 1 model-level CPU offload.
--stage-overrides '{"1":{"enable_cpu_offload":true}}'
# Stage 1 distributed layerwise offload. Tune resident layers for available RAM.
--stage-overrides '{"1":{"enable_distributed_layerwise_offload":true,"dlo_use_allgather":false,"dlo_resident_layers":20}}'
# Stage 1 online FP8 quantization of the DiT only.
--stage-overrides '{"1":{"diffusion_quantization_config":"{\"transformer\":{\"method\":\"fp8\"}}"}}'
```
The Stage 1 VAE patch-parallel options remain independent of offload and
quantization. See [MiniMax-H3.md](MiniMax-H3.md) for memory requirements and
hardware-qualified profiles before combining these options.
Stage 1 sets `model_loaded.text_encoder: false`; it must not load or download
text-encoder weights. This H3 topology explicitly keeps its single-replica
diffusion stage inline, avoiding serialization of decoded video through a
subprocess. It expects both stages to run in one deployment; cross-node payload
transport is outside this configuration.
The `/v1/videos` request schema and `extra_params.task` values (`t2va`,
`fl2va`, and `ref2va`) are unchanged from the single-stage recipe.
## Turbo LoRA
MiniMax-H3 Turbo uses five sigma points, `flow_shift=6`, and
`audio_flow_shift=3`. Start Turbo deployments with the dedicated defaults so
requests that omit sampling controls do not inherit the 50-step base schedule:
```bash
vllm-omni serve MiniMaxAI/MiniMax-H3 \
--omni \
--lora-path /path/to/MiniMax-H3-Turbo \
--deploy-config vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml
```
The base deployment intentionally retains 50 inference steps for non-LoRA
quality. Turbo LoRA supports T2VA and FL2VA requests, not Ref2VA.
+2
View File
@@ -13,6 +13,8 @@ transformers >= 5.10.1, < 5.15
av>=14.0.0
omegaconf>=2.3.0
diffusers==0.40.0
# Hub kernels, including LTX-2.5 NATTEN
kernels==0.14.1
safetensors>=0.8.0
accelerate==1.12.0
soundfile>=0.13.1
@@ -0,0 +1,143 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import hashlib
import json
import sys
from types import SimpleNamespace
import pytest
from vllm_omni.benchmarks.duplex import omni_duplex_eval_eval as eval_module
from vllm_omni.benchmarks.duplex.omni_duplex_eval_clock import (
extract_timed_sentences,
normalize_response_items,
split_text,
validate_clock,
)
from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import (
DuplexSample,
canonical_task_type,
family_for_split,
load_samples,
task_type_for_split,
)
from vllm_omni.benchmarks.duplex.omni_duplex_eval_metrics import (
PROTOCOL_PIN,
build_content_prompt,
build_reminder_prompt,
build_temporal_prompt,
parse_judge_json,
reminder_window,
summarize_pr_results,
temporal_window,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_dataset_mapping_and_manifest(tmp_path):
manifest = tmp_path / "samples.json"
manifest.write_text(json.dumps([{"id": "a", "split": "RTD_OCR", "video": "clip.mp4"}]), encoding="utf-8")
sample = load_samples(manifest, media_root=tmp_path)[0]
assert sample.family == "rtd"
assert sample.video == str(tmp_path / "clip.mp4")
assert family_for_split("PR_event_reminder") == "pr"
assert canonical_task_type("post-event-reminder") == "post_event_reminder"
assert task_type_for_split("PR_correction") == "correction"
assert load_samples([{"id": "pr", "split": "PR_correction"}])[0].task_type == "correction"
def test_hugging_face_subset_names_are_splits(monkeypatch):
calls = []
def fake_load_dataset(name, **kwargs):
calls.append((name, kwargs))
return [{"id": "pr", "question_text": "Correct this."}]
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
sample = load_samples("Hothan/Omni-DuplexEval", split="PR_correction")[0]
assert calls == [("Hothan/Omni-DuplexEval", {"split": "PR_correction"})]
assert sample.split == "PR_correction"
assert sample.task_type == "correction"
def test_response_aliases_and_clock_guard():
assert split_text("One. Two!") == ["One.", "Two!"]
assert normalize_response_items({"chunks": [{"text": "x", "current_time": 800}]}) == [
{"sentence": "x", "start": 800.0, "end": 800.0}
]
with pytest.raises(ValueError, match="clock=invalid"):
validate_clock({"clock": "invalid"})
validate_clock({"clock": "invalid"}, allow_invalid=True)
timed = extract_timed_sentences([{"type": "response.output_text.delta", "delta": "Done.", "_media_clock_ms": 800}])
assert [item.as_dict() for item in timed] == [{"sentence": "Done.", "start": 0.8, "end": 0.8}]
with pytest.raises(ValueError, match="only clock=media"):
extract_timed_sentences([], clock="wall")
def test_protocol_windows_and_parsing():
assert PROTOCOL_PIN == "ca3c122b4d4bf67afd6b18ea5e724b4561bdde48"
assert temporal_window(4, 5, 10) == (2.0, 3.0)
assert temporal_window(1, 1.2, 10) == (0.0, 0.5)
assert reminder_window(5) == (5.0, 15.0)
assert parse_judge_json('noise {"content_score": 2.345, "is_relevant": 1}') == {
"content_score": 2.35,
"is_relevant": 1,
}
assert parse_judge_json('noise "success_score": 1') == {"success_score": 1}
assert summarize_pr_results([{"task_type": "correction", "all_success": 1}])["mean_all_success"] == 1.0
def test_protocol_prompt_wording_is_pinned():
prompts = [
build_temporal_prompt(1, 2, "response", "question"),
build_content_prompt("response", "question", ["one", "two"]),
build_reminder_prompt("instruction", "response", "proactive_reminder"),
build_reminder_prompt("instruction", "response", "post_event_reminder"),
build_reminder_prompt("instruction", "response", "correction", "answer"),
]
assert [hashlib.sha256(prompt.encode()).hexdigest() for prompt in prompts] == [
"5f7f5fcb4a37e4d2eea8ba23beabf74630e6fc2fb5b2661ceb94403254bb04d5",
"1fccbed6c1ed8ea1c76d5564d7289ed6adca3678f6a0da2cc735942123d64457",
"d0a0e3a9d873d2f368c5c99b9fe46681f683cc36a9fe0e09b99d6cbca1d6dd0c",
"08f9951243c859d271f82ee53da22cd0bdab1fe93c6ca04f4870d3b959c0803a",
"3de131a7e36cb4b8ac369bb05e87ede518be20221e214df3644211c179131a73",
]
def test_frame_sample_content_passes_frames(tmp_path, monkeypatch):
response = tmp_path / "response.json"
response.write_text(json.dumps([{"sentence": "A person moves.", "start": 1, "end": 2}]), encoding="utf-8")
response.with_name("response.meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8")
video = tmp_path / "video.mp4"
video.write_bytes(b"video")
monkeypatch.setattr(eval_module, "extract_jpeg", lambda *args, **kwargs: b"\xff\xd8jpeg")
class Judge:
def temporal(self, *args, **kwargs):
return '{"temporal_score": 3, "is_relevant": 1}'
def content(self, prompt, video, frames=None, *, mode="video_url"):
assert mode == "frame-sample"
assert frames == [b"\xff\xd8jpeg", b"\xff\xd8jpeg"]
return '{"content_score": 3}'
sample = DuplexSample(
id="rtd",
split="RTD_OCR",
family="rtd",
task_type=None,
video=video,
video_duration=4,
)
score = eval_module.evaluate_sample(
sample,
response,
tmp_path / "score.json",
Judge(),
judge_video_mode="frame-sample",
)
assert score["content"]["frame_count"] == 2
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for Daily-Omni ``--daily-omni-pack-mode`` packing recipes.
The ``minicpm-interleave`` mode is an accuracy-critical protocol: MiniCPM-o only reaches
@@ -27,6 +30,8 @@ _MODULE_NAME = "vllm_omni.benchmarks.data_modules.daily_omni_dataset"
if _MODULE_NAME not in sys.modules:
_spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
if _spec is None or _spec.loader is None:
raise ImportError(f"Cannot load {_MODULE_NAME} from {_MODULE_PATH}")
_mod = importlib.util.module_from_spec(_spec)
sys.modules[_MODULE_NAME] = _mod
_spec.loader.exec_module(_mod)
@@ -87,7 +92,9 @@ def _stub_extract(monkeypatch, dataset: DailyOmniDataset, counter: list[int]) ->
def fake_extract(video_path, *, audio_path, include_audio, **_kw):
counter.append(1)
frames = [Image.new("RGB", (16, 16), color=(i, i, i)) for i in range(_NUM_FRAMES)]
segments = [np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else []
segments: list[np.ndarray] = (
[np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else []
)
return frames, segments
monkeypatch.setattr(dataset, "_extract_minicpm_frame_audio_segments", fake_extract)
+154 -3
View File
@@ -34,6 +34,7 @@ from vllm_omni.config.stage_config import (
build_stage_runtime_overrides,
load_deploy_config,
merge_pipeline_deploy,
normalize_pipeline_cli_overrides,
pipeline_cfg_resolver,
)
from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
@@ -429,6 +430,20 @@ class TestStageConfigFactory:
assert cfg["final_output"] is True
assert cfg["final_output_type"] == "image"
@pytest.mark.parametrize(
"model_class_name",
[
"MiniMaxH3Pipeline",
"WanPipeline",
"LTX2Pipeline",
"Cosmos3OmniDiffusersPipeline",
],
)
def test_default_video_diffusion_output_type(self, model_class_name):
configs = StageConfigFactory.create_default_diffusion({"model_class_name": model_class_name})
assert configs[0]["final_output_type"] == "video"
def test_default_diffusion_with_parallel_config(self):
"""Test diffusion config calculates devices from parallel_config."""
@@ -534,6 +549,7 @@ class TestPipelineDiscovery:
"""Check that specific models are in OMNI_PIPELINES."""
assert "qwen2_5_omni" in OMNI_PIPELINES
assert "qwen3_omni_moe" in OMNI_PIPELINES
assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
assert "qwen3_tts" in OMNI_PIPELINES
def test_registry_resolver_qwen3_omni_all_stages(self):
@@ -558,6 +574,18 @@ class TestPipelineDiscovery:
assert pipeline.default_deploy_config_name is None
assert len(pipeline.stages) == 1 # thinker only
def test_registry_qwen3_omni_thinker_only_key_is_static(self):
"""Explicit registry key must not re-enter the HF-config resolver."""
pipeline = resolve_pipeline_config(
"qwen3_omni_moe_thinker_only",
Q3_OMNI_ALL_STAGES_HF_CONFIG,
)
assert pipeline is OMNI_PIPELINES["qwen3_omni_moe_thinker_only"]
assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
assert len(pipeline.stages) == 1
assert pipeline.stages[0].engine_output_type == "text"
assert pipeline.stages[0].final_output_type == "text"
@pytest.mark.parametrize(
"pipeline",
[registered for registered in OMNI_PIPELINES.values() if isinstance(registered, PipelineConfig)],
@@ -1225,14 +1253,22 @@ stages:
assert deploy.connectors is not None
assert deploy.platforms is not None
def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self):
@pytest.mark.parametrize(
("hf_config", "model"),
[
(Q3_OMNI_THINKER_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Thinking"),
(Q3_OMNI_ALL_STAGES_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Instruct"),
],
ids=["thinker_checkpoint", "instruct_forced_thinker_only"],
)
def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self, hf_config, model):
deploy_path = Path(get_deploy_config_path("qwen3_omni_moe_thinking.yaml"))
with patch(
"vllm_omni.config.config_factory.get_config",
return_value=Q3_OMNI_THINKER_HF_CONFIG,
return_value=hf_config,
):
pipeline = StageConfigFactory.get_pipeline_config(
"Qwen/Qwen3-Omni-30B-A3B-Thinking",
model,
trust_remote_code=False,
deploy_config_path=str(deploy_path),
)
@@ -1241,6 +1277,7 @@ stages:
assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
deploy = load_deploy_config(deploy_path)
assert deploy.pipeline == "qwen3_omni_moe_thinker_only"
stages = merge_pipeline_deploy(pipeline, deploy)
assert len(stages) == 1
@@ -1292,6 +1329,8 @@ stages:
@pytest.mark.parametrize(
("deploy_name", "pipeline_name", "stage_count", "final_output_type", "declared_pipeline"),
[
("mammoth_moda2.yaml", "mammoth_moda2", 2, "image", "mammoth_moda2"),
("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text", "mammoth_moda2_ar"),
("omnivoice.yaml", "omnivoice", 1, "audio", "omnivoice"),
("mimo_audio.yaml", "mimo_audio", 2, "audio", None),
("step_audio_2.yaml", "step_audio_2", 2, "audio", None),
@@ -1300,6 +1339,13 @@ stages:
("step_audio2_ci.yaml", "step_audio_2_asr", 1, "text", "step_audio_2_asr"),
("hunyuan_video_15.yaml", "hunyuan_video_15", 1, "video", None),
("wan2_2_ti2v.yaml", "wan2_2_ti2v", 1, "video", None),
(
"minimax_h3_disaggregated.yaml",
"minimax_h3_disaggregated",
2,
"video",
"minimax_h3_disaggregated",
),
],
)
def test_load_new_registry_backed_deploy_configs(
@@ -1320,9 +1366,91 @@ stages:
assert len(stages) == stage_count
assert stages[-1].final_output is True
assert stages[-1].final_output_type == final_output_type
if pipeline_name == "minimax_h3_disaggregated":
assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
assert (
stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
)
assert "model_path_resolver" not in stages[0].yaml_extras
assert "model_path_resolver" not in stages[1].yaml_extras
if deploy.trust_remote_code is not None:
assert {s.yaml_engine_args.get("trust_remote_code") for s in stages} == {deploy.trust_remote_code}
def test_minimax_h3_resolvers_reach_production_engine_args(self, tmp_path: Path, monkeypatch):
from vllm_omni.engine import stage_init_utils
model_root = tmp_path / "MiniMax-H3"
(model_root / "FL2VA" / "text_encoder").mkdir(parents=True)
(model_root / "Ref2VA" / "text_encoder").mkdir(parents=True)
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None)
deploy = load_deploy_config(get_deploy_config_path("minimax_h3_disaggregated.yaml"))
assert deploy.stages[1].engine_extras["model_loaded"] == {"text_encoder": False}
stages = merge_pipeline_deploy(OMNI_PIPELINES["minimax_h3_disaggregated"], deploy)
resolved = [stage_init_utils.build_engine_args_dict(stage.to_omegaconf(), str(model_root)) for stage in stages]
assert resolved[0]["model"] == str(model_root / "FL2VA" / "text_encoder")
assert resolved[1]["model"] == str(model_root / "FL2VA")
assert all("model_path_resolver" not in engine_args for engine_args in resolved)
def test_minimax_h3_text_encoder_tp_reaches_only_stage_zero(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
stages, _ = StageConfigFactory._create_legacy_from_registry(
pipeline,
{"text_encoder_tp_size": 4},
)
assert stages[0].runtime_overrides["tensor_parallel_size"] == 4
assert "text_encoder_tp_size" not in stages[0].runtime_overrides
assert "text_encoder_tp_size" not in stages[1].runtime_overrides
assert stages[1].yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1
def test_minimax_h3_disaggregated_defaults_match_validated_topology(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml")))
stages = merge_pipeline_deploy(pipeline, deploy)
assert stages[0].yaml_engine_args["max_num_seqs"] == 1
assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
assert stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
parallel = stages[1].yaml_engine_args["parallel_config"]
assert parallel["tensor_parallel_size"] == 1
assert parallel["ulysses_degree"] == 4
assert parallel["vae_patch_parallel_size"] == 4
assert stages[1].yaml_engine_args["inline_diffusion"] is True
assert stages[1].yaml_extras["default_sampling_params"]["num_inference_steps"] == 50
turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml")))
turbo_stages = merge_pipeline_deploy(pipeline, turbo)
turbo_sampling = turbo_stages[1].yaml_extras["default_sampling_params"]
assert turbo_sampling["num_inference_steps"] == 5
assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}
def test_minimax_h3_text_encoder_tp_alias_targets_stage_zero(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
assert normalize_pipeline_cli_overrides(pipeline, {}) == {}
assert normalize_pipeline_cli_overrides(pipeline, {"text_encoder_tp_size": 4}) == {
"stage_0_tensor_parallel_size": 4
}
def test_minimax_h3_stage_zero_tp_override_wins_alias_conflict(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
with pytest.warns(UserWarning, match="stage_0_tensor_parallel_size=2 takes precedence"):
normalized = normalize_pipeline_cli_overrides(
pipeline,
{"text_encoder_tp_size": 4, "stage_0_tensor_parallel_size": 2},
)
assert normalized == {"stage_0_tensor_parallel_size": 2}
def test_minimax_h3_rejects_text_encoder_tp_on_diffusion_stage(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
with pytest.raises(ValueError, match="stage_1_text_encoder_tp_size cannot be set"):
normalize_pipeline_cli_overrides(pipeline, {"stage_1_text_encoder_tp_size": 4})
@pytest.mark.parametrize(
("config_json", "model_index", "expected_pipeline"),
[
@@ -1358,6 +1486,29 @@ stages:
assert pipeline is not None
assert pipeline.model_type == expected_pipeline
def test_minimax_h3_disaggregation_is_not_auto_discovered(self):
def get_model_file(filename, _model, revision=None):
del revision
if filename == "config.json":
return None
if filename == "model_index.json":
return {"_class_name": "MiniMaxH3Pipeline"}
return None
with (
patch.object(StageConfigFactory, "get_hf_config", return_value=None),
patch(
"vllm_omni.config.config_factory.get_hf_file_to_dict",
side_effect=get_model_file,
),
):
pipeline = StageConfigFactory.get_pipeline_config(
model="/models/unrelated-checkpoint-name",
trust_remote_code=False,
)
assert pipeline is None
@pytest.mark.parametrize("deploy_name", ["step_audio_2.yaml", "step_audio_2_async_chunk.yaml"])
def test_step_audio2_deploy_configs_fit_two_gpus(self, deploy_name):
deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
+46 -1
View File
@@ -92,10 +92,23 @@ def _from_pipeline_key(
)
def test_minimax_h3_text_encoder_tp_targets_only_structured_stage_zero() -> None:
config = _from_pipeline_key(
"minimax_h3_disaggregated",
cli_overrides={"text_encoder_tp_size": 4},
)
assert config.stage_by_id(0).parallel_config.tensor_parallel_size == 4
diffusion = config.stage_by_id(1)
assert diffusion.parallel_config.tensor_parallel_size == 1
assert diffusion.parallel_config.text_encoder_tp_size == 1
def test_duplex_session_capacity_propagates_to_every_structured_stage() -> None:
omni_config = _from_pipeline_key("personaplex")
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [2, 2]
assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["duplex", "duplex"]
def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -> None:
@@ -105,6 +118,7 @@ def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -
omni_config = _from_pipeline_key("personaplex", deploy_config_path=str(deploy_path))
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [1, 1]
assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["turn", "turn"]
@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
@@ -132,6 +146,7 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
engine_args = legacy_stage.yaml_engine_args
assert omni_stage.model_config.duplex_max_sessions == engine_args.get("duplex_max_sessions", 1)
assert omni_stage.model_config.session_mode == engine_args.get("session_mode", "turn")
assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
@@ -579,6 +594,7 @@ def test_sub_config_fields_match_structured_scopes():
"interleave_mm_strings",
"media_io_kwargs",
"active_stream_window",
"session_mode",
"duplex_max_sessions",
"enable_sleep_mode",
"default_sampling_params",
@@ -639,8 +655,8 @@ def test_sub_config_fields_match_structured_scopes():
"ulysses_mode",
"cfg_parallel_size",
"vae_patch_parallel_size",
"text_encoder_tp_size",
"vae_parallel_mode",
"text_encoder_tp_size",
"use_hsdp",
"mask_sp_padding",
"hsdp_shard_size",
@@ -936,6 +952,35 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
def test_from_pipeline_config_routes_ltx2_conv_vae_extra(tmp_path):
deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml"
deploy_path.write_text(
"\n".join(
[
"pipeline: dreamzero",
"async_chunk: false",
"stages:",
" - stage_id: 0",
" extras:",
" ltx2_use_conv_vae: true",
]
)
)
stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True
def test_stage_override_routes_ltx2_conv_vae_extra():
stage = _from_pipeline_key(
"dreamzero",
cli_overrides={"stage_0_extras": {"ltx2_use_conv_vae": True}},
).stage_by_id(0)
assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True
def test_structured_diffusion_config_rejects_invalid_compile_granularity():
with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")
+7
View File
@@ -66,3 +66,10 @@ def test_register_resolver_requires_model_type(custom_resolver, clean_pipeline_r
"""Ensure that registering a custom resolver to OMNI_PIPELINES requires an explicit model_type."""
with pytest.raises(ValueError):
register_pipeline(custom_resolver)
def test_minimax_h3_disaggregation_is_explicit_opt_in():
assert "minimax_h3" not in OMNI_PIPELINES
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
assert isinstance(pipeline, PipelineConfig)
assert pipeline.model_type == "minimax_h3_disaggregated"
@@ -25,10 +25,10 @@ from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler:
def _make_scheduler(*, stage_id: int = 0, session_mode: str = "turn") -> OmniARScheduler:
sched = OmniARScheduler.__new__(OmniARScheduler)
sched._new_prompt_len_snapshot = {}
sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id))
sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id, session_mode=session_mode))
sched.num_waiting_for_streaming_input = 0
sched.log_stats = False
sched.chunk_transfer_adapter = None
@@ -446,11 +446,15 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
return result
sched = OmniARScheduler.__new__(OmniARScheduler)
sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode="duplex"))
sched.running = []
sched.waiting = queue()
sched.skipped_waiting = queue(session)
sched.num_waiting_for_streaming_input = 1
sched.chunk_transfer_adapter = SimpleNamespace(segment_finished_requests={session.request_id})
sched.chunk_transfer_adapter = SimpleNamespace(
receives_chunks=True,
segment_finished_requests={session.request_id},
)
sched._resume_downstream_chunk_receiver(session)
sched._remove_stopped_requests_from_queues(set(), {session})
@@ -460,3 +464,37 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
assert session not in sched.skipped_waiting.requests
assert sched.num_waiting_for_streaming_input == 0
assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
@pytest.mark.parametrize(
("receives_chunks", "session_mode"),
[
(False, "duplex"),
(True, "turn"),
],
)
def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
receives_chunks: bool,
session_mode: str,
) -> None:
"""Only duplex connector-driven receivers resume without an update."""
session = _make_request()
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
sched = OmniARScheduler.__new__(OmniARScheduler)
sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode=session_mode))
sched.num_waiting_for_streaming_input = 1
sched.chunk_transfer_adapter = SimpleNamespace(
receives_chunks=receives_chunks,
segment_finished_requests={session.request_id},
)
sched.skipped_waiting = MagicMock()
sched._enqueue_waiting_request = MagicMock()
sched._resume_downstream_chunk_receiver(session)
assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
assert sched.num_waiting_for_streaming_input == 1
assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
sched.skipped_waiting.remove_requests.assert_not_called()
sched._enqueue_waiting_request.assert_not_called()
@@ -0,0 +1,106 @@
[
{
"test_name": "test_hunyuanvideo15_i2v_single_device",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 1
}
},
"full_model",
"diffusion"
],
"description": "Single-device baseline",
"server_type": "vllm-omni",
"server_params": {
"model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v",
"serve_args": {
"flow-shift": 5.0,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "832x480_frames33_steps4",
"dataset": "random",
"task": "i2v",
"num-prompts": 10,
"max-concurrency": 1,
"num-input-images": 1,
"seed": 42,
"enable-negative-prompt": true,
"random-request-config": [
{
"width": 832,
"height": 480,
"num_inference_steps": 4,
"num_frames": 33,
"fps": 24,
"weight": 1
}
],
"extra-body": {
"guidance_scale": 6.0,
"flow_shift": 5.0
}
}
]
},
{
"test_name": "test_hunyuanvideo15_i2v_cachedit_tp2_vae_patch2",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 2
}
},
"full_model",
"diffusion"
],
"description": "CacheDiT + TP=2 + VAE patch parallel=2 + VAE tiling",
"server_type": "vllm-omni",
"server_params": {
"model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v",
"serve_args": {
"flow-shift": 5.0,
"cache-backend": "cache_dit",
"tensor-parallel-size": 2,
"vae-patch-parallel-size": 2,
"vae-use-tiling": true,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "832x480_frames33_steps4",
"dataset": "random",
"task": "i2v",
"num-prompts": 10,
"max-concurrency": 1,
"num-input-images": 1,
"seed": 42,
"enable-negative-prompt": true,
"random-request-config": [
{
"width": 832,
"height": 480,
"num_inference_steps": 4,
"num_frames": 33,
"fps": 24,
"weight": 1
}
],
"extra-body": {
"guidance_scale": 6.0,
"flow_shift": 5.0
}
}
]
}
]
@@ -0,0 +1,104 @@
[
{
"test_name": "test_hunyuanvideo15_t2v_single_device",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 1
}
},
"full_model",
"diffusion"
],
"description": "Single-device baseline",
"server_type": "vllm-omni",
"server_params": {
"model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
"serve_args": {
"flow-shift": 5.0,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "832x480_frames33_steps4",
"dataset": "random",
"task": "t2v",
"num-prompts": 10,
"max-concurrency": 1,
"seed": 42,
"enable-negative-prompt": true,
"random-request-config": [
{
"width": 832,
"height": 480,
"num_inference_steps": 4,
"num_frames": 33,
"fps": 24,
"weight": 1
}
],
"extra-body": {
"guidance_scale": 6.0,
"flow_shift": 5.0
}
}
]
},
{
"test_name": "test_hunyuanvideo15_t2v_cachedit_tp2_vae_patch2",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 2
}
},
"full_model",
"diffusion"
],
"description": "CacheDiT + TP=2 + VAE patch parallel=2 + VAE tiling",
"server_type": "vllm-omni",
"server_params": {
"model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
"serve_args": {
"flow-shift": 5.0,
"cache-backend": "cache_dit",
"tensor-parallel-size": 2,
"vae-patch-parallel-size": 2,
"vae-use-tiling": true,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "832x480_frames33_steps4",
"dataset": "random",
"task": "t2v",
"num-prompts": 10,
"max-concurrency": 1,
"seed": 42,
"enable-negative-prompt": true,
"random-request-config": [
{
"width": 832,
"height": 480,
"num_inference_steps": 4,
"num_frames": 33,
"fps": 24,
"weight": 1
}
],
"extra-body": {
"guidance_scale": 6.0,
"flow_shift": 5.0
}
}
]
}
]
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
@@ -194,6 +194,47 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
@pytest.mark.cpu
@pytest.mark.core_model
@pytest.mark.parametrize(
("fully_shard_degree", "error_message"),
[(0, "fully_shard_degree must be positive"), (3, "must be divisible by fully_shard_degree")],
)
def test_invalid_hsdp_shard_size_fails_before_group_creation(monkeypatch, fully_shard_degree: int, error_message: str):
"""Reject invalid HSDP shard sizes before allocating any process groups."""
world_size = 4
created_groups: list[object] = []
def fail_if_group_created(*args, **kwargs):
del args, kwargs
created_groups.append(object())
raise AssertionError("a process group was created before HSDP validation")
fake_world_group = SimpleNamespace(
rank_in_group=0,
local_rank=0,
device_group=object(),
)
monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
monkeypatch.setattr(parallel_state, "init_model_parallel_group", fail_if_group_created)
for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_HSDP_REPLICATE", "_EXPERT_PARALLEL_GROUP_RANKS"):
monkeypatch.setattr(parallel_state, name, None)
for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
with pytest.raises(ValueError, match=error_message):
parallel_state.initialize_model_parallel(
fully_shard_degree=fully_shard_degree,
use_hsdp=True,
backend="gloo",
)
assert created_groups == []
@pytest.mark.cpu
@pytest.mark.core_model
def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
@@ -9,6 +9,7 @@ Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model.
from __future__ import annotations
import gc
import os
import time
from pathlib import Path
@@ -75,8 +76,13 @@ def _run_generate(
if num_requests < 2:
raise ValueError("num_requests must be >= 2 (1 warmup + >=1 timed)")
gc.collect()
current_omni_platform.empty_cache()
device_index = current_omni_platform.current_device()
with current_omni_platform.device(device_index):
free_bytes, total_bytes = current_omni_platform.mem_get_info()
initial_used_mb = (total_bytes - free_bytes) / (1024**2)
monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
monitor.start()
try:
@@ -113,7 +119,10 @@ def _run_generate(
pass
median_time_s = float(np.median(per_request_times_s))
peak_memory_mb = monitor.peak_used_mb
# DeviceMemoryMonitor samples absolute device usage. Normalize by
# this run's starting usage so compiler and allocator caches retained
# by the preceding TP=1 run are not charged to the TP=2 run.
peak_memory_mb = max(0.0, monitor.peak_used_mb - initial_used_mb)
return _extract_single_image([last_output]), median_time_s, peak_memory_mb
finally:
monitor.stop()
@@ -1,11 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Focused LTX-2.5 pipeline correctness tests."""
import json
from dataclasses import replace
from types import SimpleNamespace
from typing import Any
import pytest
import torch
@@ -170,7 +171,10 @@ def _resolve_request_inputs_for_test(
def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch):
pipe = SimpleNamespace(component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None))
pipe = SimpleNamespace(
component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None),
use_diffusion_decoder=True,
)
od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16)
monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
@@ -187,7 +191,19 @@ def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatc
def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
revision = "pinned-revision"
calls = {"components": []}
calls: dict[str, Any] = {"components": []}
natten_processor = object()
class FakeDiffusionDecoder:
def set_attn_processor(self, processor):
calls["diffusion_decoder_processor"] = processor
def enable_tiling(self):
calls["diffusion_decoder_tiling"] = True
def set_parallel_size(self, parallel_size, mode):
calls["diffusion_decoder_parallel"] = (parallel_size, mode)
profile = replace(
LTX25_DISTILLED_COMPONENT_PROFILE,
text_encoder_cls=object,
@@ -196,12 +212,19 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
vocoder_fallback_cls=None,
scheduler_use_dynamic_shifting=False,
)
pipeline = SimpleNamespace(component_profile=profile, pipeline_kind="distilled_two_stage")
pipeline = SimpleNamespace(
component_profile=profile,
pipeline_kind="distilled_two_stage",
use_diffusion_decoder=True,
)
od_config = SimpleNamespace(
model="org/converted-ltx25",
revision=revision,
dtype=torch.bfloat16,
quantization_config=None,
extras={},
vae_use_tiling=False,
parallel_config=SimpleNamespace(vae_patch_parallel_size=2, vae_parallel_mode="tile"),
)
monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
@@ -227,6 +250,10 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
)
return object()
def fake_native_diffusion_decoder(model, **kwargs):
calls["native_diffusion_decoder"] = (model, kwargs)
return FakeDiffusionDecoder()
def fake_transformer_config(model, subfolder, local_files_only, *, revision):
calls["transformer_config"] = (model, subfolder, local_files_only, revision)
return {"component": "transformer"}
@@ -234,7 +261,13 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch)
monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer)
monkeypatch.setattr(ltx2_components, "_load_component", fake_component)
monkeypatch.setattr(ltx2_components, "_load_ltx25_native_diffusion_decoder", fake_native_diffusion_decoder)
monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config)
monkeypatch.setattr(
ltx2_components,
"LTX2VideoVaeNeighborhoodNattenProcessor",
lambda: natten_processor,
)
monkeypatch.setattr(
ltx2_components,
"create_transformer_from_config",
@@ -261,6 +294,14 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
"latent_upsampler",
}
assert all(call[3]["revision"] == revision for call in calls["components"])
assert "diffusion_decoder" not in calls["prefetch"][1]
assert calls["native_diffusion_decoder"] == (
od_config.model,
{"local_files_only": False, "dtype": torch.bfloat16, "revision": revision},
)
assert calls["diffusion_decoder_processor"] is natten_processor
assert calls["diffusion_decoder_tiling"] is True
assert calls["diffusion_decoder_parallel"] == (2, "tile")
assert calls["transformer_config"] == (
od_config.model,
profile.transformer_subfolder,
@@ -270,6 +311,85 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
assert calls["scheduler"][1]["revision"] == revision
def test_ltx25_native_diffusion_decoder_uses_diffusers_config_and_canonical_artifact(monkeypatch):
revision = "diffusers-revision"
config = {"decoder_stage_channels": [8, 8, 8, 8, 8]}
calls: dict[str, Any] = {}
class FakeDecoder:
def init_distributed(self):
calls["init_distributed"] = True
expected = FakeDecoder()
def fake_load_config(model, **kwargs):
calls["config"] = (model, kwargs)
return config
def fake_resolve(model, repo_id, filename, **kwargs):
calls["artifact"] = (model, repo_id, filename, kwargs)
return "/models/native-diffvae.safetensors"
def fake_load(cls, path, passed_config, dtype):
calls["load"] = (path, passed_config, dtype)
return expected
monkeypatch.setattr(
ltx2_components.DistributedLTX2VideoDiffusionDecoderModel,
"load_config",
staticmethod(fake_load_config),
)
monkeypatch.setattr(ltx2_components, "resolve_ltx_artifact", fake_resolve)
monkeypatch.setattr(
ltx2_components.DistributedLTX2VideoDiffusionDecoderModel,
"from_ltx25_native_checkpoint",
classmethod(fake_load),
)
actual = ltx2_components._load_ltx25_native_diffusion_decoder(
"Lightricks/LTX-2.5-Diffusers",
local_files_only=False,
dtype=torch.bfloat16,
revision=revision,
)
assert actual is expected
assert calls["config"] == (
"Lightricks/LTX-2.5-Diffusers",
{
"subfolder": "diffusion_decoder",
"local_files_only": False,
"revision": revision,
},
)
assert calls["artifact"] == (
"Lightricks/LTX-2.5-Diffusers",
"Lightricks/LTX-2.5",
"vae/ltx-2.5-video-vae-bf16.safetensors",
{
"model_revision": revision,
"artifact_revision": ltx2_components.LTX25_NATIVE_ARTIFACT_REVISION,
},
)
assert calls["load"] == ("/models/native-diffvae.safetensors", config, torch.bfloat16)
assert calls["init_distributed"] is True
def test_ltx25_natten_failure_has_actionable_omni_remedy(monkeypatch):
class MissingNattenProcessor:
def __init__(self):
raise ImportError("kernels is unavailable")
monkeypatch.setattr(
ltx2_components,
"LTX2VideoVaeNeighborhoodNattenProcessor",
MissingNattenProcessor,
)
with pytest.raises(RuntimeError, match=r"kernels==0\.14\.1.*supported GPU.*allow Hub access"):
ltx2_components._create_ltx25_natten_processor()
def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path):
(tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"}))
(tmp_path / "transformer").mkdir()
@@ -411,6 +531,46 @@ def test_ltx25_four_public_pipeline_semantics_are_disjoint():
assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage"
@pytest.mark.parametrize(
"pipeline_cls",
[
LTX2Pipeline,
LTX2TwoStagePipeline,
LTX2DistilledOneStagePipeline,
LTX2DistilledTwoStagePipeline,
],
)
@pytest.mark.parametrize(
("extras", "use_diffusion_decoder"),
[({}, True), ({"ltx2_use_conv_vae": True}, False)],
)
def test_ltx25_all_public_pipelines_default_to_diffvae_with_conv_opt_in(
tmp_path, monkeypatch, pipeline_cls, extras, use_diffusion_decoder
):
(tmp_path / "model_index.json").write_text(
json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
)
def stub_components(pipe, od_config):
pipe.od_config = od_config
pipe.vae_spatial_compression_ratio = 32
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda _pipe: None)
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
pipe = pipeline_cls(
od_config=SimpleNamespace(
model=str(tmp_path),
extras=extras,
enable_diffusion_pipeline_profiler=False,
)
)
assert pipe.use_diffusion_decoder is use_diffusion_decoder
assert pipe._vae_modules.count("diffusion_decoder") == int(use_diffusion_decoder)
@pytest.mark.parametrize(
("pipeline_kind", "expected"),
[
@@ -673,7 +833,7 @@ def test_ltx25_distilled_two_stage_executes_custom_phase_schedules():
max_sequence_length=16,
)
prompt_context = object()
phase_calls = []
phase_calls: list[Any] = []
def resolve_request_inputs(req, **kwargs):
return request_inputs
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for LTX image-to-video input and conditioning behavior."""
@@ -258,6 +258,39 @@ class TestLTXImageToVideoConditioning:
assert LTX2Pipeline.support_image_input is True
def test_ltx25_i2v_encode_uses_diffusion_decoder_statistics(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.use_diffusion_decoder = True
pipe.vae = SimpleNamespace(
encode=lambda image: image,
latents_mean=torch.full((2,), 100.0),
latents_std=torch.full((2,), 10.0),
config=SimpleNamespace(scaling_factor=1.0),
)
pipe.diffusion_decoder = SimpleNamespace(
latents_mean=torch.tensor([1.0, 3.0]),
latents_std=torch.tensor([2.0, 4.0]),
config=SimpleNamespace(scaling_factor=2.0),
)
monkeypatch.setattr(
ltx2_conditioning,
"retrieve_latents",
lambda *_args, **_kwargs: torch.tensor([[[[[3.0]]], [[[7.0]]]]]),
)
actual = pipe._encode_i2v_image_latents(
torch.zeros(1, 3, 1, 1),
batch_size=1,
generator=None,
dtype=torch.float32,
)
torch.testing.assert_close(actual, torch.full((1, 2, 1, 1, 1), 2.0))
def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
import pytest
@@ -47,6 +50,39 @@ def test_prepare_video_latents_matches_official_values_and_token_major_layout():
assert actual.stride()[1:] == (1, actual.shape[1])
def test_prepare_video_latents_uses_diffusion_decoder_statistics():
pipeline = _make_pipeline(LTX2Pipeline)
pipeline.use_diffusion_decoder = True
pipeline.vae_spatial_compression_ratio = 1
pipeline.vae_temporal_compression_ratio = 1
pipeline.transformer_spatial_patch_size = 1
pipeline.transformer_temporal_patch_size = 1
pipeline.vae = SimpleNamespace(
latents_mean=torch.full((2,), 100.0),
latents_std=torch.full((2,), 10.0),
config=SimpleNamespace(scaling_factor=1.0),
)
pipeline.diffusion_decoder = SimpleNamespace(
latents_mean=torch.tensor([1.0, 3.0]),
latents_std=torch.tensor([2.0, 4.0]),
config=SimpleNamespace(scaling_factor=2.0),
)
actual = pipeline.prepare_latents(
batch_size=1,
num_channels_latents=2,
height=1,
width=1,
num_frames=1,
noise_scale=0.0,
dtype=torch.float32,
device=torch.device("cpu"),
latents=torch.tensor([[[[[3.0]]], [[[7.0]]]]]),
)
torch.testing.assert_close(actual, torch.full((1, 1, 2), 2.0))
def test_prepare_audio_latents_samples_directly_in_packed_token_space():
pipeline = _make_pipeline(LTX2Pipeline)
@@ -1,9 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for the fixed LTX refinement-phase adapter."""
from types import SimpleNamespace
from typing import Any
import pytest
import torch
@@ -148,10 +149,13 @@ def test_ltx_phase_adapter_selects_execution_from_quantization(
):
pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config)
manifest = object()
captured = {}
captured: dict[str, Any] = {}
def resolve_artifact(*_args, **kwargs):
assert not kwargs
assert kwargs == {
"model_revision": None,
"artifact_revision": None,
}
return "adapter"
def parse(transformer, path):
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for the shared LTX pipeline runtime and public contracts."""
@@ -8,6 +8,7 @@ import os
import tempfile
from dataclasses import replace
from types import SimpleNamespace
from typing import Any
import pytest
import torch
@@ -165,7 +166,19 @@ def test_ltx_checkpoint_version_detection_uses_metadata(tmp_path):
assert detect_ltx_model_version(str(tmp_path)) == "2"
def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch):
@pytest.mark.parametrize(
("model", "repo_id", "expected_revision"),
[
("Lightricks/test", "Lightricks/test", "model-revision"),
("Lightricks/test-Diffusers", "Lightricks/test", "artifact-revision"),
],
)
def test_ltx_artifact_uses_source_revision_and_hub_fallback(
monkeypatch,
model,
repo_id,
expected_revision,
):
filename = "ltx-sidecar.safetensors"
calls = []
@@ -176,9 +189,22 @@ def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch)
monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
assert (
resolve_ltx_artifact(str(tmp_path / "model"), "Lightricks/test", filename) == "/cache/ltx-sidecar.safetensors"
resolve_ltx_artifact(
model,
repo_id,
filename,
model_revision="model-revision",
artifact_revision="artifact-revision",
)
== "/cache/ltx-sidecar.safetensors"
)
assert calls == [{"repo_id": "Lightricks/test", "filename": filename}]
assert calls == [
{
"repo_id": repo_id,
"filename": filename,
"revision": expected_revision,
}
]
def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
@@ -187,7 +213,41 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
expected.write_bytes(b"sidecar")
monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup"))
assert resolve_ltx_artifact(str(tmp_path), "Lightricks/test", filename) == str(expected)
assert resolve_ltx_artifact(
str(tmp_path),
"Lightricks/test",
filename,
model_revision="unused-model-revision",
artifact_revision="unused-artifact-revision",
) == str(expected)
def test_ltx_artifact_local_model_missing_sidecar_falls_back_to_hub(tmp_path, monkeypatch):
calls = []
def fake_download(**kwargs):
calls.append(kwargs)
return "/cache/ltx-sidecar.safetensors"
monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
assert (
resolve_ltx_artifact(
str(tmp_path),
"Lightricks/test",
"ltx-sidecar.safetensors",
model_revision="local-model-revision",
artifact_revision="artifact-revision",
)
== "/cache/ltx-sidecar.safetensors"
)
assert calls == [
{
"repo_id": "Lightricks/test",
"filename": "ltx-sidecar.safetensors",
"revision": "artifact-revision",
}
]
@pytest.mark.parametrize(
@@ -900,7 +960,7 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
max_sequence_length=16,
)
prompt_context = object()
phase_calls = []
phase_calls: list[Any] = []
def resolve_request_inputs(req, **kwargs):
return request_inputs
@@ -959,7 +1019,7 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8
pipeline.latent_upsampler = FakeUpsampler()
activated_slots = []
activated_slots: list[Any] = []
if pipeline_cls is LTX2TwoStagePipeline:
pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append)
object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
@@ -1133,6 +1193,59 @@ def test_ltx_two_stage_recipe_rejects_sigmas():
)
def test_ltx_request_applies_diffvae_minimum_spatial_size_only_when_requested():
request = LTXRequestInputs(
prompt="prompt",
negative_prompt="",
height=128,
width=128,
num_frames=9,
frame_rate=24.0,
num_inference_steps=30,
guidance=LTX2_ONE_STAGE_RECIPE.request_guidance,
num_videos_per_prompt=1,
generator=None,
latents=None,
audio_latents=None,
prompt_embeds=None,
negative_prompt_embeds=None,
prompt_attention_mask=None,
negative_prompt_attention_mask=None,
decode_timestep=0.0,
decode_noise_scale=None,
output_type="np",
max_sequence_length=16,
)
with pytest.raises(ValueError, match="DiffVAE requires video latent spatial dimensions of at least 7x7"):
validate_pipeline_request(
request,
pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
vae_spatial_compression_ratio=32,
vae_temporal_compression_ratio=8,
pipeline_name="LTX2Pipeline",
min_video_latent_spatial_size=(7, 7),
)
# ConvVAE does not use NATTEN, so its existing 128x128 boundary remains valid.
validate_pipeline_request(
request,
pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
vae_spatial_compression_ratio=32,
vae_temporal_compression_ratio=8,
pipeline_name="LTX2Pipeline",
)
validate_pipeline_request(
replace(request, height=256, width=256),
pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
vae_spatial_compression_ratio=32,
vae_temporal_compression_ratio=8,
pipeline_name="LTX2Pipeline",
min_video_latent_spatial_size=(7, 7),
)
@pytest.mark.parametrize(
("direct_kwargs", "sampling_kwargs", "error"),
[
@@ -1284,7 +1397,7 @@ def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, step
def test_ltx_two_stage_entries_select_the_official_adapter_slots():
phase_switches = []
phase_switches: list[Any] = []
ordinary_pipeline = object.__new__(LTX2TwoStagePipeline)
ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append)
for phase in LTX2_TWO_STAGE_RECIPE.phases:
+439 -3
View File
@@ -1,9 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for LTX-2.3 VAE tiling and distributed decode behavior."""
"""Unit tests for LTX video VAE tiling and distributed decode behavior."""
from types import SimpleNamespace
from typing import Any
import pytest
import torch
@@ -36,6 +37,441 @@ def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch):
assert output.dtype == torch.bfloat16
class TestLTXDiffusionDecoder:
@pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width"])
def test_distributed_diffusion_decoder_rejects_non_tile_parallel_modes(self, mode):
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
DistributedLTX2VideoDiffusionDecoderModel,
)
model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
torch.nn.Module.__init__(model)
model.distributed_executor = SimpleNamespace(set_parallel_size=lambda *_args, **_kwargs: None)
with pytest.raises(ValueError, match="only supports vae_parallel_mode='tile'"):
model.set_parallel_size(2, mode=mode)
def test_distributed_diffusion_decoder_accepts_tile_parallel_mode(self):
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
DistributedLTX2VideoDiffusionDecoderModel,
)
calls = []
model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
torch.nn.Module.__init__(model)
model.distributed_executor = SimpleNamespace(
set_parallel_size=lambda parallel_size, mode: calls.append((parallel_size, mode))
)
model.set_parallel_size(4, mode="tile")
assert calls == [(4, "tile")]
def test_short_clip_keeps_stage5_temporal_context_then_crops_output(self):
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
LTX2VideoDiffusionDecoder3d,
LTX2VideoDiffusionDecoderModel,
)
class IdentityAttention:
def build_block_mask(self, _hidden_states):
return None
class IdentityBlock(torch.nn.Module):
def __init__(self):
super().__init__()
self.attn = IdentityAttention()
def forward(self, hidden_states, _block_mask):
return hidden_states
class IdentityUpsample(torch.nn.Module):
def forward(self, hidden_states, *, drop_leading_frame):
assert drop_leading_frame
return hidden_states
decoder = object.__new__(LTX2VideoDiffusionDecoder3d)
torch.nn.Module.__init__(decoder)
decoder.det_stages = torch.nn.ModuleList([torch.nn.ModuleList([IdentityBlock()])])
decoder.upsamples = torch.nn.ModuleList([IdentityUpsample()])
decoder.trailing_pad_latent_frames = 2
decoder.temporal_compression_ratio = 8
decoder.stage5_kernel = (11, 3, 3)
# 17 frames contain one real frame followed by 16 propagated ghost
# frames. The reference keeps 11 for stage-5 NATTEN instead of
# cropping the tensor immediately back to one.
context = decoder.forward_stage_4(torch.zeros(1, 17, 3, 3, 4))
assert context.shape == (1, 11, 3, 3, 4)
class FakeDecoder(torch.nn.Module):
def forward(self, z, *, generator, num_inference_steps):
assert generator is None
assert num_inference_steps is None
return torch.zeros(z.shape[0], 3, 11, z.shape[3] * 2, z.shape[4] * 2)
model = object.__new__(LTX2VideoDiffusionDecoderModel)
torch.nn.Module.__init__(model)
model.decoder = FakeDecoder()
model.use_tiling = False
model.tile_sample_min_num_frames = 80
model.tile_sample_min_height = 768
model.tile_sample_min_width = 768
model.temporal_compression_ratio = 8
model.spatial_compression_ratio = 2
output = model.decode(torch.zeros(1, 4, 1, 2, 3), generator=None, return_dict=False)[0]
assert output.shape == (1, 3, 1, 4, 6)
@pytest.mark.parametrize(
("model_version", "extras", "expected"),
[
("2.5", {}, True),
("2.5", {"ltx2_use_conv_vae": True}, False),
("2.3", {}, False),
],
)
def test_ltx25_diffusion_decoder_is_default_with_conv_vae_opt_in(self, model_version, extras, expected):
from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder
assert _ltx2_use_diffusion_decoder(SimpleNamespace(extras=extras), model_version) is expected
def test_conv_vae_opt_in_rejects_non_boolean_model_extra(self):
from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder
with pytest.raises(TypeError, match="ltx2_use_conv_vae"):
_ltx2_use_diffusion_decoder(SimpleNamespace(extras={"ltx2_use_conv_vae": "true"}), "2.5")
def test_native_diffusion_decoder_conversion_splits_qkv_and_folds_gates(self):
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
convert_ltx25_native_diffusion_decoder_state_dict,
)
qkv_weight = torch.arange(24, dtype=torch.float32).reshape(6, 4)
qkv_bias = torch.arange(6, dtype=torch.float32)
projection_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
projection_bias = torch.ones(2)
gate = torch.tensor([2.0, 3.0])
mean = torch.tensor([0.1, 0.2])
std = torch.tensor([0.3, 0.4])
native = {
"encoder.ignored": torch.ones(1),
"per_channel_statistics.mean-of-means": mean,
"vae.per_channel_statistics.std-of-means": std,
"decoder.type_emb": torch.ones(2),
"decoder.coarse_head.weight": torch.ones(1),
"decoder.det_stages.0.0.attn.qkv.weight": qkv_weight,
"decoder.det_stages.0.0.attn.qkv.bias": qkv_bias,
"decoder.det_stages.0.0.attn.q_norm.weight": torch.tensor([5.0, 6.0]),
"decoder.det_stages.0.0.attn.k_norm.weight": torch.tensor([7.0, 8.0]),
"decoder.diff_blocks.0.attn.proj.weight": projection_weight,
"decoder.diff_blocks.0.attn.proj.bias": projection_bias,
"decoder.diff_blocks.0.gate_msa": gate,
"vae.decoder.t_embedder.mlp.0.weight": torch.ones(2, 2),
}
converted = convert_ltx25_native_diffusion_decoder_state_dict(native)
torch.testing.assert_close(converted["latents_mean"], mean)
torch.testing.assert_close(converted["latents_std"], std)
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.weight"], qkv_weight[:2])
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.weight"], qkv_weight[2:4])
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.weight"], qkv_weight[4:])
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.bias"], qkv_bias[:2])
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.bias"], qkv_bias[2:4])
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.bias"], qkv_bias[4:])
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_q.weight"], torch.tensor([5.0, 6.0]))
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_k.weight"], torch.tensor([7.0, 8.0]))
torch.testing.assert_close(
converted["decoder.diff_blocks.0.attn.to_out.0.weight"], gate.unsqueeze(1) * projection_weight
)
torch.testing.assert_close(converted["decoder.diff_blocks.0.attn.to_out.0.bias"], gate * projection_bias)
assert "decoder.t_embedder.timestep_embedder.linear_1.weight" in converted
assert not any("type_emb" in key or "coarse" in key or "gate_msa" in key for key in converted)
def test_native_diffusion_decoder_conversion_rejects_invalid_fused_qkv(self):
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
convert_ltx25_native_diffusion_decoder_state_dict,
)
with pytest.raises(ValueError, match="not divisible by 3"):
convert_ltx25_native_diffusion_decoder_state_dict(
{"decoder.det_stages.0.0.attn.qkv.weight": torch.ones(5, 2)}
)
def test_decode_uses_diffusion_decoder_without_conv_vae_conditioning(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
seen = {}
class ConvVae:
dtype = torch.float32
config = SimpleNamespace(timestep_conditioning=True)
def decode(self, *_args, **_kwargs):
raise AssertionError("the convolutional VAE decoder must not run")
class DiffusionDecoder:
dtype = torch.float32
def decode(self, latents, *, generator, return_dict):
seen["latents"] = latents
seen["generator"] = generator
seen["return_dict"] = return_dict
return (latents + 1,)
class AudioVae:
dtype = torch.float32
def decode(self, audio_latents, *, return_dict):
assert return_dict is False
return (audio_latents + 2,)
class VideoProcessor:
def postprocess_video(self, video, *, output_type):
seen["output_type"] = output_type
return video
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.use_diffusion_decoder = True
pipe.vae = ConvVae()
pipe.diffusion_decoder = DiffusionDecoder()
pipe.audio_vae = AudioVae()
pipe.vocoder = torch.nn.Identity()
pipe.video_processor = VideoProcessor()
generator = torch.Generator().manual_seed(123)
latents = torch.ones(1, 1)
output = pipe._decode_output(
latents=latents,
audio_latents=torch.ones(1, 1),
output_type="pt",
connector_prompt_embeds=torch.ones(1, 1),
generator=generator,
device=torch.device("cpu"),
decode_timestep=0.7,
decode_noise_scale=0.5,
prompt_batch_size=1,
)
torch.testing.assert_close(seen["latents"], latents)
assert seen["generator"] is generator
assert seen["return_dict"] is False
assert seen["output_type"] == "pt"
torch.testing.assert_close(output.output[0], latents + 1)
torch.testing.assert_close(output.output[1], torch.full((1, 1), 3.0))
def test_diffusion_decode_denormalizes_with_decoder_statistics(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.use_diffusion_decoder = True
pipe.transformer_spatial_patch_size = 1
pipe.transformer_temporal_patch_size = 1
pipe.vae = SimpleNamespace(
latents_mean=torch.full((2,), 100.0),
latents_std=torch.full((2,), 10.0),
config=SimpleNamespace(scaling_factor=1.0),
)
pipe.diffusion_decoder = SimpleNamespace(
latents_mean=torch.tensor([10.0, 20.0]),
latents_std=torch.tensor([2.0, 4.0]),
config=SimpleNamespace(scaling_factor=2.0),
)
pipe.audio_vae = SimpleNamespace(
latents_mean=torch.tensor(0.0),
latents_std=torch.tensor(1.0),
)
forward_ctx = SimpleNamespace(
latent_num_frames=1,
latent_height=1,
latent_width=1,
original_audio_num_frames=1,
latent_mel_bins=2,
)
video, audio = pipe._unpack_and_denormalize_stage(
forward_ctx,
torch.tensor([[[1.0, 2.0]]]),
torch.zeros(1, 1, 2),
)
torch.testing.assert_close(video.flatten(), torch.tensor([11.0, 24.0]))
torch.testing.assert_close(audio, torch.zeros(1, 1, 1, 2))
def test_non_output_rank_does_not_enter_diffusion_decoder(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
class ConvVae:
config = SimpleNamespace(timestep_conditioning=False)
def is_distributed_enabled(self):
raise AssertionError("DiffVAE decode must not inspect ConvVAE collectives")
class DiffusionDecoder:
dtype = torch.float32
def decode(self, *_args, **_kwargs):
raise AssertionError("only the output rank may run DiffVAE decode")
monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True)
monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1)
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.use_diffusion_decoder = True
pipe.vae = ConvVae()
pipe.diffusion_decoder = DiffusionDecoder()
output = pipe._decode_output(
latents=torch.ones(1, 1),
audio_latents=torch.ones(1, 1),
output_type="pt",
connector_prompt_embeds=torch.ones(1, 1),
generator=None,
device=torch.device("cpu"),
decode_timestep=0.0,
decode_noise_scale=None,
prompt_batch_size=1,
)
assert output.output[0].numel() == 0
assert output.output[1].numel() == 0
def test_non_output_rank_enters_distributed_diffusion_decoder(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
class ConvVae:
config = SimpleNamespace(timestep_conditioning=False)
def is_distributed_enabled(self):
raise AssertionError("DiffVAE decode must query the active decoder")
class DiffusionDecoder:
dtype = torch.float32
def __init__(self):
self.decode_calls = 0
def is_distributed_enabled(self):
return True
def decode(self, *_args, **_kwargs):
self.decode_calls += 1
return (torch.empty(0, 3, 0, 0, 0),)
monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True)
monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1)
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.use_diffusion_decoder = True
pipe.vae = ConvVae()
pipe.diffusion_decoder = DiffusionDecoder()
output = pipe._decode_output(
latents=torch.ones(1, 1),
audio_latents=torch.ones(1, 1),
output_type="pt",
connector_prompt_embeds=torch.ones(1, 1),
generator=None,
device=torch.device("cpu"),
decode_timestep=0.0,
decode_noise_scale=None,
prompt_batch_size=1,
)
assert pipe.diffusion_decoder.decode_calls == 1
assert output.output[0].numel() == 0
assert output.output[1].numel() == 0
def test_diffusion_decoder_patch_parallel_size_one_uses_native_tiling(self, monkeypatch):
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
DistributedLTX2VideoDiffusionDecoderModel,
)
expected = torch.ones(1, 3, 1, 2, 2)
seen = {}
def native_tiled_decode(self, z, generator=None, num_inference_steps=None):
seen["z"] = z
seen["generator"] = generator
seen["num_inference_steps"] = num_inference_steps
return expected
monkeypatch.setattr(LTX2VideoDiffusionDecoderModel, "tiled_decode", native_tiled_decode)
model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
torch.nn.Module.__init__(model)
model.is_distributed_enabled = lambda: False
generator = torch.Generator().manual_seed(17)
z = torch.zeros(1, 1, 1, 1, 1)
result = DistributedLTX2VideoDiffusionDecoderModel.tiled_decode(
model,
z,
generator=generator,
num_inference_steps=1,
)
assert result is expected
assert seen["z"] is z
assert seen["generator"] is generator
assert seen["num_inference_steps"] == 1
def test_distributed_diffusion_tiles_preserve_serial_noise_order(self, monkeypatch):
from diffusers.utils.torch_utils import randn_tensor
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
DistributedLTX2VideoDiffusionDecoderModel,
)
model = DistributedLTX2VideoDiffusionDecoderModel(
out_channels=1,
latent_channels=1,
patch_size=1,
decoder_head_dim=8,
decoder_stage_channels=(8, 8, 8, 8, 8),
decoder_stage_depths=(1, 1, 1, 1, 1),
decoder_stage_kernels=((1, 1, 1),) * 4,
decoder_upsample_strides=((1, 1, 1),) * 4,
decoder_upsample_channel_reductions=(1, 1, 1, 1),
decoder_stage5_kernel=(1, 1, 1),
spatial_compression_ratio=1,
temporal_compression_ratio=1,
)
model.tile_sample_min_num_frames = 2
model.tile_sample_stride_num_frames = 1
model.tile_sample_min_height = 2
model.tile_sample_stride_height = 1
model.tile_sample_min_width = 2
model.tile_sample_stride_width = 1
model.distributed_executor = SimpleNamespace(group=None)
model.decoder.forward_stages_1_to_3 = lambda z: z.permute(0, 2, 3, 4, 1)
monkeypatch.setattr(torch.distributed, "broadcast", lambda *_args, **_kwargs: None)
z = torch.zeros(1, 1, 2, 3, 3)
distributed_generator = torch.Generator().manual_seed(123)
reference_generator = torch.Generator().manual_seed(123)
tasks, grid_spec = model._distributed_tile_split(z, distributed_generator, num_inference_steps=1)
assert grid_spec.grid_shape == (1, 3, 3)
assert grid_spec.split_dims == (2, 3, 4)
for task in tasks:
tile_shape = model._tiled_pixel_shape_from_features(
task.tensor,
drop_leading_frame=task.drop_leading_frame,
crop_trailing_ghost=task.crop_trailing_ghost,
)
expected = randn_tensor(tile_shape, generator=reference_generator, device=z.device, dtype=z.dtype)
actual = randn_tensor(tile_shape, generator=task.noise_generator, device=z.device, dtype=z.dtype)
torch.testing.assert_close(actual, expected)
assert torch.equal(distributed_generator.get_state(), reference_generator.get_state())
class TestLTXOutputRank:
@pytest.mark.parametrize(
("distributed_vae_state", "expected_decode_calls"),
@@ -138,7 +574,7 @@ class TestLTX23VaeDistributedDecode:
},
output_dtype=torch.float32,
)
seen = {}
seen: dict[str, Any] = {}
def exec_tile(task):
return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1))
@@ -38,6 +38,7 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
pipeline.default_audio_shift = 3.0
pipeline.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace()
pipeline.text_encoder = object()
cache_spec = object()
quality_plan = SimpleNamespace(cache_dit=cache_spec)
pipeline._quality_policy = Mock()
@@ -110,6 +111,40 @@ def test_pipeline_import_registry_and_component_discovery():
assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
def test_encoder_free_stage_skips_text_encoder_during_dlo_discovery():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.transformer = torch.nn.Linear(1, 1)
pipeline.video_vae = torch.nn.Linear(1, 1)
pipeline.audio_vae = torch.nn.Linear(1, 1)
pipeline.text_encoder = None
pipeline._dit_modules = ["transformer"]
pipeline._encoder_modules = []
pipeline._vae_modules = ["video_vae", "audio_vae"]
discovered = ModuleDiscovery.discover(pipeline)
assert discovered.dit_names == ["transformer"]
assert discovered.encoder_names == []
assert discovered.vae_names == ["video_vae", "audio_vae"]
def test_encoder_free_stage_skips_missing_component_during_dlo_registration():
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
cache = Mock()
video_vae = Mock()
audio_vae = Mock()
pipeline_module._register_dlo_component_cache(cache, None, video_vae, audio_vae)
video_vae.set_omni_component_cache.assert_called_once_with(cache)
audio_vae.set_omni_component_cache.assert_called_once_with(cache)
def _write_partition_index(path, *, partition, tasks):
path.mkdir(parents=True)
(path / "model_index.json").write_text(
@@ -575,6 +610,7 @@ def _distilled_pipeline(diffuse_calls, base_schedule_by_partition):
pipeline.default_audio_shift = 3.0
pipeline.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace()
pipeline.text_encoder = Mock()
pipeline._base_schedule_by_partition = schedules
pipeline._quality_policy = Mock()
pipeline._quality_policy.resolve.return_value = SimpleNamespace(cache_dit=None)
@@ -998,7 +1034,7 @@ def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch):
config.hidden_size,
bias=False,
)
attn_call = {}
attn_call: dict[str, object] = {}
def fake_attention(query, key, value):
attn_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape)
@@ -1020,7 +1056,7 @@ def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch):
def test_text_attention_sdpa_helper_preserves_expanded_kv_fallback(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
sdpa_call = {}
sdpa_call: dict[str, object] = {}
def fake_sdpa(query, key, value, **kwargs):
sdpa_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape, kwargs=kwargs)
@@ -1095,7 +1131,7 @@ def test_encoder_forward_forwards_video_inputs():
def test_reference_video_shape_uses_h3_adapt_shape_policy():
from vllm_omni.diffusion.models.minimax_h3.reference_video import (
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
_reference_video_shape,
)
@@ -1681,6 +1717,85 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti
assert transcode_calls[0][1]["target_frame_count"] == 209
def test_ref2va_transcode_zero_frame_count_keeps_video_stream(monkeypatch, tmp_path):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
commands: list[list[str]] = []
monkeypatch.setattr(
reference_video_module.subprocess,
"run",
lambda command, **_kwargs: commands.append(command),
)
reference_video_module._transcode_reference_video(
"input.mp4",
target_width=768,
target_height=448,
target_frame_count=0,
workdir=str(tmp_path),
duration_seconds=4.0,
)
assert "-frames:v" not in commands[0]
def test_prepared_reference_video_descriptor_round_trip(tmp_path):
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
deserialize_prepared_reference_videos,
serialize_prepared_reference_videos,
)
item = {
"original_path": str(tmp_path / "source.mp4"),
"prepared_path": str(tmp_path / "prepared.mp4"),
"input_has_audio": True,
"width": 1344,
"height": 768,
"start_time_seconds": 0.0,
"duration_seconds": 5.2,
"audio_duration_seconds": 5.2,
}
descriptor = serialize_prepared_reference_videos([item], str(tmp_path))
assert deserialize_prepared_reference_videos(descriptor) == (str(tmp_path), [item])
def test_prepared_reference_video_descriptor_rejects_unknown_fields():
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
deserialize_prepared_reference_videos,
)
with pytest.raises(ValueError, match="invalid MiniMax H3 prepared-reference-video descriptor"):
deserialize_prepared_reference_videos('{"artifact_dir":"/tmp","videos":[{"path":"x"}]}')
def test_stage_one_reuses_existing_prepared_reference_video(tmp_path):
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_reuse_prepared_reference_videos,
)
prepared_path = tmp_path / "prepared.mp4"
prepared_path.touch()
prepared = [{"prepared_path": str(prepared_path)}]
assert _reuse_prepared_reference_videos(prepared, expected_count=1) is prepared
def test_stage_one_rejects_invalid_prepared_reference_video(tmp_path):
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_reuse_prepared_reference_videos,
)
with pytest.raises(ValueError, match="count does not match"):
_reuse_prepared_reference_videos([], expected_count=1)
with pytest.raises(ValueError, match="prepared reference video is unavailable"):
_reuse_prepared_reference_videos(
[{"prepared_path": str(tmp_path / "missing.mp4")}],
expected_count=1,
)
def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
@@ -2120,7 +2235,7 @@ def test_g4_reference_image_file_format_and_size_contract(tmp_path):
def test_g4_standalone_audio_duration_and_total_duration_contract():
from vllm_omni.diffusion.models.minimax_h3.reference_video import (
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
validate_reference_audio_waveforms,
)
@@ -2195,7 +2310,7 @@ def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration():
],
)
def test_g4_reference_video_metadata_validation(field, value, message, tmp_path):
from vllm_omni.diffusion.models.minimax_h3.reference_video import (
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
_validate_reference_video_metadata,
)
@@ -254,18 +254,3 @@ def test_pipeline_resolves_transformer_component_quant_config():
assert transformer_config.ignored_layers == ignored_layers
assert _resolve_component_quant_config(component_config, "transformer") is transformer_config
assert _resolve_component_quant_config(transformer_config, "transformer") is transformer_config
def test_pipeline_strips_prequantized_text_encoder_config():
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _resolve_minimax_h3_text_encoder_quant_config
from vllm_omni.quantization import ComponentQuantizationConfig
prequantized = Mock()
prequantized.get_name.return_value = "modelopt"
online_fp8 = Mock()
online_fp8.get_name.return_value = "fp8"
component_config = ComponentQuantizationConfig({"text_encoder": prequantized})
assert _resolve_minimax_h3_text_encoder_quant_config(online_fp8) is online_fp8
assert _resolve_minimax_h3_text_encoder_quant_config(prequantized) is None
assert _resolve_minimax_h3_text_encoder_quant_config(component_config) is None
@@ -187,7 +187,6 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
"model": model_ref,
"enforce_eager": True,
"tensor_parallel_size": 2,
"text_encoder_tp_size": 2,
"vae_use_tiling": True,
}
@@ -1,13 +1,18 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
import torch
import torch.nn.functional as F
from tests.helpers.mark import hardware_marks
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
pytestmark = [
pytest.mark.core_model,
pytest.mark.diffusion,
*hardware_marks(res={"npu": "A3"}, num_cards=1),
]
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
@@ -0,0 +1,385 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from unittest.mock import Mock
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.gpu, pytest.mark.diffusion]
def _selected_operators():
from vllm_omni.diffusion.models.minimax_h3.ops.vae.dispatch import (
resolve_h3_vae_operators,
)
if not current_omni_platform.is_available():
pytest.skip("No accelerator is available")
device = current_omni_platform.get_torch_device()
operators = resolve_h3_vae_operators(device)
if operators is None:
pytest.skip("No H3 VAE operator implementation is registered for this target")
return device, operators
def _operator_set(*, supports=lambda _device: True):
from vllm_omni.diffusion.models.minimax_h3.ops.vae.dispatch import H3VAEOperatorSet
return H3VAEOperatorSet(
supports=supports,
qk_norm_rope=lambda *_args: None,
scaled_residual=lambda *_args: None,
)
def _qk_reference(x, cos, sin):
normalized = F.rms_norm(x.float(), (64,), None, 1e-5).to(x.dtype)
rotary = normalized[..., :48]
first, second = rotary.chunk(2, dim=-1)
rotated = torch.cat((-second, first), dim=-1)
return torch.cat(
(rotary * cos + rotated * sin, normalized[..., 48:]),
dim=-1,
)
def _vit_norm_input(_module, hidden_states):
return hidden_states.float()
def _same_dtype_norm_input(_module, hidden_states):
return hidden_states
def _failing_norm_input(_module, _hidden_states):
raise RuntimeError("unsupported remote normalization semantics")
@pytest.mark.parametrize(("batch", "sequence"), [(1, 1), (1, 195), (2, 1797)])
def test_h3_vae_qk_norm_rope_is_bit_exact(batch, sequence):
device, operators = _selected_operators()
torch.manual_seed(17)
qkv = torch.randn(
batch,
sequence,
32,
192,
device=device,
dtype=torch.float16,
)
q, k, _ = qkv.chunk(3, dim=-1)
cos = torch.randn(batch, sequence, 1, 48, device=device, dtype=torch.float16)
sin = torch.randn_like(cos)
expected_q = _qk_reference(q, cos, sin)
expected_k = _qk_reference(k, cos, sin)
with torch.inference_mode():
actual = operators.qk_norm_rope(q, k, (cos, sin), 1e-5)
assert actual is not None
actual_q, actual_k = actual
assert torch.equal(actual_q, expected_q)
assert torch.equal(actual_k, expected_k)
def test_h3_vae_scaled_residual_is_bit_exact():
device, operators = _selected_operators()
torch.manual_seed(29)
residual = torch.randn(195, 2048, device=device, dtype=torch.float32)
branch = torch.randn(195, 2048, device=device, dtype=torch.float16)
scale = torch.randn(2048, device=device, dtype=torch.float32)
expected = residual + branch * scale
with torch.inference_mode():
actual = operators.scaled_residual(residual, branch, scale)
assert actual is not None
assert torch.equal(actual, expected)
def test_h3_vae_exact_ops_reject_unsupported_inputs():
from vllm_omni.diffusion.models.minimax_h3.ops.vae.qk_norm_rope import (
try_qk_norm_rope_exact,
)
from vllm_omni.diffusion.models.minimax_h3.ops.vae.scaled_residual import try_scaled_residual_exact
q = torch.randn(1, 2, 32, 64, dtype=torch.float32)
cos = torch.randn(1, 2, 1, 48, dtype=torch.float32)
residual = torch.randn(2, 2048)
branch = torch.randn(2, 2048, dtype=torch.float16)
scale = torch.randn(2048)
with torch.inference_mode():
assert try_qk_norm_rope_exact(q, q, (cos, cos), 1e-5) is None
assert try_scaled_residual_exact(residual, branch, scale) is None
def _make_decoder():
class Attention(nn.Module):
def __init__(self):
super().__init__()
self.to_qkv = nn.Linear(8, 24)
self.to_out = nn.Linear(8, 8)
self.norm_q = nn.RMSNorm(8, elementwise_affine=False)
self.norm_k = nn.RMSNorm(8, elementwise_affine=False)
self.spatial_parallel = False
self.dim_head = 8
def perform_attention(self, query, _key, _value, _pack_info):
return query
def forward(self, hidden_states, rotary_pos_emb=None, pack_info=None):
return hidden_states
class FeedForward(nn.Module):
def __init__(self):
super().__init__()
self.w1 = nn.Linear(8, 32)
self.w2 = nn.Linear(16, 8)
self.use_gated = True
self.act_fn = nn.SiLU()
self._compile_forward_enabled = False
self._compile_forward_fatal = False
self._compiled_forward = None
def forward(self, hidden_states):
hidden_states = self.w1(hidden_states)
gate, hidden_states = hidden_states.chunk(2, dim=-1)
return self.w2(self.act_fn(gate) * hidden_states)
class Block(nn.Module):
def __init__(self):
super().__init__()
self.attn = Attention()
self.ff = FeedForward()
self.norm1 = nn.RMSNorm(8)
self.norm2 = nn.RMSNorm(8)
self.scale1 = nn.Parameter(torch.zeros(8))
self.scale2 = nn.Parameter(torch.zeros(8))
self.use_scale = True
def forward(self, hidden_states, rotary_pos_emb=None, pack_info=None):
return hidden_states
decoder = nn.Module()
decoder.transformer_blocks = nn.ModuleList([Block(), Block()])
decoder.proj_out = nn.Linear(8, 8)
return decoder
def test_video_vae_installs_exact_optimizations(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
class FakeModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.decoder = torch.nn.Module()
class FakeRemote(torch.nn.Module):
def __init__(self):
super().__init__()
self.model = FakeModel()
monkeypatch.setattr(
vae_module,
"_load_component_config",
lambda _path: {
"latent_channels": 1,
"latents_mean": [0.0],
"latents_std": [1.0],
},
)
monkeypatch.setattr(
vae_module,
"_load_remote_component",
lambda _path, _config: FakeRemote(),
)
install = Mock(return_value=True)
monkeypatch.setattr(vae_module, "install_h3_vae_optimizations", install)
monkeypatch.setattr(vae_module, "PinnedModuleStager", Mock())
video_vae = vae_module.MiniMaxH3VideoVAE(
"unused",
device=torch.device("cuda"),
load_device=torch.device("cpu"),
)
install.assert_called_once_with(
video_vae.model.decoder,
device=torch.device("cuda"),
)
def test_h3_vae_install_precasts_only_block_linears(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
operators = _operator_set()
monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: operators)
decoder = _make_decoder()
assert vae_ops.install_h3_vae_optimizations(
decoder,
device=torch.device("meta"),
)
for block in decoder.transformer_blocks:
assert block.attn.to_qkv.weight.dtype == torch.float16
assert block.attn.to_out.weight.dtype == torch.float16
assert block.ff.w1.weight.dtype == torch.float16
assert block.ff.w2.weight.dtype == torch.float16
assert block.forward.__func__.__name__ == "_optimized_transformer_block"
assert block.attn.forward.__func__.__name__ == "_optimized_attention"
assert block.ff.forward.__func__.__name__ == "_optimized_feed_forward"
assert block.attn._omni_qk_norm_rope is operators.qk_norm_rope
assert block._omni_scaled_residual is operators.scaled_residual
assert decoder.proj_out.weight.dtype == torch.float32
# Repeated installation is idempotent.
assert vae_ops.install_h3_vae_optimizations(
decoder,
device=torch.device("meta"),
)
def test_h3_vae_install_accepts_remote_integer_parallel_flag(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
monkeypatch.setattr(
vae_ops,
"resolve_h3_vae_operators",
lambda _device: _operator_set(),
)
decoder = _make_decoder()
for block in decoder.transformer_blocks:
# The official checkpoint loads this JSON boolean-like field as int 0.
block.attn.spatial_parallel = 0
assert vae_ops.install_h3_vae_optimizations(
decoder,
device=torch.device("meta"),
)
assert getattr(decoder, "_omni_h3_vae_optimizations_installed", False)
def test_h3_vae_swiglu_uses_post_linear_fp16_output(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
device, operators = _selected_operators()
monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: operators)
decoder = _make_decoder().to(device)
feed_forward = decoder.transformer_blocks[0].ff
reference_forward = type(feed_forward).forward
assert vae_ops.install_h3_vae_optimizations(
decoder,
device=device,
)
hidden_states = torch.randn(4, 8, device=device, dtype=torch.float32)
with torch.inference_mode(), torch.autocast(device.type, dtype=torch.float16):
expected = reference_forward(feed_forward, hidden_states)
actual = feed_forward(hidden_states)
assert expected.dtype == torch.float16
assert torch.equal(actual, expected)
def test_h3_vae_install_leaves_unsupported_target_untouched(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: None)
decoder = _make_decoder()
assert not vae_ops.install_h3_vae_optimizations(
decoder,
device=torch.device("meta"),
)
assert not hasattr(decoder, "_omni_h3_vae_optimizations_installed")
for block in decoder.transformer_blocks:
assert block.attn.to_qkv.weight.dtype == torch.float32
assert block.ff.w1.weight.dtype == torch.float32
assert block.forward.__func__.__name__ == "forward"
assert block.attn.forward.__func__.__name__ == "forward"
@pytest.mark.parametrize("norm_input", [None, _same_dtype_norm_input, _failing_norm_input])
def test_h3_vae_install_requires_fp32_attention_norm_semantics(monkeypatch, norm_input):
from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
monkeypatch.setattr(
vae_ops,
"resolve_h3_vae_operators",
lambda _device: _operator_set(),
)
decoder = _make_decoder()
forward_globals = type(decoder.transformer_blocks[0].attn).forward.__globals__
monkeypatch.setitem(forward_globals, "_vit_norm_input", norm_input)
assert not vae_ops.install_h3_vae_optimizations(
decoder,
device=torch.device("meta"),
)
assert not hasattr(decoder, "_omni_h3_vae_optimizations_installed")
assert decoder.transformer_blocks[0].attn.to_qkv.weight.dtype == torch.float32
@pytest.mark.parametrize(
"break_contract",
[
lambda block: delattr(block.attn, "spatial_parallel"),
lambda block: setattr(block.attn, "spatial_parallel", 2),
lambda block: setattr(block.attn, "perform_attention", None),
lambda block: setattr(block, "scale1", nn.Parameter(torch.zeros(7))),
lambda block: setattr(block.ff, "_compile_forward_enabled", True),
],
)
def test_h3_vae_install_rejects_incompatible_remote_contract(monkeypatch, break_contract):
from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
monkeypatch.setattr(
vae_ops,
"resolve_h3_vae_operators",
lambda _device: _operator_set(),
)
decoder = _make_decoder()
break_contract(decoder.transformer_blocks[0])
assert not vae_ops.install_h3_vae_optimizations(
decoder,
device=torch.device("meta"),
)
assert decoder.transformer_blocks[0].attn.to_qkv.weight.dtype == torch.float32
def test_h3_vae_dispatch_is_extended_by_adding_an_operator_set(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3.ops.vae import dispatch
first = _operator_set(supports=lambda _device: False)
added = _operator_set(supports=lambda device: device.type == "meta")
monkeypatch.setattr(dispatch, "H3_VAE_OPERATOR_TABLE", (first, added))
assert dispatch.resolve_h3_vae_operators(torch.device("meta")) is added
assert dispatch.resolve_h3_vae_operators(torch.device("cpu")) is None
def test_h3_vae_dispatch_selects_supported_cuda_capabilities(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3.ops.vae import dispatch
platform = Mock()
platform.is_cuda.return_value = True
platform.is_available.return_value = True
monkeypatch.setattr(dispatch, "HAS_TRITON", True)
monkeypatch.setattr(dispatch, "current_omni_platform", platform)
for capability in (90, 100, 103):
platform.get_device_capability.return_value.to_int.return_value = capability
assert dispatch.resolve_h3_vae_operators(torch.device("cuda:0")) is not None
for capability in (89, 101, 110):
platform.get_device_capability.return_value.to_int.return_value = capability
assert dispatch.resolve_h3_vae_operators(torch.device("cuda:0")) is None
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import gc
from typing import Any
import numpy as np
@@ -15,7 +16,11 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
AUDIO_MODEL: dict[str, dict[str, int | None]] = {
"stabilityai/stable-audio-open-1.0": {"cuda": 1500, "rocm": 1500},
# The inference peak includes backend workspaces as well as model weights.
# On ROCm, large MIOpen workspaces and allocator fragmentation can mask the
# resident-weight reduction, so use a conservative floor that still catches
# a disabled/no-op layerwise offloader.
"stabilityai/stable-audio-open-1.0": {"cuda": 1500, "rocm": 512},
}
IMAGE_VIDEO_MODELS: dict[str, dict[str, int | None]] = {
@@ -63,8 +68,9 @@ def run_inference(
) -> tuple[float, Any]:
current_omni_platform.empty_cache()
device_index = current_omni_platform.current_device()
monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
monitor.start()
with current_omni_platform.device(device_index):
free_bytes, total_bytes = current_omni_platform.mem_get_info()
initial_used_mb = (total_bytes - free_bytes) / (1024**2)
if model_name in AUDIO_MODEL:
params = AUDIO_MODEL_PARAMS
@@ -78,22 +84,37 @@ def run_inference(
# cache_backend="cache_dit",
**params["runner_params"],
) as runner:
# Measure steady-state inference memory, not model construction. Enabling
# layerwise offload first loads the model and then replaces each block's
# device storage with CPU-backed weights. Monitoring that transition
# captures both the original model and temporary staging allocations,
# which is not representative of layerwise-offloaded inference.
monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
current_omni_platform.reset_peak_memory_stats()
monitor.start()
# Refer to tests/e2e/offline_inference/test_wan22.py
# Use minimal settings for testing
output = runner.omni.generate(
"A cat sitting on a table",
OmniDiffusionSamplingParams(
generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
guidance_scale=1.0,
num_inference_steps=num_inference_steps,
**params["sampler_params"],
),
)
try:
# Refer to tests/e2e/offline_inference/test_wan22.py
# Use minimal settings for testing
output = runner.omni.generate(
"A cat sitting on a table",
OmniDiffusionSamplingParams(
generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
guidance_scale=1.0,
num_inference_steps=num_inference_steps,
**params["sampler_params"],
),
)
finally:
monitor.stop()
peak = monitor.peak_used_mb
monitor.stop()
# DeviceMemoryMonitor reports absolute device usage. Subtract this run's
# starting usage so process-wide compiler/workspace caches retained from a
# previous run do not make the second measurement order-dependent.
peak = max(0.0, monitor.peak_used_mb - initial_used_mb)
gc.collect()
current_omni_platform.empty_cache()
return peak, output
@@ -148,8 +169,10 @@ def test_layerwise_offload_diffusion_model(model_name: str):
assert expected_saved_memory is not None
# Verify that layerwise offloading significantly reduces memory usage
# Passes only if the actual savings exceeds the expected savings
assert layerwise_offload_peak_memory + expected_saved_memory < no_offload_peak_memory, (
# Passes only if the actual savings meets the expected savings
actual_saved_memory = no_offload_peak_memory - layerwise_offload_peak_memory
assert layerwise_offload_peak_memory + expected_saved_memory <= no_offload_peak_memory, (
f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB "
f"should be significantly less than no offload peak memory {no_offload_peak_memory} MB"
f"should be at least {expected_saved_memory} MB less than no offload peak memory "
f"{no_offload_peak_memory} MB (actual savings: {actual_saved_memory} MB)"
)
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -76,7 +76,7 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata(
result = results[0]
assert result.images == ["frame-0"]
assert result.prompt == "prompt-0"
assert result.final_output_type == "image"
assert result.final_output_type == "video"
assert result.multimodal_output == {
"metadata": {
"audio": {"sample_rate": 48000},
@@ -315,7 +315,7 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
assert result.images == ["frame-0"]
assert result.prompt == "watch and listen"
assert result.final_output_type == "image"
assert result.final_output_type == "video"
assert result.multimodal_output == {
"metadata": {
"audio": {"sample_rate": 16000},
@@ -311,8 +311,8 @@ class TestPipelineStreamingOutputToEntrypoint:
lambda *, output_format, fps, video_codec_options=None: FakeStreamingVideoEncoder(),
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
return_value="diffusion",
"vllm_omni.entrypoints.openai.serving_video_output_stream.is_video_generation_pipeline",
return_value=True,
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
@@ -0,0 +1,255 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import importlib.util
import shutil
import subprocess
import sys
import textwrap
import time
from dataclasses import dataclass
from unittest.mock import Mock
import pytest
from pytest_mock import MockerFixture
import vllm_omni.diffusion.worker.diffusion_worker as worker_module
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@dataclass(frozen=True)
class _FakeGroup:
world_size: int = 1
rank_in_group: int = 0
_GROUP_GETTERS = (
"get_dp_group",
"get_pp_group",
"get_sp_group",
"get_cfg_group",
"get_tp_group",
"get_fs_group",
"get_hsdp_replicate_group",
"get_ep_group",
)
def _patch_groups(
mocker: MockerFixture,
groups: dict[str, _FakeGroup],
) -> dict[str, Mock]:
patched_getters = {}
for getter_name in _GROUP_GETTERS:
patched_getters[getter_name] = mocker.patch.object(
worker_module,
getter_name,
return_value=groups.get(getter_name, _FakeGroup()),
)
return patched_getters
def test_setup_uses_default_name_before_model_parallel_init(
mocker: MockerFixture,
) -> None:
mocker.patch.object(
worker_module,
"model_parallel_is_initialized",
return_value=False,
)
group_getters = _patch_groups(mocker, {})
set_process_title = mocker.patch.object(worker_module, "set_process_title")
decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=True, use_hsdp=True)
set_process_title.assert_called_once_with(
name="DiffusionWorker",
prefix="vLLM-Omni",
)
decorate_logs.assert_called_once_with("DiffusionWorker")
for group_getter in group_getters.values():
group_getter.assert_not_called()
@pytest.mark.parametrize(
("groups", "enable_ep", "use_hsdp", "hsdp_replicate_size", "expected_name"),
[
({}, False, False, 1, "DiffusionWorker"),
({"get_tp_group": _FakeGroup(2, 1)}, False, False, 1, "DiffusionWorker_TP1"),
({"get_dp_group": _FakeGroup(4, 2)}, False, False, 1, "DiffusionWorker_DP2"),
({"get_pp_group": _FakeGroup(2, 1)}, False, False, 1, "DiffusionWorker_PP1"),
(
{
"get_cfg_group": _FakeGroup(2, 1),
"get_tp_group": _FakeGroup(2, 1),
},
False,
False,
1,
"DiffusionWorker_CFG1_TP1",
),
(
{
"get_dp_group": _FakeGroup(2, 1),
"get_sp_group": _FakeGroup(4, 2),
"get_tp_group": _FakeGroup(2, 1),
},
False,
False,
1,
"DiffusionWorker_DP1_SP2_TP1",
),
(
{"get_fs_group": _FakeGroup(4, 3)},
False,
True,
1,
"DiffusionWorker_FS3",
),
(
{
"get_sp_group": _FakeGroup(2, 1),
"get_fs_group": _FakeGroup(4, 3),
},
False,
True,
1,
"DiffusionWorker_SP1_FS3",
),
(
{
"get_fs_group": _FakeGroup(1, 0),
"get_hsdp_replicate_group": _FakeGroup(2, 1),
},
False,
True,
2,
"DiffusionWorker_RP1",
),
({"get_ep_group": _FakeGroup(4, 2)}, True, False, 1, "DiffusionWorker_EP2"),
({}, True, False, 1, "DiffusionWorker"),
],
ids=[
"all-singleton",
"tp-only",
"dp-only",
"pp-only",
"tp-cfg",
"dp-sp-tp",
"hsdp-only",
"hsdp-sp",
"hsdp-replicated",
"expert-parallel",
"singleton-ep",
],
)
def test_setup_uses_initialized_parallel_groups(
mocker: MockerFixture,
groups: dict[str, _FakeGroup],
enable_ep: bool,
use_hsdp: bool,
hsdp_replicate_size: int,
expected_name: str,
) -> None:
mocker.patch.object(
worker_module,
"model_parallel_is_initialized",
return_value=True,
)
group_getters = _patch_groups(mocker, groups)
set_process_title = mocker.patch.object(worker_module, "set_process_title")
decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
worker_module._setup_diffusion_worker_proc_title_and_log_prefix(
enable_ep=enable_ep,
use_hsdp=use_hsdp,
hsdp_replicate_size=hsdp_replicate_size,
)
set_process_title.assert_called_once_with(
name=expected_name,
prefix="vLLM-Omni",
)
decorate_logs.assert_called_once_with(expected_name)
if use_hsdp:
group_getters["get_fs_group"].assert_called_once_with()
if hsdp_replicate_size > 1:
group_getters["get_hsdp_replicate_group"].assert_called_once_with()
else:
group_getters["get_hsdp_replicate_group"].assert_not_called()
else:
group_getters["get_fs_group"].assert_not_called()
group_getters["get_hsdp_replicate_group"].assert_not_called()
if enable_ep:
group_getters["get_ep_group"].assert_called_once_with()
else:
group_getters["get_ep_group"].assert_not_called()
def test_missing_setproctitle_is_non_fatal(mocker: MockerFixture) -> None:
mocker.patch.object(
worker_module,
"model_parallel_is_initialized",
return_value=False,
)
mocker.patch.dict(sys.modules, {"setproctitle": None})
decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=False, use_hsdp=False)
decorate_logs.assert_called_once_with("DiffusionWorker")
@pytest.mark.skipif(
sys.platform != "linux" or shutil.which("ps") is None or importlib.util.find_spec("setproctitle") is None,
reason="requires Linux ps and setproctitle",
)
def test_process_title_is_visible_through_ps() -> None:
expected_name = "vLLM-Omni::DiffusionWorker_TP1"
child_script = textwrap.dedent(
"""
import time
from vllm.utils.system_utils import set_process_title
set_process_title(
name="DiffusionWorker_TP1",
prefix="vLLM-Omni",
)
time.sleep(30)
"""
)
process = subprocess.Popen(
[sys.executable, "-c", child_script],
stderr=subprocess.PIPE,
text=True,
)
observed_title = ""
try:
for _ in range(100):
if process.poll() is not None:
stderr = process.stderr.read() if process.stderr is not None else ""
pytest.fail(f"child process exited before title check: {stderr}")
result = subprocess.run(
["ps", "-o", "args=", "-p", str(process.pid)],
check=False,
capture_output=True,
text=True,
)
observed_title = result.stdout.strip()
if expected_name in observed_title:
break
time.sleep(0.1)
else:
pytest.fail(f"expected {expected_name!r} in process title, got {observed_title!r}")
finally:
process.terminate()
try:
process.wait(timeout=5)
except subprocess.TimeoutExpired:
process.kill()
process.wait(timeout=5)
@@ -52,6 +52,11 @@ def client(mock_engine):
c.shutdown()
def test_inline_client_implements_shared_prompt_hook_protocol(client):
assert client.prompt_transform_func is None
assert client.prompt_expand_func is None
@pytest.mark.asyncio
async def test_inline_dispatch_request_success(client, mock_engine):
mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()])
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""E2E accuracy guards against pinned official Lightricks LTX references.
@@ -43,7 +43,7 @@ LTX25_OFFICIAL_MODEL_ID = "Lightricks/LTX-2.5"
LTX25_OFFICIAL_MODEL_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1"
LTX25_OFFICIAL_COMMON_FILES = {
"text_encoder": "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors",
"video_vae": "vae/ltx-2.5-video-vae-conv-bf16.safetensors",
"video_vae": "vae/ltx-2.5-video-vae-bf16.safetensors",
"audio_vae": "vae/ltx-2.5-audio-vae-bf16.safetensors",
}
LTX25_OFFICIAL_VARIANT_FILES = {
@@ -208,6 +208,9 @@ def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") ->
model = Path(configured_model).resolve()
snapshot_revision = model.name if model.parent.name == "snapshots" else None
assert (model / transformer_subfolder).is_dir(), f"LTX-2.5 component not found: {model / transformer_subfolder}"
assert (model / "diffusion_decoder" / "config.json").is_file(), (
f"LTX-2.5 DiffVAE config not found: {model / 'diffusion_decoder' / 'config.json'}"
)
return model, configured_model, configured_revision or snapshot_revision
model_id = configured_model or LTX25_OMNI_MODEL_ID
revision = configured_revision
@@ -223,6 +226,7 @@ def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") ->
"connectors/config.json",
"connectors/diffusion_pytorch_model.safetensors.index.json",
"connectors/diffusion_pytorch_model-*.safetensors",
"diffusion_decoder/config.json",
"scheduler/*",
"text_encoder/config.json",
"text_encoder/generation_config.json",
@@ -470,7 +474,7 @@ def _assert_strict_similarity(
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@pytest.mark.parametrize("task", ["t2v", "i2v"])
def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path, task: str) -> None:
"""Compare Omni with the official runtime using the same checkpoint connector weights."""
"""Compare Omni with official canonical DiffVAE and matching connector weights."""
_require_ltx25_model_access()
artifact_parent = accuracy_artifact_root / "ltx_official"
output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_distilled_{task}")
@@ -545,6 +549,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
omni_metadata = json.loads((omni_output / "metadata.json").read_text())
assert official_metadata["official_revision"] == official_revision
assert official_metadata["pipeline"] == "DistilledPipeline"
assert official_metadata["video_vae_path"] == str(official_artifacts["video_vae"])
assert official_metadata["attention_backend"] == ATTENTION_BACKEND
assert official_metadata["connector_model"] == str(omni_model)
assert omni_metadata["model_class_name"] == "LTX2DistilledTwoStagePipeline"
@@ -594,7 +599,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
@pytest.mark.parametrize("task", ["t2v", "i2v"])
@pytest.mark.parametrize("pipeline_mode", ["full_one_stage", "full_two_stage"])
def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pipeline_mode: str) -> None:
"""Compare raw official Full/SFT weights with Omni's converted Full pipeline."""
"""Compare raw official Full/SFT and DiffVAE weights with Omni."""
_require_ltx25_model_access()
artifact_parent = accuracy_artifact_root / "ltx_official"
output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_{pipeline_mode}_{task}")
@@ -697,6 +702,7 @@ def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pi
"full_two_stage": "TI2VidTwoStagesPipeline",
}[pipeline_mode]
)
assert official_metadata["video_vae_path"] == str(official_artifacts["video_vae"])
assert official_metadata["attention_backend"] == ATTENTION_BACKEND
assert official_metadata["connector_model"] == str(omni_model)
assert omni_metadata["model_class_name"] == model_class_name
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -119,6 +119,9 @@ def _model_name(model_env_var: str = MODEL_ENV_VAR, component: str = "FL2VA") ->
def _server_args() -> list[str]:
return [
"--trust-remote-code",
# Both official-reference requests use fixed shapes. Keep regional
# compilation static so cold Inductor codegen is stable across CI nodes.
"--no-diffusion-compile-dynamic",
"--num-gpus",
"4",
"--usp",
@@ -537,7 +537,7 @@ def api_server(server_case: ServerCase, mock_async_omni, tmp_path):
cmd = OmniServeCommand()
cmd.subparser_init(subparsers)
port = get_open_port()
port = get_open_port(host="0.0.0.0")
args = parser.parse_args(["serve", server_case.served_model, "--omni", "--port", str(port)])
def run_server():
@@ -814,7 +814,7 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
ServerCase(
served_model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
stage_list=["diffusion"],
stage_configs=[{"stage_type": "diffusion"}],
stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
outputs=[_build_diffusion_video_output()],
),
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
@@ -825,7 +825,7 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
ServerCase(
served_model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
stage_list=["diffusion"],
stage_configs=[{"stage_type": "diffusion"}],
stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
outputs=[_build_diffusion_video_output()],
),
"Wan-AI/Wan2.2-I2V-A14B-Diffusers",
@@ -886,7 +886,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b
ServerCase(
served_model="MiniMaxAI/MiniMax-H3",
stage_list=["diffusion"],
stage_configs=[{"stage_type": "diffusion"}],
stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
outputs=[_build_diffusion_video_output()],
),
SamplingCase(kind=SamplingKind.VIDEO_REF2VA_IMAGE_AUDIO, sampling_params=None),
@@ -897,7 +897,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b
ServerCase(
served_model="MiniMaxAI/MiniMax-H3",
stage_list=["diffusion"],
stage_configs=[{"stage_type": "diffusion"}],
stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
outputs=[_build_diffusion_video_output()],
),
SamplingCase(kind=SamplingKind.VIDEO_REF2VA_MULTI_VIDEO, sampling_params=None),
@@ -8,6 +8,7 @@ from dataclasses import FrozenInstanceError
import pytest
from tests.helpers.stage_config import get_deploy_duplex_max_sessions
from vllm_omni.config.stage_config import (
DuplexSessionRuntimeConfig,
load_deploy_config,
@@ -66,3 +67,18 @@ def test_duplex_session_runtime_rejects_non_positive_values(tmp_path, name: str,
with pytest.raises(ValueError, match=rf"duplex_session\.{name} must be positive"):
load_deploy_config(deploy_path)
@pytest.mark.parametrize(
("deploy_yaml", "expected"),
[
("minicpmo_4_5.yaml", 4),
("minicpmo_4_5_8x4090.yaml", 1),
("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4),
],
)
def test_deploy_duplex_max_sessions_tracks_the_deploy_config(deploy_yaml: str, expected: int) -> None:
# Guards the admission probe. A capacity edit, a config that declares none,
# or an overlay inheriting one from its base must surface here rather than
# as another nightly duplex admission timeout.
assert get_deploy_duplex_max_sessions(deploy_yaml) == expected
@@ -18,6 +18,10 @@ from vllm_omni.experimental.fullduplex.client import (
from vllm_omni.experimental.fullduplex.minicpmo45.policy import (
MiniCPMO45DuplexPolicy,
)
from vllm_omni.experimental.fullduplex.video_stacking import (
concat_frames,
unit_subframe_offsets,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -418,3 +422,62 @@ def test_realtime_client_pcm16_wav_round_trip(tmp_path):
assert wav_file.getnchannels() == 1
assert wav_file.getframerate() == 16_000
assert read_pcm16_wav(path) == pcm16
def _stack_test_image(color, size=(40, 30)):
from PIL import Image
return Image.new("RGB", size, color)
def test_video_stacking_tiles_a_units_subframes_into_one_image():
frames = [_stack_test_image((255, 0, 0)), _stack_test_image((0, 255, 0))]
composite = concat_frames(frames)
# Landscape cells stack into a column because that canvas is squarer, and
# the interior seam carries the separator band that lets the model tell the
# sub-frames apart.
assert composite.size == (40, 30 * 2 + 6)
assert composite.getpixel((20, 30 + 3)) == (0, 0, 0)
def test_video_stacking_picks_the_squarest_grid():
portrait = [_stack_test_image((255, 0, 0), size=(20, 60)) for _ in range(2)]
four = [_stack_test_image((0, 0, 255)) for _ in range(4)]
# A portrait clip tiles side by side, four frames always tile 2x2.
assert concat_frames(portrait).size == (20 * 2 + 6, 60)
assert concat_frames(four).size == (40 * 2 + 6, 30 * 2 + 6)
def test_video_stacking_skips_the_subframe_that_duplicates_the_base_frame():
# Official samples stack_frames=5 as 0.2/0.4/0.6/0.8 s: offset 0 is the base
# frame, already sent as frame_list[0].
assert unit_subframe_offsets(5) == pytest.approx([0.2, 0.4, 0.6, 0.8])
assert unit_subframe_offsets(1) == []
def test_realtime_client_sends_each_units_composite_beside_its_base_frame():
sent: list[dict] = []
class _Client(RealtimeDuplexClient):
def __init__(self):
pass
async def send(self, event):
sent.append(event)
asyncio.run(
_Client().stream_pcm16(
b"\x01\x00" * (16_000 * 3),
chunk_ms=200,
realtime=False,
video_frames=["f0", "f1"],
stacked_video_frames=["s0", None],
)
)
# A composite belongs to the unit it was captured in, so it rides the same
# append as that unit's base frame; a unit without one sends the base alone.
assert [event["video_frames"] for event in sent if "video_frames" in event] == [["f0", "s0"], ["f1"]]
@@ -0,0 +1,194 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import argparse
import base64
import json
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
import pytest
from vllm_omni.benchmarks.duplex import omni_duplex_eval_runner as runner
from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import DuplexSample
from vllm_omni.benchmarks.duplex.omni_duplex_eval_judge import DuplexJudge
from vllm_omni.entrypoints.cli.benchmark import omni_duplex_eval as cli
from vllm_omni.entrypoints.cli.benchmark.omni_duplex_eval import OmniDuplexEvalSubcommand
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
def test_cli_generate_evaluate_summarize_flow(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys):
manifest = tmp_path / "manifest.json"
manifest.write_text(
json.dumps(
[
{
"id": "sample-1",
"split": "PR_correction",
"question_text": "What changed?",
"answer1": "The object moved.",
}
]
),
encoding="utf-8",
)
response_root = tmp_path / "responses"
score_root = tmp_path / "scores"
async def fake_generate(sample, *, output_root, **kwargs):
output = Path(output_root) / sample.split / f"{sample.id}.json"
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text(json.dumps([{"sentence": "It moved.", "start": 0, "end": 1}]), encoding="utf-8")
output.with_name(output.stem + ".meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8")
return output
class FakeJudge:
def __init__(self, *args, **kwargs):
pass
def chat(self, *args, **kwargs):
return '{"success_score": 1, "is_relevant": 1}'
monkeypatch.setattr(cli, "generate_sample", fake_generate)
monkeypatch.setattr(cli, "DuplexJudge", FakeJudge)
parser = argparse.ArgumentParser()
OmniDuplexEvalSubcommand.add_cli_args(parser)
common = ["--dataset", str(manifest), "--family", "pr"]
generate = parser.parse_args(
[
"generate",
*common,
"--model",
"mock",
"--ref-audio",
str(manifest),
"--response-root",
str(response_root),
"--concurrency",
"2",
]
)
OmniDuplexEvalSubcommand.cmd(generate)
evaluate = parser.parse_args(
[
"evaluate",
*common,
"--response-root",
str(response_root),
"--score-root",
str(score_root),
"--judge-model",
"mock-judge",
"--eval-workers",
"2",
]
)
OmniDuplexEvalSubcommand.cmd(evaluate)
OmniDuplexEvalSubcommand.cmd(parser.parse_args(["summarize", "--score-root", str(score_root)]))
summary = json.loads(capsys.readouterr().out)
assert summary["samples"] == 1
assert summary["pr"]["mean_all_success"] == 1.0
@pytest.mark.asyncio
async def test_generate_exercises_realtime_socket_and_media_clock(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
import websockets
received = []
async def handler(websocket):
response_started = False
async for raw in websocket:
event = json.loads(raw)
received.append(event)
if event["type"] == "session.update":
await websocket.send(json.dumps({"type": "session.created"}))
elif event["type"] == "input_audio_buffer.append" and not response_started:
response_started = True
await websocket.send(json.dumps({"type": "response.created", "response": {"id": "r1"}}))
await websocket.send(
json.dumps({"type": "response.output_text.delta", "response_id": "r1", "delta": "Done."})
)
await websocket.send(
json.dumps(
{
"type": "response.audio.delta",
"response_id": "r1",
"delta": base64.b64encode(b"\0\0" * 240).decode(),
"sample_rate_hz": 24_000,
}
)
)
elif event["type"] == "input_audio_buffer.commit":
await websocket.send(json.dumps({"type": "response.done", "response": {"id": "r1"}}))
elif event["type"] == "session.close":
await websocket.send(json.dumps({"type": "session.closed"}))
return
monkeypatch.setattr(runner, "read_audio_pcm16", lambda path: b"\0\0" * (16_000 * 8 // 10))
monkeypatch.setattr(runner, "video_duration", lambda path: 0.8)
monkeypatch.setattr(runner, "iter_jpegs", lambda *args, **kwargs: iter([(0.0, b"jpeg")]))
audio = tmp_path / "question.wav"
video = tmp_path / "video.mp4"
ref = tmp_path / "ref.wav"
for path in (audio, video, ref):
path.write_bytes(b"data")
sample = DuplexSample("sample", "PR_correction", "pr", "correction", video, audio)
async with websockets.serve(handler, "127.0.0.1", 0) as server:
port = server.sockets[0].getsockname()[1]
output = await runner.generate_sample(
sample,
url=f"ws://127.0.0.1:{port}/v1/realtime?duplex=1",
model="mock",
ref_audio=ref,
output_root=tmp_path / "responses",
)
assert json.loads(output.read_text(encoding="utf-8")) == [{"sentence": "Done.", "start": 0.8, "end": 0.8}]
meta = json.loads(output.with_name("sample.meta.json").read_text(encoding="utf-8"))
assert meta["response_done"] is True
assert meta["drain_timeout"] is None
assert any(event["type"] == "playback.ack" for event in received)
def test_judge_exercises_openai_http_schema():
requests = []
class Handler(BaseHTTPRequestHandler):
def do_POST(self):
length = int(self.headers["Content-Length"])
requests.append((self.path, self.headers["Authorization"], json.loads(self.rfile.read(length))))
payload = json.dumps({"choices": [{"message": {"content": '{"success_score": 1}'}}]}).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(payload)))
self.end_headers()
self.wfile.write(payload)
def log_message(self, *args):
pass
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
judge = DuplexJudge(f"http://127.0.0.1:{server.server_port}", "judge", api_key="token")
assert judge.chat("prompt") == '{"success_score": 1}'
finally:
server.shutdown()
thread.join()
server.server_close()
path, authorization, payload = requests[0]
assert path == "/v1/chat/completions"
assert authorization == "Bearer token"
assert payload["model"] == "judge"
assert payload["messages"] == [{"role": "user", "content": "prompt"}]
@@ -0,0 +1,204 @@
"""
End-to-end test for MammothModa2 text-to-image generation.
Verifies that the AR->DiT pipeline produces an image tensor whose pixel values
match a golden reference.
Model Hub repo id: ``bytedance-research/MammothModa2-Preview``.
Deploy config: ``get_deploy_config_path("mammoth_moda2.yaml")`` -> ``vllm_omni/deploy/mammoth_moda2.yaml``
Golden pixel file: ``tests/e2e/offline_inference/fixtures/mammoth_moda2_t2i_golden.json``
Regenerate with: ``UPDATE_GOLDEN=1 pytest tests/e2e/offline_inference/test_mammoth_moda2_expansion.py``
"""
from __future__ import annotations
import json
import os
from pathlib import Path
import pytest
import torch
from huggingface_hub import snapshot_download
from vllm.sampling_params import SamplingParams
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path
# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
_IMAGE_TOKEN_ID = 151655 # "<|image_pad|>"
_VIDEO_TOKEN_ID = 151656 # "<|video_pad|>"
_VISION_START_TOKEN_ID = 151652 # "<|vision_start|>"
_VISION_END_TOKEN_ID = 151653 # "<|vision_end|>"
_AR_PATCH_SIZE = 16
MODEL_PATH = "bytedance-research/MammothModa2-Preview"
T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG)
# Golden pixel reference file. Set UPDATE_GOLDEN=1 to regenerate.
_GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"
# Fixed sampling coordinates: (channel, row_fraction, col_fraction)
# Covers corners, centre, and mid-edges across all 3 channels.
_PIXEL_SAMPLE_COORDS = [
(0, 0.0, 0.0),
(0, 0.5, 0.5),
(0, 1.0, 1.0),
(0, 0.25, 0.75),
(1, 0.0, 1.0),
(1, 0.5, 0.0),
(1, 0.75, 0.25),
(1, 1.0, 0.5),
(2, 0.0, 0.5),
(2, 0.5, 1.0),
(2, 0.75, 0.75),
(2, 1.0, 0.0),
]
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _load_t2i_gen_config(repo_id: str) -> dict:
weights_dir = Path(snapshot_download(repo_id))
cfg_path = weights_dir / "t2i_generation_config.json"
if not cfg_path.exists():
pytest.skip(f"t2i_generation_config.json not found at {cfg_path}")
with cfg_path.open() as f:
return json.load(f)
def _format_t2i_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str:
return (
"<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n"
f"<|im_start|>user\n{user_prompt}<|im_end|>\n"
"<|im_start|>assistant\n"
f"<|image start|>{ar_width}*{ar_height}<|image token|>"
)
def _sample_pixels(img_tensor: torch.Tensor) -> list[float]:
"""Sample pixel values at fixed fractional coordinates from a (C, H, W) tensor."""
t = img_tensor.float().clamp(0.0, 1.0)
if t.ndim == 4:
t = t[0] # unbatch
C, H, W = t.shape
values = []
for c, rh, rw in _PIXEL_SAMPLE_COORDS:
ri = min(int(rh * (H - 1)), H - 1)
ci = min(int(rw * (W - 1)), W - 1)
values.append(round(float(t[c, ri, ci]), 6))
return values
# ---------------------------------------------------------------------------
# End-to-end test
# ---------------------------------------------------------------------------
pytestmark = [
pytest.mark.slow,
pytest.mark.diffusion,
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
]
@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3201")
@hardware_test(res={"cuda": "H100"})
def test_mammothmoda2_t2i_e2e(omni_runner: OmniRunner):
"""
End-to-end text-to-image generation with MammothModa2 (AR -> DiT).
Verifies:
- Omni pipeline initialises with the two-stage YAML config.
- DiT stage outputs an image tensor with the correct shape.
- A fixed set of pixel values matches a golden reference
(regenerate with ``UPDATE_GOLDEN=1``).
"""
gen_cfg = _load_t2i_gen_config(MODEL_PATH)
eol_token_id = int(gen_cfg["eol_token_id"])
visual_start = int(gen_cfg["visual_token_start_id"])
visual_end = int(gen_cfg["visual_token_end_id"])
height, width = 256, 256 # small for CI speed
ar_height, ar_width = height // _AR_PATCH_SIZE, width // _AR_PATCH_SIZE
expected_grid_tokens = ar_height * (ar_width + 1)
prompt_text = "A cat sitting on a laptop keyboard"
formatted_prompt = _format_t2i_prompt(prompt_text, ar_width, ar_height)
omni = omni_runner.omni
ar_sampling = SamplingParams(
temperature=0.0,
top_k=1,
max_tokens=max(1, expected_grid_tokens + 1),
detokenize=False,
)
dit_sampling = SamplingParams(temperature=0.0, max_tokens=1, detokenize=False)
outputs = list(
omni.generate(
[
{
"prompt": formatted_prompt,
"additional_information": {
"omni_task": ["t2i"],
"ar_width": [ar_width],
"ar_height": [ar_height],
"eol_token_id": [eol_token_id],
"visual_token_start_id": [visual_start],
"visual_token_end_id": [visual_end],
"image_height": [height],
"image_width": [width],
"num_inference_steps": [2],
"text_guidance_scale": [1.0],
"cfg_range": [0.0, 1.0],
"visual_ids": [
_IMAGE_TOKEN_ID,
_VIDEO_TOKEN_ID,
_VISION_START_TOKEN_ID,
_VISION_END_TOKEN_ID,
],
},
}
],
[ar_sampling, dit_sampling],
)
)
assert len(outputs) > 0, "Pipeline produced no outputs"
found_image = False
for out in outputs:
ro_list = out
if not isinstance(ro_list, list):
ro_list = [ro_list]
for ro in ro_list:
completion_outputs = getattr(ro, "outputs", None)
if not isinstance(completion_outputs, list):
continue
for completion in completion_outputs:
mm = getattr(completion, "multimodal_output", None)
if not (isinstance(mm, dict) and "image" in mm):
continue
img_list = mm["image"] if isinstance(mm["image"], list) else [mm["image"]]
for img_tensor in img_list:
assert isinstance(img_tensor, torch.Tensor), f"Expected image tensor, got {type(img_tensor)}"
assert img_tensor.ndim in (3, 4), f"Expected 3D or 4D image tensor, got {img_tensor.ndim}D"
sampled = _sample_pixels(img_tensor)
if os.environ.get("UPDATE_GOLDEN"):
_GOLDEN_T2I_PATH.parent.mkdir(parents=True, exist_ok=True)
_GOLDEN_T2I_PATH.write_text(json.dumps({"pixels": sampled}, indent=2))
print(f"\nGolden file written to {_GOLDEN_T2I_PATH}")
elif _GOLDEN_T2I_PATH.exists():
golden = json.loads(_GOLDEN_T2I_PATH.read_text())["pixels"]
for i, (got, exp) in enumerate(zip(sampled, golden)):
assert abs(got - exp) < 1e-4, f"Pixel {i} mismatch: got {got}, expected {exp}"
found_image = True
assert found_image, "No image tensor found in pipeline output"
@@ -15,10 +15,11 @@ import pytest
from huggingface_hub import snapshot_download
from tests.helpers.runtime import OmniServerParams, get_model_prefix
from tests.helpers.stage_config import get_deploy_config_path
from tests.helpers.stage_config import get_deploy_config_path, get_deploy_duplex_max_sessions
MODEL = "openbmb/MiniCPM-o-4_5"
DEPLOY_CONFIG = get_deploy_config_path("minicpmo_4_5.yaml")
DEPLOY_CONFIG_REL = "minicpmo_4_5.yaml"
DEPLOY_CONFIG = get_deploy_config_path(DEPLOY_CONFIG_REL)
ASSET_DIR = Path(__file__).resolve().parents[3] / "assets" / "minicpmo_4_5"
RESPONSE_REQUIRED_WAV = ASSET_DIR / "response_required_16k.wav"
RESPONSE_REQUIRED_SHA256 = "2e5fd4eb3ee434ce107ee3a0591fa624a33f7683c7462f45fe651c443c9af941"
@@ -39,6 +40,16 @@ SERVER_PARAMS = [
]
def deploy_max_sessions() -> int:
"""Concurrent duplex sessions the deploy config under test admits.
The admission probe has to expect the capacity the server is actually
started with. Hardcoding it silently drifts the moment the deploy config
changes its capacity, turning that change into an unrelated probe timeout.
"""
return get_deploy_duplex_max_sessions(DEPLOY_CONFIG_REL)
def validated_wav(path: Path, expected_sha256: str) -> Path:
actual_sha256 = hashlib.sha256(path.read_bytes()).hexdigest()
if actual_sha256 != expected_sha256:
@@ -80,6 +91,19 @@ def resolve_ref_audio() -> Path:
return ref_audio
def duplex_camera_frames(*, seconds: int, cache_dir: Path) -> list[str]:
"""Base64 JPEG frames of a moving synthetic scene, one per second.
Mirrors what the realtime web client captures (1 fps, no client-side
resize: the server normalizes at ``scale_resolution=448``).
"""
from tests.e2e.online_serving.helpers.minicpmo_realtime_duplex_scenarios import _video_frames_from_file
from tests.helpers.media import generate_synthetic_video
video = generate_synthetic_video(448, 448, 30 * seconds, cache_dir=cache_dir)
return _video_frames_from_file(Path(video["file_path"]))
def realtime_url(omni_server) -> str:
return f"ws://{omni_server.host}:{omni_server.port}/v1/realtime?duplex=1"
@@ -102,6 +126,8 @@ def demo_args(
output_audio_format="pcm16",
chunk_ms=200,
realtime_input=True,
input_video=None,
stack_frames=1,
first_turn_ms=1400,
turn_duration_ms=[],
first_turn_transcript="duplex CI speech",
@@ -37,8 +37,10 @@ import json
import sys
import time
import wave
from collections.abc import Sequence
from dataclasses import dataclass, field
from pathlib import Path
from types import SimpleNamespace
_REPO_ROOT = Path(__file__).resolve().parents[4]
if str(_REPO_ROOT) not in sys.path:
@@ -55,15 +57,100 @@ from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
PCM16_SAMPLE_RATE,
RealtimeEventCollector,
build_realtime_url,
duplex_unit_boundary_ms,
read_pcm16_wav,
summarize_session_request_metrics,
)
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
reference_audio_data_url as _ref_audio_data_url,
)
from vllm_omni.experimental.fullduplex.video_stacking import ( # noqa: E402
concat_frames_b64,
unit_subframe_offsets,
)
_url_with_model = build_realtime_url
_read_wav_pcm16 = read_pcm16_wav
DemoArgs = argparse.Namespace | SimpleNamespace
def _video_frames_from_file(
path: Path,
*,
fps: float = 1.0,
jpeg_quality: int = 70,
stack_frames: int = 1,
) -> tuple[list[str], list[str | None]]:
"""Decode ``path`` into base64 JPEG frames at the omni duplex cadence.
One frame per second matches what the realtime web client captures, so the
stream stays representative of a live camera. ``stack_frames > 1`` also
samples each unit's interior sub-frames and tiles them into one composite,
the official way to raise visual refresh rate without changing the audio
cadence. Returns ``(base_frames, stacked_frames)`` with the tracks parallel.
"""
import cv2
offsets = unit_subframe_offsets(stack_frames)
capture = cv2.VideoCapture(str(path))
if not capture.isOpened():
raise ValueError(f"cannot open video: {path}")
def encode(frame_bgr, index: int) -> str:
ok, jpeg = cv2.imencode(".jpg", frame_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality])
if not ok:
raise ValueError(f"cannot JPEG-encode frame {index} of {path}")
return base64.b64encode(jpeg.tobytes()).decode("ascii")
try:
source_fps = capture.get(cv2.CAP_PROP_FPS) or 0.0
stride = max(1, round(source_fps / fps)) if source_fps > 0 else 1
# Sub-frame k of a unit sits ``offset * stride`` frames past its base.
sub_strides = {max(1, round(offset * stride)): position for position, offset in enumerate(offsets)}
frames: list[str] = []
interiors: list[list[str]] = []
index = 0
while True:
ok, frame_bgr = capture.read()
if not ok:
break
position, within_unit = divmod(index, stride)
if within_unit == 0:
frames.append(encode(frame_bgr, index))
interiors.append([])
elif within_unit in sub_strides and position < len(interiors):
interiors[position].append(encode(frame_bgr, index))
index += 1
finally:
capture.release()
if not frames:
raise ValueError(f"video has no decodable frames: {path}")
stacked = [concat_frames_b64(interior) if interior else None for interior in interiors]
return frames, stacked
def _resolve_video_frames(args: DemoArgs) -> tuple[list[str], list[str | None]]:
"""Base64 JPEG frames for the duplex camera track, newest last.
Returns the base track plus the parallel stacked composites. Only a video
source has sub-frames to stack; an explicit frame list or a still image
stacks nothing.
"""
frames = list(getattr(args, "video_frames_b64", None) or [])
if frames:
stacked = list(getattr(args, "video_stacked_frames_b64", None) or [])
stacked += [None] * (len(frames) - len(stacked))
return frames, stacked[: len(frames)]
input_video = getattr(args, "input_video", None)
if input_video:
return _video_frames_from_file(
Path(input_video).expanduser(),
stack_frames=max(1, int(getattr(args, "stack_frames", 1) or 1)),
)
frame_image = getattr(args, "frame_image", None)
if frame_image:
return [base64.b64encode(Path(frame_image).expanduser().read_bytes()).decode("ascii")], [None]
return [], []
@dataclass
@@ -373,17 +460,20 @@ class DemoState:
commit_indices = [
index for index, event in enumerate(self.events) if event.get("type") == "input_audio_buffer.committed"
]
decision_indices = [
index
for index, event in enumerate(self.events)
if event.get("type") == "response.created"
or (
event.get("type") == "response.listen"
and isinstance(event.get("response"), dict)
and isinstance(event["response"].get("metadata"), dict)
and event["response"]["metadata"].get("model_listen") is True
)
]
decision_indices: list[int] = []
for index, event in enumerate(self.events):
event_type = event.get("type")
if event_type == "response.created":
decision_indices.append(index)
continue
if event_type != "response.listen":
continue
response = event.get("response")
if not isinstance(response, dict):
continue
metadata = response.get("metadata")
if isinstance(metadata, dict) and metadata.get("model_listen") is True:
decision_indices.append(index)
first_input_index = self.first_index("input_audio_buffer.speech_started")
if first_input_index is None:
first_input_index = commit_indices[0] if commit_indices else None
@@ -500,7 +590,7 @@ class DemoState:
return stale
def _session_update_event(args: argparse.Namespace) -> dict[str, object]:
def _session_update_event(args: DemoArgs) -> dict[str, object]:
session_payload: dict[str, object] = {
"model": args.model,
"modalities": ["audio", "text"],
@@ -836,7 +926,8 @@ async def _send_pcm16(
hints: dict[str, object] | None = None,
first_chunk_hints: dict[str, object] | None = None,
on_model_unit_ready=None,
frame_b64: str | None = None,
frames_b64: Sequence[str] | None = None,
stacked_frames_b64: Sequence[str | None] | None = None,
) -> None:
hints = hints or {}
first_chunk_hints = first_chunk_hints or {}
@@ -844,6 +935,8 @@ async def _send_pcm16(
model_unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
audio_ms = 0
next_model_unit_bytes = model_unit_bytes
frames = list(frames_b64 or [])
stacked = list(stacked_frames_b64 or [])
frames_sent = 0
for offset in range(0, len(pcm16), chunk_bytes):
chunk = pcm16[offset : offset + chunk_bytes]
@@ -852,9 +945,18 @@ async def _send_pcm16(
chunk_hints = dict(hints)
if offset == 0:
chunk_hints.update(first_chunk_hints)
if frame_b64 is not None and audio_ms > frames_sent * 1000:
# Omni duplex cadence: one camera frame per 1 s of audio.
chunk_hints["video_frames"] = [frame_b64]
if frames and audio_ms >= duplex_unit_boundary_ms(frames_sent):
# Omni duplex cadence: one camera frame per model unit, riding the
# append that closes it, so the frame and the second of audio it was
# captured during enter the same unit. A video advances one frame per
# unit and holds its last frame when the audio outlives the clip; a
# still image is a one-element list and therefore repeats.
# A stacked composite of that unit's interior sub-frames rides the
# same append right after the base frame (wire maximum of 2).
index = min(frames_sent, len(frames) - 1)
composite = stacked[index] if index < len(stacked) else None
base = frames[index]
chunk_hints["video_frames"] = [base] if composite is None else [base, composite]
frames_sent += 1
await ws.send(
json.dumps(
@@ -935,7 +1037,8 @@ async def _send_clean_turn(
realtime_input: bool = False,
model_policy_settle_s: float = 2.0,
commit_input: bool = True,
frame_b64: str | None = None,
frames_b64: Sequence[str] | None = None,
stacked_frames_b64: Sequence[str | None] | None = None,
) -> tuple[str | None, str]:
before_created = state.count("response.created")
before_model_listen = state.model_listen_count
@@ -957,7 +1060,8 @@ async def _send_clean_turn(
chunk_ms=chunk_ms,
realtime_delay=realtime_input,
hints={"transcript": transcript} if send_transcript_hint else {},
frame_b64=frame_b64,
frames_b64=frames_b64,
stacked_frames_b64=stacked_frames_b64,
)
if commit_input:
state.input_commit_sent_at_s.append(time.monotonic())
@@ -1163,7 +1267,8 @@ async def _send_listen_only_overlap_pair(
model_policy_settle_s: float,
validation_mode: str,
silence_ms: int,
frame_b64: str | None = None,
frames_b64: Sequence[str] | None = None,
stacked_frames_b64: Sequence[str | None] | None = None,
) -> tuple[list[str | None], list[str], bool]:
if not realtime_input:
raise ValueError("listen-only-overlap requires --realtime-input")
@@ -1194,7 +1299,8 @@ async def _send_listen_only_overlap_pair(
realtime_delay=True,
hints={"transcript": transcripts[0]} if send_transcript_hint else {},
on_model_unit_ready=continue_until_first_speak,
frame_b64=frame_b64,
frames_b64=frames_b64,
stacked_frames_b64=stacked_frames_b64,
)
state.input_commit_sent_at_s.append(time.monotonic())
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
@@ -1233,7 +1339,8 @@ async def _send_listen_only_overlap_pair(
realtime_delay=True,
hints={"transcript": transcripts[1]} if send_transcript_hint else {},
on_model_unit_ready=record_model_unit_ready,
frame_b64=frame_b64,
frames_b64=frames_b64,
stacked_frames_b64=stacked_frames_b64,
)
state.input_commit_sent_at_s.append(time.monotonic())
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
@@ -1297,10 +1404,8 @@ async def _send_listen_only_overlap_pair(
return [first_response_id, second_response_id], ["speak", second_outcome], overlap_ok
async def run_demo(args: argparse.Namespace) -> dict[str, object]:
demo_frame_b64: str | None = None
if getattr(args, "frame_image", None):
demo_frame_b64 = base64.b64encode(Path(args.frame_image).read_bytes()).decode("ascii")
async def run_demo(args: DemoArgs) -> dict[str, object]:
demo_frames_b64, demo_stacked_b64 = _resolve_video_frames(args)
turn_input_paths = _turn_input_paths(
Path(args.input_wav),
list(getattr(args, "turn_input_wav", []) or []),
@@ -1371,7 +1476,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000),
validation_mode=validation_mode,
silence_ms=max(0, args.silence_ms),
frame_b64=demo_frame_b64,
frames_b64=demo_frames_b64,
stacked_frames_b64=demo_stacked_b64,
)
for turn_index in range(2, len(turn_specs)):
transcript, duration_ms = turn_specs[turn_index]
@@ -1389,7 +1495,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
realtime_input=realtime_input,
model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000),
commit_input=not continuous_input,
frame_b64=demo_frame_b64,
frames_b64=demo_frames_b64,
stacked_frames_b64=demo_stacked_b64,
)
turn_response_ids.append(response_id)
turn_outcomes.append(outcome)
@@ -1409,7 +1516,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
realtime_input=realtime_input,
model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000),
commit_input=not continuous_input,
frame_b64=demo_frame_b64,
frames_b64=demo_frames_b64,
stacked_frames_b64=demo_stacked_b64,
)
turn_response_ids.append(response_id)
turn_outcomes.append(outcome)
@@ -1584,6 +1692,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
"continuous_input": continuous_input,
"continuous_input_ok": continuous_input_ok,
"input_chunk_ms": args.chunk_ms,
"video_frame_count": len(demo_frames_b64),
"video_stacked_frame_count": sum(frame is not None for frame in demo_stacked_b64),
"model_policy_settle_ms": args.model_policy_settle_ms,
"turn_outcomes": turn_outcomes,
"turn_outcomes_ok": turn_outcomes_ok,
@@ -1641,6 +1751,19 @@ def parse_args() -> argparse.Namespace:
default=None,
help="Optional image file sent as an omni-duplex camera frame (one per 1 s of audio)",
)
parser.add_argument(
"--input-video",
default=None,
help="Optional video file streamed as omni-duplex camera frames (decoded at 1 fps, "
"overrides --frame-image; the last frame is held if the audio outlives the clip)",
)
parser.add_argument(
"--stack-frames",
type=int,
default=1,
choices=[1, 2],
help="Frames per append: 2 also resends the previous frame for motion context.",
)
parser.add_argument(
"--turn-input-wav",
action="append",
@@ -0,0 +1,85 @@
import base64
from pathlib import Path
from vllm_omni.experimental.fullduplex.client import (
PCM16_BYTES_PER_SAMPLE,
PCM16_SAMPLE_RATE,
RealtimeDuplexClient,
RealtimeEventCollector,
build_realtime_url,
read_pcm16_wav,
wait_for,
)
def _events(events: list[dict[str, object]], response_id: str, kind: str) -> list[dict[str, object]]:
return [e for e in events if e.get("type") == kind and RealtimeEventCollector.response_id(e) == response_id]
async def run_server_vad_interrupt(args) -> dict[str, object]:
initial = read_pcm16_wav(Path(args.input_wav))
interrupt = read_pcm16_wav(Path(args.interrupt_wav))[: PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * 47 // 10]
session_id = f"server-vad-hard-interrupt-{id(args)}"
client = RealtimeDuplexClient(build_realtime_url(args.url, args.model, autostart=False, session_id=session_id))
async def until(predicate, label: str) -> None:
await wait_for(predicate, timeout_s=args.timeout_s, label=label)
try:
await client.__aenter__()
await client.configure(
args.model,
ref_audio="data:audio/wav;base64," + base64.b64encode(Path(args.ref_audio).read_bytes()).decode(),
session_id=session_id,
turn_detection={"type": "server_vad", "interrupt_response": True},
timeout_s=args.timeout_s,
)
await client.stream_pcm16(initial + bytes(16_000 * 2 * 1600 // 1000), chunk_ms=200, realtime=True)
await client.commit()
await until(lambda: any(e.get("type") == "response.created" for e in client.events.events), "response.created")
target_id = RealtimeEventCollector.response_id(
next(e for e in client.events.events if e.get("type") == "response.created")
)
assert target_id
await until(lambda: _events(client.events.events, target_id, "response.audio.delta"), "response.audio.delta")
cursor = len(client.events.events)
await client.stream_pcm16(interrupt + bytes(16_000 * 2 * 800 // 1000), chunk_ms=args.chunk_ms, realtime=True)
await until(lambda: _events(client.events.events, target_id, "response.done"), "cancelled response.done")
await client.commit()
await until(
lambda: any(
e.get("type") == "response.done" and RealtimeEventCollector.response_id(e) != target_id
for e in client.events.events[cursor:]
),
"follow-up response.done",
)
events = client.events.events
finally:
try:
await client.close_session(timeout_s=args.timeout_s)
finally:
await client.__aexit__(None, None, None)
done = _events(events, target_id, "response.done")
terminal = done[0] if len(done) == 1 else None
trailing = events[events.index(terminal) + 1 :] if terminal is not None else events
stale = any(
e.get("type") == "response.audio.delta" and RealtimeEventCollector.response_id(e) == target_id for e in trailing
)
followup_audio = any(
e.get("type") == "response.audio.delta" and RealtimeEventCollector.response_id(e) != target_id
for e in events[cursor:]
)
ok = (
terminal is not None
and isinstance(terminal.get("response"), dict)
and terminal["response"].get("status") == "cancelled"
and terminal["response"].get("status_details", {}).get("reason") == "turn_detected"
and not stale
and sum(e.get("type") == "input_audio_buffer.speech_started" for e in events[cursor:]) == 1
and sum(e.get("type") == "input_audio_buffer.speech_stopped" for e in events[cursor:]) == 1
and any(e.get("type") == "input_audio_buffer.committed" for e in events[cursor:])
and followup_audio
and not any(e.get("type") == "error" for e in events)
)
return {"ok": ok}
@@ -5,8 +5,20 @@ Tests of common diffusion feature combinations in online serving mode
for HunyuanVideo-1.5-T2V (480p).
Coverage (H100, since model cannot fit L4):
- CacheDiT + Layerwise CPU offloading (1 GPU)
- CacheDiT + TP=2 + VAE patch parallel=2 (2 GPUs)
- CPU offloading (1 GPU) — ``core_model`` + ``advanced_model``
- CacheDiT + Layerwise CPU offloading (1 GPU) — ``full_model``
- CacheDiT + TP=2 + VAE patch parallel=2 (2 GPUs) — ``full_model``
HunyuanVideo-1.5 is a high-priority model, so only the most basic single-card deployment
row runs on every PR (L2) and on merge (L3). The heavyweight feature combinations stay
nightly-only (L4), together with the video similarity suites in
``tests/e2e/accuracy/hunyuanvideo15_{t2v,i2v}/``.
From ``tests/``::
pytest -s -v e2e/online_serving/test_hunyuan_video_15_expansion.py -m "core_model and diffusion" --run-level=core_model
pytest -s -v e2e/online_serving/test_hunyuan_video_15_expansion.py -m "advanced_model and diffusion" --run-level=advanced_model
pytest -s -v e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and diffusion" --run-level=full_model
"""
import pytest
@@ -14,7 +26,7 @@ import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
pytestmark = [pytest.mark.diffusion]
PROMPT = "A cat walking across a sunlit garden, cinematic lighting, slow motion."
NEGATIVE_PROMPT = "low quality, blurry, distorted"
@@ -24,9 +36,14 @@ MODEL = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
# Only the most basic single-card deployment row is cheap enough for PR (L2) / merge (L3)
# CI; the CacheDiT / parallel feature combinations run nightly (L4).
_CORE_MARKS = [pytest.mark.core_model, pytest.mark.advanced_model]
_FULL_MARKS = [pytest.mark.full_model]
def _get_diffusion_feature_cases(model: str):
"""Return L4 diffusion feature cases for HunyuanVideo-1.5.
"""Return diffusion feature cases for HunyuanVideo-1.5.
Designed for 2x H100 environment per issue #1832.
"""
@@ -40,7 +57,7 @@ def _get_diffusion_feature_cases(model: str):
],
),
id="single_card_cpu_offload",
marks=SINGLE_CARD_MARKS,
marks=SINGLE_CARD_MARKS + _CORE_MARKS,
),
# (1 GPU) CacheDiT + Layerwise CPU offloading
pytest.param(
@@ -53,7 +70,7 @@ def _get_diffusion_feature_cases(model: str):
],
),
id="single_card_cachedit_layerwise",
marks=SINGLE_CARD_MARKS,
marks=SINGLE_CARD_MARKS + _FULL_MARKS,
),
# (2 GPUs) CacheDiT + TP=2 + VAE patch parallel=2
pytest.param(
@@ -70,7 +87,7 @@ def _get_diffusion_feature_cases(model: str):
],
),
id="parallel_cachedit_tp2_vae2",
marks=PARALLEL_MARKS,
marks=PARALLEL_MARKS + _FULL_MARKS,
),
]
@@ -84,7 +101,7 @@ def test_hunyuan_video_15_t2v(
omni_server: OmniServer,
online_client: OnlineOmniClient,
):
"""L4 diffusion feature coverage for HunyuanVideo-1.5-T2V on H100."""
"""Diffusion feature coverage for HunyuanVideo-1.5-T2V on H100."""
form_data = {
"prompt": PROMPT,
"negative_prompt": NEGATIVE_PROMPT,
+3 -3
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Online serving smokes for all four canonical LTX-2.5 pipelines."""
@@ -73,8 +73,8 @@ def test_ltx25_pipeline_entries(
"form_data": {
"model": omni_server.model,
"prompt": PROMPT,
"height": 128,
"width": 128,
"height": 256,
"width": 256,
"num_frames": 9,
"fps": 24,
"num_inference_steps": num_inference_steps,
@@ -16,6 +16,7 @@ import websockets
from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
SERVER_PARAMS,
demo_args,
duplex_camera_frames,
multi_session_args,
realtime_url,
resolve_ref_audio,
@@ -30,10 +31,16 @@ from tests.e2e.online_serving.run_minicpmo_realtime_duplex_multi_session import
)
from tests.helpers.mark import hardware_test
from vllm_omni.experimental.fullduplex.client import build_realtime_url
from vllm_omni.experimental.fullduplex.video_stacking import concat_frames_b64
pytestmark = pytest.mark.omni
def _assert_positive_int(value: object) -> None:
assert isinstance(value, int)
assert value > 0
def _assert_request_metrics(metrics: object, *, expected_count: int) -> None:
assert isinstance(metrics, list)
assert len(metrics) == expected_count
@@ -142,7 +149,45 @@ def test_duplex_single_session_response_required(omni_server, tmp_path: Path) ->
args.turn_duration_ms = [args.first_turn_ms] * args.turns
result = asyncio.run(run_demo(args))
assert result["ok"] is True
assert result["audio_delta_count"] > 0
_assert_positive_int(result["audio_delta_count"])
assert result["done_count"] == 2
assert result["error_count"] == 0
assert result["all_audio_responses_have_transcript"] is True
assert result["transcript_delta_done_ok"] is True
_assert_request_metrics(result["request_metrics"], expected_count=2)
_assert_session_metrics(result["session_metrics"], expected_count=2)
@pytest.mark.advanced_model
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
def test_duplex_single_session_video_input(omni_server, tmp_path: Path) -> None:
"""Audio plus a 1 fps camera track, the omni-duplex video contract.
Frames ride the audio appends, so the response contract is unchanged: what
this covers is that interleaved vision input keeps the turn intact instead
of stalling or erroring out mid-segment.
"""
args = demo_args(
omni_server=omni_server,
input_wav=validated_input_wav(),
ref_audio=resolve_ref_audio(),
output_dir=tmp_path / "video_input",
)
args.turns = 2
args.turn_duration_ms = [args.first_turn_ms] * args.turns
frames = duplex_camera_frames(seconds=4, cache_dir=tmp_path / "camera")
args.video_frames_b64 = frames
# Each unit also carries a composite of its interior sub-frames, so the
# append exercises the official two-image frame_list that carries motion.
args.video_stacked_frames_b64 = [concat_frames_b64([frames[index]] * 2) for index in range(len(frames))]
result = asyncio.run(run_demo(args))
assert result["ok"] is True
assert result["video_frame_count"] == 4
assert result["video_stacked_frame_count"] == 4
_assert_positive_int(result["audio_delta_count"])
assert result["done_count"] == 2
assert result["error_count"] == 0
assert result["all_audio_responses_have_transcript"] is True
@@ -15,6 +15,7 @@ import pytest
from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
SERVER_PARAMS,
SOFT_INTERRUPT_SHA256,
deploy_max_sessions,
multi_session_args,
realtime_url,
resolve_ref_audio,
@@ -24,6 +25,7 @@ from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
from tests.e2e.online_serving.run_minicpmo_realtime_duplex_multi_session import (
run_lifecycle_probes,
)
from tests.e2e.online_serving.run_minicpmo_realtime_duplex_server_vad import run_server_vad_interrupt
from tests.e2e.online_serving.run_minicpmo_realtime_duplex_soft_interrupt import (
run_soft_interrupt,
)
@@ -46,7 +48,7 @@ def test_duplex_admission_and_expiry_reaper(omni_server, tmp_path: Path) -> None
args.disconnect_session_index = None
args.takeover_session_index = None
args.expire_session_index = 0
args.verify_admission_limit = 2
args.verify_admission_limit = deploy_max_sessions()
result = asyncio.run(run_lifecycle_probes(args))
assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
@@ -91,3 +93,22 @@ def test_duplex_soft_interrupt(omni_server, tmp_path: Path) -> None:
assert result["response_lifecycle_ok"] is True
assert result["response_audio_contract_ok"] is True
assert result["followup_response_transcript_ok"] is True
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
def test_duplex_server_vad_hard_interrupt(omni_server) -> None:
result = asyncio.run(
run_server_vad_interrupt(
SimpleNamespace(
url=realtime_url(omni_server),
model=omni_server.model,
input_wav=str(validated_input_wav()),
interrupt_wav=str(validated_soft_interrupt_wav()),
ref_audio=str(resolve_ref_audio()),
chunk_ms=200,
timeout_s=180.0,
)
)
)
assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
@@ -18,6 +18,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
DRIVER_DIR = Path(__file__).resolve().parent
HELPERS_DIR = DRIVER_DIR / "helpers"
DEMO_PATH = HELPERS_DIR / "minicpmo_realtime_duplex_scenarios.py"
EXAMPLE_DEMO_PATH = DRIVER_DIR.parents[2] / "examples/online_serving/minicpmo/realtime_duplex_demo.py"
MULTI_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_multi_session.py"
PAIR_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_demo_pair.py"
SOFT_INTERRUPT_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_soft_interrupt.py"
@@ -50,6 +51,18 @@ def _load_pair_demo_module():
return module
def _load_example_demo_module():
spec = importlib.util.spec_from_file_location(
"minicpmo_realtime_duplex_example_demo_test",
EXAMPLE_DEMO_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
def _load_soft_interrupt_demo_module():
spec = importlib.util.spec_from_file_location(
"minicpmo_realtime_duplex_soft_interrupt_test",
@@ -1813,6 +1826,219 @@ def test_realtime_duplex_demo_waits_at_each_model_unit_and_stops_after_speak():
assert len(ws.messages) == 10
def _sent_video_frames(messages: list[dict[str, object]]) -> list[list[str]]:
frames: list[list[str]] = []
for message in messages:
raw = message.get("video_frames")
if not isinstance(raw, list):
continue
frames.append([frame for frame in raw if isinstance(frame, str)])
return frames
def _decode_jpeg_size(frame_b64: str) -> tuple[int, int]:
import io
from PIL import Image
return Image.open(io.BytesIO(base64.b64decode(frame_b64))).size
class _RecordingWebSocket:
def __init__(self, demo):
self._demo = demo
self.messages: list[dict[str, object]] = []
async def send(self, payload):
self.messages.append(self._demo.json.loads(payload))
def _send_seconds_of_audio(demo, *, seconds: int, frames: list[str], stacked: list[str | None] | None = None):
ws = _RecordingWebSocket(demo)
asyncio.run(
demo._send_pcm16(
ws,
b"\x01\x00" * (demo.PCM16_SAMPLE_RATE * seconds),
chunk_ms=200,
realtime_delay=False,
frames_b64=frames,
stacked_frames_b64=stacked,
)
)
return ws.messages
def test_realtime_duplex_demo_streams_one_video_frame_per_model_unit():
demo = _load_demo_module()
messages = _send_seconds_of_audio(demo, seconds=3, frames=["f0", "f1", "f2", "f3"])
# Frame k rides the append that closes model unit k, so Stage0 can bind it
# to that unit's audio. 3 s of audio closes units at 1030 ms and 2030 ms and
# leaves a residual too short for a third, so only two frames go out.
assert _sent_video_frames(messages) == [["f0"], ["f1"]]
frame_indices = [index for index, message in enumerate(messages) if "video_frames" in message]
assert frame_indices == [5, 10]
# The carrying appends are exactly the ones that reach a unit boundary.
assert [messages[index]["audio_end_ms"] for index in frame_indices] == [1200, 2200]
def test_realtime_duplex_demo_never_sends_a_frame_before_its_unit_can_close():
"""Regression: frame 0 used to ride the very first 200 ms append.
Stage0 cannot close a unit until 1030 ms of audio has arrived, and it does
not carry frames across appends, so a frame sent that early was silently
dropped and every later frame ended up one unit ahead of its audio.
"""
demo = _load_demo_module()
messages = _send_seconds_of_audio(demo, seconds=3, frames=["f0", "f1", "f2"])
carried = [message for message in messages if "video_frames" in message]
assert carried, "expected at least one frame to be sent"
for message in carried:
assert message["audio_end_ms"] >= demo.duplex_unit_boundary_ms(0)
def test_realtime_duplex_demo_holds_the_last_video_frame_when_audio_outlives_the_clip():
demo = _load_demo_module()
# 4 s of audio closes three units (1030/2030/3030 ms); the two-frame clip
# holds its last frame for the third.
assert _sent_video_frames(_send_seconds_of_audio(demo, seconds=4, frames=["a", "b"])) == [["a"], ["b"], ["b"]]
# A still image is a one-element clip and therefore repeats every unit.
assert _sent_video_frames(_send_seconds_of_audio(demo, seconds=3, frames=["still"])) == [["still"], ["still"]]
def test_realtime_duplex_demo_sends_each_units_stacked_composite_next_to_its_base_frame():
demo = _load_demo_module()
messages = _send_seconds_of_audio(
demo,
seconds=4,
frames=["f0", "f1", "f2"],
stacked=["s0", "s1", None],
)
# Official pairing is frame_list=[base, composite of that unit's interior],
# so the composite belongs to the same unit as the base beside it -- never
# to the previous one. A unit without interior sub-frames sends base alone.
assert _sent_video_frames(messages) == [["f0", "s0"], ["f1", "s1"], ["f2"]]
def test_realtime_duplex_demo_decodes_a_video_file_into_one_frame_per_second(tmp_path):
pytest.importorskip("cv2")
pytest.importorskip("imageio")
from tests.helpers.media import generate_synthetic_video
demo = _load_demo_module()
# 30 fps synthetic clip: 90 frames is 3 s of camera input.
video = generate_synthetic_video(64, 64, 90, cache_dir=tmp_path)
frames, stacked = demo._resolve_video_frames(SimpleNamespace(input_video=video["file_path"]))
assert len(frames) == 3
assert all(base64.b64decode(frame, validate=True)[:3] == b"\xff\xd8\xff" for frame in frames)
assert len(set(frames)) == 3
# stack_frames defaults to 1, so there is nothing to stack.
assert stacked == [None, None, None]
def test_realtime_duplex_demo_stacks_a_units_interior_subframes_into_one_composite(tmp_path):
pytest.importorskip("cv2")
pytest.importorskip("imageio")
from tests.helpers.media import generate_synthetic_video
demo = _load_demo_module()
video = generate_synthetic_video(64, 64, 90, cache_dir=tmp_path)
frames, stacked = demo._resolve_video_frames(SimpleNamespace(input_video=video["file_path"], stack_frames=5))
# stack_frames=5 raises the visual refresh rate to 5 fps, but the four
# interior sub-frames of each second collapse into a single composite, so
# the wire still carries 2 images per unit and the audio cadence is
# untouched.
assert len(frames) == 3
assert len(stacked) == 3
assert all(frame is not None for frame in stacked)
assert all(base64.b64decode(frame, validate=True)[:3] == b"\xff\xd8\xff" for frame in stacked)
composite = _decode_jpeg_size(stacked[0])
assert composite != _decode_jpeg_size(frames[0])
def test_example_demo_uses_external_wav_duration_and_skips_video_soundtrack(tmp_path, monkeypatch):
demo = _load_example_demo_module()
wav = tmp_path / "external.wav"
demo.write_pcm16_wav(wav, b"\x00\x00" * (demo.PCM16_SAMPLE_RATE * 2), sample_rate_hz=demo.PCM16_SAMPLE_RATE)
soundtrack_calls: list[object] = []
def fake_soundtrack(*args, **kwargs):
soundtrack_calls.append(1)
raise AssertionError("video soundtrack must not be loaded when --input-wav is set")
def fake_frames(video_path, work_dir, *, fps, max_side, duration_s, stack_frames):
assert video_path == Path("/tmp/silent.mp4")
assert fps == 1.0
assert duration_s == pytest.approx(2.0)
return ["f0", "f1"], [None, None]
monkeypatch.setattr(demo, "_extract_video_soundtrack", fake_soundtrack)
monkeypatch.setattr(demo, "_extract_video_frames", fake_frames)
wav_out, frames, stacked = demo._resolve_duplex_av_inputs(
input_wav=str(wav),
input_video="/tmp/silent.mp4",
work_dir=tmp_path / "video_input",
fps=1.0,
max_side=0,
)
assert wav_out == str(wav)
assert frames == ["f0", "f1"]
assert stacked == [None, None]
assert soundtrack_calls == []
def test_example_demo_extracts_video_soundtrack_only_when_wav_is_absent(tmp_path, monkeypatch):
demo = _load_example_demo_module()
extracted = tmp_path / "from_video.wav"
demo.write_pcm16_wav(extracted, b"\x00\x00" * demo.PCM16_SAMPLE_RATE, sample_rate_hz=demo.PCM16_SAMPLE_RATE)
def fake_extract(video_path, work_dir, *, fps, max_side, stack_frames):
assert video_path == Path("/tmp/clip.mp4")
return extracted, ["v0"], ["s0"]
monkeypatch.setattr(demo, "_extract_video_input", fake_extract)
wav_out, frames, stacked = demo._resolve_duplex_av_inputs(
input_wav=None,
input_video="/tmp/clip.mp4",
work_dir=tmp_path / "video_input",
fps=1.0,
max_side=0,
)
assert wav_out == str(extracted)
assert frames == ["v0"]
assert stacked == ["s0"]
def test_realtime_duplex_demo_prefers_a_video_over_a_still_frame_image(tmp_path):
demo = _load_demo_module()
still = tmp_path / "frame.jpg"
still.write_bytes(b"\xff\xd8\xff-still")
# A still has no interior sub-frames, so it never stacks.
assert demo._resolve_video_frames(SimpleNamespace(frame_image=str(still))) == (
[base64.b64encode(still.read_bytes()).decode("ascii")],
[None],
)
assert demo._resolve_video_frames(
SimpleNamespace(frame_image=str(still), video_frames_b64=["preset"]),
) == (["preset"], [None])
assert demo._resolve_video_frames(SimpleNamespace()) == ([], [])
def test_realtime_duplex_demo_listen_only_overlap_accepts_silence_unit_before_first_done(monkeypatch):
demo = _load_demo_module()
state = demo.DemoState()
+5
View File
@@ -34,6 +34,11 @@ def test_sync_config_is_omni():
assert isinstance(cfg, OmniModelConfig)
def test_session_mode_reaches_omni_model_config():
assert OmniEngineArgs().create_model_config().session_mode == "turn"
assert OmniEngineArgs(session_mode="duplex").create_model_config().session_mode == "duplex"
def test_default_stage_id_is_concrete_int():
"""Ensure `stage_id` stays safe for downstream arithmetic/indexing."""
engine_args = OmniEngineArgs()
@@ -67,6 +67,39 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock
output_processor.add_request.assert_not_called()
def test_build_add_request_message_injects_global_id_before_prompt_transform(mocker: MockerFixture):
engine = object.__new__(AsyncOmniEngine)
params = SamplingParams(max_tokens=8)
engine.default_sampling_params_list = [params]
engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
engine.supported_tasks = ("speech",)
input_processor = mocker.Mock()
input_processor.process_inputs.return_value = _make_engine_core_request()
engine.input_processor = input_processor
engine.output_processors = [mocker.Mock()]
seen_global_ids: list[list[str]] = []
def transform(prompt, _sampling_params):
seen_global_ids.append(prompt["additional_information"]["global_request_id"])
return prompt
engine.prompt_transform_func = transform
prompt = {"prompt_token_ids": [1, 1, 1]}
message = engine._build_add_request_message(
request_id="req-1",
prompt=prompt,
sampling_params_list=[params],
final_stage_id=0,
arrival_time=0.0,
)
assert seen_global_ids == [["req-1"]]
assert message.original_prompt["additional_information"]["global_request_id"] == ["req-1"]
def test_build_add_request_message_preserves_model_intermediate_buffer(mocker: MockerFixture):
engine = object.__new__(AsyncOmniEngine)
params = SamplingParams(max_tokens=8)
@@ -433,11 +433,35 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline
"stage_id": 0,
"stage_init_timeout": 302,
"batch_size": 4,
"use_inline": True,
"use_inline": False,
"omni_master_server": None,
}
def test_initialize_local_diffusion_replica_uses_explicit_inline_opt_in(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
from vllm_omni.engine.stage_engine_startup import StageReplicaResources
runtime = _make_stage_runtime()
plan = _make_diffusion_plan(0, stage_id=1).replicas[0]
plan.stage_cfg.engine_args["inline_diffusion"] = True
captured: dict[str, object] = {}
monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None)
monkeypatch.setattr(
runtime_mod,
"launch_diffusion_stage_replica",
lambda **kwargs: (
captured.update(use_inline=kwargs["use_inline"]),
(types.SimpleNamespace(), StageReplicaResources()),
)[1],
)
runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1)
assert captured["use_inline"] is True
def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
@@ -1140,6 +1164,21 @@ def test_build_engine_args_keeps_stage_owned_tokenizer_subdir(tmp_path):
assert engine_args["tokenizer"] == os.path.join(str(tmp_path), "tokenizer")
def test_model_path_resolver_is_generic_and_model_owned(tmp_path):
from vllm_omni.engine.stage_init_utils import _resolve_model_path
engine_args = {
"model_path_resolver": ("vllm_omni.model_executor.models.minimax_h3.checkpoint.resolve_minimax_h3_model_root"),
"revision": None,
"task_type": "ref2va",
}
resolved = _resolve_model_path(str(tmp_path), engine_args)
assert resolved == str(tmp_path / "Ref2VA" / "text_encoder")
assert "model_path_resolver" not in engine_args
def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch):
import vllm_omni.engine.stage_init_utils as init_mod
+28
View File
@@ -74,6 +74,25 @@ class TestOmniPayloadStruct:
for top, sub in original.items():
assert sorted(d[top].keys()) == sorted(sub.keys())
def test_text_conditioning_uses_shared_payload_fields(self):
hidden = torch.zeros(4, 5120)
token_role_ids = torch.tensor([[1], [1], [0], [0]])
payload = to_dict(
to_struct(
{
"hidden_states": {"output": hidden},
"meta": {"token_role_ids": token_role_ids},
}
)
)
assert torch.equal(payload["hidden_states"]["output"], hidden)
assert torch.equal(payload["meta"]["token_role_ids"], token_role_ids)
for legacy_key in ("encoder_hidden_states", "token_tags"):
with pytest.raises(msgspec.ValidationError, match="unknown field"):
to_struct({legacy_key: torch.zeros(1)})
def test_to_dict_drops_unset_fields(self):
s = OmniPayloadStruct(meta=MetaStruct(left_context_size=10))
d = to_dict(s)
@@ -397,3 +416,12 @@ class TestSerializeDeserializePayload:
original: OmniPayload = {"meta": {"finished": None}}
wire = serialize_payload(original)
assert wire is None
def test_minimax_h3_prepared_reference_video_descriptor_round_trip(self):
descriptor = '{"artifact_dir":"/tmp/h3","videos":[]}'
wire = serialize_payload({"meta": {"minimax_h3_prepared_reference_videos": descriptor}})
assert wire is not None
restored = deserialize_payload(wire)
assert restored["meta"]["minimax_h3_prepared_reference_videos"] == descriptor
@@ -27,7 +27,11 @@ from vllm_omni.engine.messages import (
ShutdownRequestMessage,
StageSubmissionMessage,
)
from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
from vllm_omni.engine.orchestrator import (
Orchestrator,
OrchestratorRequestState,
cleanup_request_artifact_dirs,
)
from vllm_omni.engine.stage_pool import StageUnavailableError
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
@@ -44,6 +48,18 @@ from .test_orchestrator import (
_wait_for,
)
def test_request_artifact_cleanup_is_idempotent(tmp_path):
artifact_dir = tmp_path / "request-artifacts"
artifact_dir.mkdir()
(artifact_dir / "prepared.mp4").touch()
cleanup_request_artifact_dirs([str(artifact_dir)])
cleanup_request_artifact_dirs([str(artifact_dir)])
assert not artifact_dir.exists()
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -919,6 +935,36 @@ async def test_diffusion_client_error_output_propagates_status_code(orchestrator
orchestrator_fixture.thread.join(timeout=5)
@pytest.mark.asyncio
async def test_duplex_input_processor_failure_is_request_scoped(orchestrator_factory, monkeypatch) -> None:
class FailingStage(FakeStageClient):
def process_engine_inputs(self, *_args, **_kwargs):
raise ValueError("No latent or hidden_states found in thinker output")
fixture = orchestrator_factory(
[FakeStageClient(stage_type="llm"), FailingStage(stage_type="llm", final_output=True)]
)
state = OrchestratorRequestState(
request_id="bad",
prompt=SimpleNamespace(request_id="bad", prompt_token_ids=[1]),
sampling_params_list=[_sampling_params(), _sampling_params()],
final_stage_id=1,
duplex_identity=SimpleNamespace(),
)
async def no_cleanup(*_args, **_kwargs):
pass
monkeypatch.setattr(fixture.orchestrator, "_cleanup_request_ids", no_cleanup)
try:
await fixture.orchestrator._forward_to_next_stage_unguarded("bad", 0, _build_request_output("raw"), state)
error = await _wait_for_error_message(fixture, request_id="bad")
assert error.fatal is False
finally:
fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
fixture.thread.join(timeout=5)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("status_code", "error_type"),
+42
View File
@@ -1207,6 +1207,48 @@ class TestSingleStageReplicaInitialization:
)
assert mock_from_addresses.call_args.kwargs["proc_manager"].proc is proc
def test_initialize_local_custom_diffusion_replica_uses_inline_client(self, mocker: MockerFixture):
import vllm_omni.engine.stage_runtime as runtime_mod
from vllm_omni.platforms import current_omni_platform
runtime = DistStageRuntime(
stage_configs=[],
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=4,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
omni_master_port=26000,
)
runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer)
runtime._coordinator_runtime = None
plan = _make_diffusion_plan(0, stage_id=0, launch_mode="local").replicas[0]
plan.stage_cfg.engine_args = {"custom_pipeline_args": {"pipeline_class": "test.CustomPipeline"}}
sentinel_client = SimpleNamespace()
mock_launch = mocker.patch.object(
runtime_mod,
"launch_diffusion_stage_replica",
return_value=(sentinel_client, StageReplicaResources()),
)
mocker.patch.object(runtime_mod, "inject_kv_stage_info")
device_env_var = current_omni_platform.device_control_env_var
prev_device_env = os.environ.get(device_env_var)
os.environ[device_env_var] = "0"
runtime._init_visible_devices_baseline = "0"
try:
result = runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=60)
finally:
if prev_device_env is None:
os.environ.pop(device_env_var, None)
else:
os.environ[device_env_var] = prev_device_env
assert result is sentinel_client
assert mock_launch.call_args.kwargs["use_inline"] is True
def test_initialize_local_diffusion_replica_failure_terminates_proc(self, mocker: MockerFixture):
import vllm_omni.engine.stage_runtime as runtime_mod
from vllm_omni.platforms import current_omni_platform
+5 -1
View File
@@ -72,6 +72,7 @@ def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str,
_LLM_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniEngineArgs))
_DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusionConfig))
_TOPOLOGY_ONLY_ENGINE_ARGS = frozenset({"inline_diffusion"})
_OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
{
"active_stream_window",
@@ -368,7 +369,9 @@ def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
)
typed_args_by_stage[stage_id] = typed_args
assert {name: typed_args[name] for name in legacy_args} == legacy_args
expected_args = {name: value for name, value in legacy_args.items() if name not in _TOPOLOGY_ONLY_ENGINE_ARGS}
assert {name: typed_args[name] for name in expected_args} == expected_args
assert _TOPOLOGY_ONLY_ENGINE_ARGS.isdisjoint(typed_args)
thinker_args = typed_args_by_stage[0]
inherited_vllm_fields = {
@@ -658,6 +661,7 @@ def test_typed_engine_args_match_current_registry_backend_semantics(model_type):
backend_fields = _LLM_BACKEND_FIELDS
backend_config_cls = OmniEngineArgs
backend_fields -= _TOPOLOGY_ONLY_ENGINE_ARGS
missing_fields = (legacy_args.keys() & backend_fields) - typed_args.keys()
assert not missing_fields, f"{model_type} stage {stage_id} lost backend fields: {sorted(missing_fields)}"
@@ -452,7 +452,7 @@ def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease():
assert caps["supports_core_resumable_request"] is True
assert caps["supports_stage_connector_handoff"] is True
assert caps["supports_audio_truncate"] is True
assert caps["supports_barge_in"] is False
assert caps["supports_barge_in"] is True
assert caps["target_barge_in_latency_ms"] is None
assert caps["supports_multi_session"] is True
assert caps["supports_multi_session_same_replica"] is True
@@ -0,0 +1,37 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from vllm_omni.entrypoints.openai import api_server
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.mark.parametrize(
("path", "engine_method"),
[
("/start_profile", "start_profile"),
("/stop_profile", "stop_profile"),
],
)
def test_profiler_endpoint_forwards_selected_stages(path: str, engine_method: str) -> None:
engine_client = SimpleNamespace(
start_profile=AsyncMock(),
stop_profile=AsyncMock(),
)
app = FastAPI()
app.include_router(api_server.profiler_router)
app.state.engine_client = engine_client
response = TestClient(app).post(path, json={"stages": [1]})
assert response.status_code == 200
assert response.json() == {"status": "SUCCESS"}
getattr(engine_client, engine_method).assert_awaited_once_with(stages=[1])
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""API-server surface guards for the OpenAI Omni entrypoint.
Purpose
@@ -20,8 +20,8 @@ Common wrong changes these tests are meant to catch
* Route accidentally removed, renamed, method-changed, or double-registered
after a move (census / assembled-app checks).
* OpenAPI drops an Omni HTTP path while the router still has it (or the reverse).
* Assembly forgets to strip upstream ``/v1/chat/completions`` (batch) /
``/v1/models``, or also deletes unrelated upstream routes.
* Assembly forgets to strip overridden upstream chat / models / profiler
routes, or also deletes unrelated upstream routes.
* Profiler / Omni router not mounted; storage not started; Omni engine
exception handlers not registered.
* Timestamp middleware stops stamping HTTP, or starts mutating WebSocket scopes.
@@ -351,9 +351,10 @@ def test_router_openapi_paths_cover_http_manifest() -> None:
async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_router(monkeypatch) -> None:
"""Lock worker assembly: override, mount, storage, handlers, full census.
Fails if assembly stops replacing upstream chat/batch/models, deletes
unrelated upstream routes, forgets Omni/profiler mounts, skips storage
start, or drops Omni ``EngineDeadError`` / ``EngineGenerateError`` handlers.
Fails if assembly stops replacing upstream chat/batch/models/profiler,
deletes unrelated upstream routes, forgets Omni/profiler mounts, skips
storage start, or drops Omni ``EngineDeadError`` / ``EngineGenerateError``
handlers.
"""
captured: dict[str, object] = {}
@@ -372,6 +373,14 @@ async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_rout
async def upstream_models():
return {"owner": "upstream"}
@app.post("/start_profile")
async def upstream_start_profile():
return {"owner": "upstream"}
@app.post("/stop_profile")
async def upstream_stop_profile():
return {"owner": "upstream"}
@app.get("/v1/upstream-only")
async def upstream_only():
return {"owner": "upstream"}
@@ -437,6 +446,8 @@ async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_rout
is api_server.create_batch_chat_completion
)
assert _single_http_route(routes, "GET", "/v1/models").endpoint is api_server.show_available_models
assert _single_http_route(routes, "POST", "/start_profile").endpoint is api_server.start_profile
assert _single_http_route(routes, "POST", "/stop_profile").endpoint is api_server.stop_profile
# Assembled app includes Omni router HTTP/WS surface + profiler.
for item in _EXPECTED_ROUTER_ROUTES:
@@ -33,6 +33,7 @@ from vllm_omni.experimental.fullduplex.minicpmo45.runtime import (
from vllm_omni.experimental.fullduplex.minicpmo45.session import (
MiniCPMO45ServingSessionState,
)
from vllm_omni.experimental.fullduplex.openai import vad as realtime_vad
from vllm_omni.experimental.fullduplex.openai.protocol import (
DuplexCapabilities,
DuplexOverlapPolicy,
@@ -706,7 +707,7 @@ async def test_native_realtime_protocol_preserves_input_turn_policy_hints():
@pytest.mark.asyncio
async def test_native_realtime_protocol_rejects_unimplemented_server_vad():
async def test_native_realtime_protocol_rejects_server_vad_without_interrupt():
ws = TimedWebSocket()
protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
protocol.bind_sender(ws.send_json)
@@ -729,7 +730,11 @@ async def test_native_realtime_protocol_rejects_unimplemented_server_vad():
error = ws.sent[-1]
assert error["type"] == "error"
assert error["error"]["code"] == "unsupported_turn_detection"
assert "server_vad" in error["error"]["message"]
assert "interrupt_response=false" in error["error"]["message"]
threshold_error = protocol._validate_realtime_turn_detection(
{"turn_detection": {"type": "server_vad", "threshold": 0.15}}
)
assert threshold_error == "turn_detection.threshold must be greater than 0.15 and at most 1"
@pytest.mark.asyncio
@@ -753,7 +758,7 @@ async def test_native_realtime_protocol_accepts_disabled_turn_detection():
@pytest.mark.asyncio
async def test_native_realtime_protocol_rejects_nested_vad_even_when_top_level_is_disabled():
async def test_native_realtime_protocol_prefers_nested_vad_over_top_level():
ws = TimedWebSocket()
protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
protocol.bind_sender(ws.send_json)
@@ -773,8 +778,55 @@ async def test_native_realtime_protocol_rejects_nested_vad_even_when_top_level_i
}
)
assert translated is None
assert ws.sent[-1]["error"]["code"] == "unsupported_turn_detection"
assert translated["session"]["extra_body"]["realtime_session_payload"]["overlap_policy"] == "barge_in_on_speech"
@pytest.mark.asyncio
async def test_realtime_vad_edge_order_and_tiny_chunk_resampling(monkeypatch):
ws = TimedWebSocket()
protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
protocol.bind_sender(ws.send_json)
protocol._server_vad = realtime_vad.SileroStreamingVAD(realtime_vad.SileroVADConfig())
result = realtime_vad.StreamingVADResult(True, True, True, True)
monkeypatch.setattr(protocol._server_vad, "process_base64", lambda *args, **kwargs: result)
protocol._input_speech_started = True
protocol._hold_realtime_output_until_session_created = False
await protocol._to_duplex_event(
{"type": "input_audio_buffer.append", "audio": _pcm_f32_b64(1), "format": "pcm_f32le"}
)
assert [event["type"] for event in ws.sent] == [
"input_audio_buffer.speech_stopped",
"input_audio_buffer.speech_started",
]
ws.sent.clear()
protocol._input_speech_started = False
result = realtime_vad.StreamingVADResult(True, False, True, True)
await protocol._to_duplex_event({"type": "input_audio_buffer.append", "audio": "AAAA", "format": "pcm_f32le"})
assert [event["type"] for event in ws.sent] == [
"input_audio_buffer.speech_started",
"input_audio_buffer.speech_stopped",
]
detector = realtime_vad.SileroStreamingVAD(realtime_vad.SileroVADConfig(), frame_scorer=lambda _: 0.0)
for _ in range(1024):
detector.process_base64(_pcm_f32_b64(1), fmt="pcm_f32le", sample_rate_hz=48_000)
assert detector._processed_samples + detector._pending.size == 341
@pytest.mark.asyncio
async def test_realtime_server_vad_runtime_failures_are_terminal(monkeypatch):
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient()))
ws = TimedWebSocket()
ws.put(_server_vad_update("generic-vad", model="test-model", session_id="generic-vad"))
await handler.handle_session(ws, realtime_protocol=NativeRealtimeSessionProtocol({}))
assert ws.sent[-1]["error"]["code"] == "server_vad_requires_native_duplex"
monkeypatch.setattr(realtime_vad.SileroStreamingVAD, "process_base64", lambda *args, **kwargs: 1 / 0)
ws = TimedWebSocket()
create = _native_realtime_session_update("vad-inference-failure")
create["session"]["turn_detection"] = {"type": "server_vad"}
ws.put(create)
ws.put({"type": "input_audio_buffer.append", "audio": _pcm_f32_b64(1), "format": "pcm_f32le"})
await asyncio.wait_for(handler.handle_session(ws, realtime_protocol=NativeRealtimeSessionProtocol({})), timeout=2)
assert any(e.get("error", {}).get("code") == "realtime_input_failed" for e in ws.sent)
def test_native_duplex_handler_has_no_fixed_session_admission_cap():
@@ -1551,6 +1603,14 @@ def _native_realtime_session_update(
}
def _server_vad_update(event_id: str, **session: object) -> dict[str, object]:
return {
"type": "session.update",
"event_id": event_id,
"session": {"turn_detection": {"type": "server_vad"}, **session},
}
def test_native_stage0_request_id_matches_engine_after_same_session_id_reopen():
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient()))
first = handler._registry.create(session_id="sid-reused")
@@ -3148,7 +3208,7 @@ async def test_minicpmo_pre_response_continuation_drops_after_model_turn_ends():
@pytest.mark.asyncio
async def test_minicpmo_auto_response_continuation_has_no_semantic_unit_cap():
async def test_minicpmo_auto_response_continuation_stops_at_large_safety_boundary():
request_id = "duplex-sid-long-response-e0-stage0"
engine = FakeEngineClient()
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine))
@@ -3164,7 +3224,7 @@ async def test_minicpmo_auto_response_continuation_has_no_semantic_unit_cap():
_install_direct_silence_scheduler(handler, session)
ws = TimedWebSocket()
for _ in range(9):
for _ in range(handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS):
await handler._maybe_continue_native_response(
ws.send_json,
session=session,
@@ -3172,9 +3232,103 @@ async def test_minicpmo_auto_response_continuation_has_no_semantic_unit_cap():
)
await asyncio.sleep(0.02)
assert len(engine.appended) == 9
response_id = session.active_response_id
assert response_id is not None
assert len(engine.appended) == handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS
assert all(payload["duplex_turn_id"] == 0 for _, _, payload, _ in engine.appended)
assert all("force_speak" not in payload for _, _, payload, _ in engine.appended)
assert all(not {"force_speak", "force_listen"} & payload.keys() for _, _, payload, _ in engine.appended[:-1])
assert engine.appended[-1][2]["force_listen"] is True
assert handler._native_response_continuations_remaining(session, response_id) is False
await handler._maybe_continue_native_response(
ws.send_json,
session=session,
expected_epoch=session.epoch,
)
assert len(engine.appended) == handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS
assert ws.sent_types()[-2:] == ["response.listen", "response.done"]
assert session.active_response_id is None
assert session.active_request_id == request_id
@pytest.mark.asyncio
async def test_minicpmo_auto_response_boundary_listen_closes_response():
request_id = "duplex-sid-boundary-listen-e0-stage0"
engine = FakeEngineClient()
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine))
session = DuplexSession(
session_id="sid-boundary-listen",
config=DuplexSessionConfig(extra_body={"auto_response": True}),
)
session.capabilities = DuplexCapabilities.minicpmo45_native()
response_id = session.begin_response(turn_id=0)
session.bind_request(request_id)
native = handler._minicpmo_session_state(session)
native.continuation_owner_id = f"response:{response_id}"
native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS
_install_direct_silence_scheduler(handler, session)
ws = TimedWebSocket()
close_reason, emitted = await handler._send_one_native_duplex_event(
ws.send_json,
{
"supported": True,
"stage_role": "llm",
"is_listen": True,
"model_listen": True,
"data_plane_request_id": request_id,
"end_of_turn": False,
"uses_model_runner_scheduler": True,
"runner_kv_backed": True,
"runtime_impl": "scheduler_data_plane",
},
session=session,
expected_epoch=session.epoch,
)
assert close_reason is None
assert emitted is True
assert engine.appended == []
assert ws.sent_types() == ["response.listen", "response.done"]
assert session.active_response_id is None
assert session.active_request_id == request_id
assert native.continuation_owner_id is None
@pytest.mark.asyncio
async def test_minicpmo_auto_response_boundary_does_not_close_new_epoch_response():
request_id = "duplex-sid-boundary-race-e0-stage0"
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient()))
session = DuplexSession(
session_id="sid-boundary-race",
config=DuplexSessionConfig(extra_body={"auto_response": True}),
)
session.capabilities = DuplexCapabilities.minicpmo45_native()
old_response_id = session.begin_response(turn_id=0)
session.bind_request(request_id)
native = handler._minicpmo_session_state(session)
native.continuation_owner_id = f"response:{old_response_id}"
native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS
sent: list[dict[str, Any]] = []
new_response_ids: list[str] = []
async def _barge_in_on_listen(payload: dict[str, Any]) -> None:
sent.append(payload)
if payload["type"] == "response.listen":
session.barge_in()
session.bind_request("duplex-sid-boundary-race-e1-stage0")
new_response_ids.append(session.begin_response(turn_id=1))
await handler._maybe_continue_native_response(
_barge_in_on_listen,
session=session,
expected_epoch=session.epoch,
)
assert [payload["type"] for payload in sent] == ["response.listen"]
assert len(new_response_ids) == 1
assert session.active_response_id == new_response_ids[0]
@pytest.mark.asyncio
@@ -5217,7 +5371,7 @@ async def test_audio_clear_preserves_later_wire_order_append():
@pytest.mark.asyncio
async def test_native_append_is_cancelled_by_later_wire_order_input_cancel():
async def test_native_append_is_cancelled_by_later_wire_order_input_cancel(mocker):
class SlowAppendEngine(FakeEngineClient):
def __init__(self) -> None:
super().__init__()
@@ -5254,6 +5408,7 @@ async def test_native_append_is_cancelled_by_later_wire_order_input_cancel():
fence=fence,
)
rollback = mocker.spy(MiniCPMO45PcmAppendBuffer, "_rollback_reservation")
engine = SlowAppendEngine()
handler = OmniDuplexSessionHandler(
chat_service=FakeChatService(engine),
@@ -5284,6 +5439,7 @@ async def test_native_append_is_cancelled_by_later_wire_order_input_cancel():
await asyncio.wait_for(handler_task, timeout=1)
assert engine.append_cancelled.is_set()
rollback.assert_called_once()
assert engine.signal_fences[-1] == DuplexFence("sid-cancel-slow-append")
assert engine.signal_next_fences[-1] == DuplexFence("sid-cancel-slow-append", epoch=1)
@@ -5379,33 +5535,59 @@ async def test_native_session_update_commits_config_only_after_runtime_ack():
return await super().signal_duplex_turn_async(session_id, **kwargs)
engine = BlockingUpdateEngine()
handler = OmniDuplexSessionHandler(
chat_service=FakeChatService(engine),
config_timeout_s=0.1,
idle_timeout_s=1,
)
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine))
ws = TimedWebSocket()
ws.put(_session_create("sid-update-two-phase"))
ws.put(
{
"type": "turn.signal",
"event": "session.update",
"payload": {"instructions": "candidate instructions"},
}
)
protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
ws.put(_native_realtime_session_update("sid-update-two-phase"))
ws.put(_server_vad_update("update-two-phase"))
handler_task = asyncio.create_task(handler.handle_session(ws))
handler_task = asyncio.create_task(handler.handle_session(ws, realtime_protocol=protocol))
await asyncio.wait_for(engine.update_started.wait(), timeout=1)
live_session = handler._registry.get("sid-update-two-phase")
assert live_session is not None
assert live_session.config.instructions != "candidate instructions"
assert live_session.config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value
assert protocol._pending_turn_detection_update is not None
assert protocol._server_vad is None
protocol.encode_outbound_event({"type": "error", "code": "unrelated", "error": "unrelated"})
assert protocol._pending_turn_detection_update is not None
engine.release_update.set()
ws.put({"type": "session.close"})
await asyncio.wait_for(handler_task, timeout=2)
updated = next(message for message in ws.sent if message.get("type") == "session.updated")
assert updated["session"]["instructions"] == "candidate instructions"
updated = next(message for message in reversed(ws.sent) if message.get("type") == "session.updated")
assert updated["session"]["overlap_policy"] == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value
assert isinstance(protocol._server_vad, realtime_vad.SileroStreamingVAD)
@pytest.mark.parametrize(
("append_fails", "expected_code"),
[(True, "session_update_aborted"), (False, "instructions_update_unsupported")],
)
@pytest.mark.asyncio
async def test_realtime_vad_update_rejection_does_not_stall(append_fails, expected_code):
append_result = {"operation": "append", "ok": False, "error_count": 1} if append_fails else None
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient(append_result=append_result)))
ws = TimedWebSocket()
protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
create = _native_realtime_session_update("sid-rejected-vad-update")
create["session"]["extra_body"]["auto_response"] = append_fails
ws.put(create)
if append_fails:
ws.put({"type": "input_audio_buffer.append", "audio": _pcm_f32_b64(16_000), "format": "pcm_f32le"})
update = _server_vad_update("rejected-vad-update")
if not append_fails:
update["session"]["instructions"] = "You are now a pirate."
ws.put(update)
ws.put({"type": "session.close"})
await asyncio.wait_for(handler.handle_session(ws, realtime_protocol=protocol), timeout=2)
errors = [message["error"] for message in ws.sent if isinstance(message.get("error"), dict)]
error = next(error for error in errors if error.get("code") == expected_code)
assert error["event_id"] == "rejected-vad-update"
assert protocol._pending_turn_detection_update is None
assert protocol._server_vad is None
@pytest.mark.asyncio
@@ -5905,7 +6087,7 @@ async def test_minicpmo_native_duplex_rejects_ref_audio_path():
assert engine.opened == []
def test_minicpmo_native_duplex_explicit_barge_in_request_is_deferred_as_listen():
def test_minicpmo_native_duplex_explicit_barge_in_request_interrupts():
engine = FakeEngineClient()
handler = OmniDuplexSessionHandler(
chat_service=FakeChatService(engine),
@@ -5932,12 +6114,10 @@ def test_minicpmo_native_duplex_explicit_barge_in_request_is_deferred_as_listen(
)
assert decision == {
"action": "listen",
"reason": "barge_in_unsupported",
"action": "barge_in",
"reason": "client_force_barge_in",
"duration_ms": 1000,
"overlap_speech_ms": 1000,
"buffer_audio": True,
"defer_runtime_append": True,
}
@@ -5949,7 +6129,7 @@ def test_minicpmo_native_duplex_explicit_barge_in_request_is_deferred_as_listen(
{"type": "turn.signal", "event": "barge_in"},
],
)
async def test_minicpmo_native_duplex_rejects_unadvertised_barge_in_control(
async def test_minicpmo_native_duplex_accepts_explicit_barge_in_control(
barge_in_event: dict[str, object],
):
engine = FakeEngineClient()
@@ -5965,9 +6145,8 @@ async def test_minicpmo_native_duplex_rejects_unadvertised_barge_in_control(
await handler.handle_session(ws)
error = next(message for message in ws.sent if message.get("code") == "barge_in_unsupported")
assert error["session_id"] == "sid-native-barge-control-disabled"
assert ("sid-native-barge-control-disabled", "barge_in") not in engine.signals
assert not any(message.get("code") == "barge_in_unsupported" for message in ws.sent)
assert ("sid-native-barge-control-disabled", "barge_in") in engine.signals
@pytest.mark.asyncio
@@ -47,7 +47,13 @@ def _build_test_app(
engine_client.abort = mocker.AsyncMock()
engine_client.submit_interaction_async = mocker.AsyncMock()
engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
engine_client.stage_configs = [
SimpleNamespace(
stage_type="diffusion",
final_output=True,
final_output_type="video",
)
]
engine_client.generate = mock_generate
class FakeStreamingVideoEncoder:
@@ -72,10 +78,6 @@ def _build_test_app(
"vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
encoder_factory,
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
return_value="diffusion",
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
return_value=[OmniDiffusionSamplingParams()],
@@ -398,7 +400,13 @@ def _build_handler_for_async_tests(
engine_client.abort = mocker.AsyncMock()
engine_client.submit_interaction_async = mocker.AsyncMock()
engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
engine_client.stage_configs = [
SimpleNamespace(
stage_type="diffusion",
final_output=True,
final_output_type="video",
)
]
engine_client.generate = mock_generate
chunks = streaming_chunks or [(b"mp4-chunk-0", True)]
@@ -421,10 +429,6 @@ def _build_handler_for_async_tests(
"vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
encoder_factory,
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
return_value="diffusion",
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
return_value=[OmniDiffusionSamplingParams()],
@@ -0,0 +1,66 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
import pytest
from vllm_omni.diffusion.io_support import get_diffusion_output_type
from vllm_omni.entrypoints.openai.utils import is_video_generation_pipeline
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_video_pipeline_requires_declared_final_video_stage():
assert is_video_generation_pipeline(
[
SimpleNamespace(
stage_type="llm",
final_output=False,
final_output_type=None,
),
SimpleNamespace(
stage_type="diffusion",
final_output=True,
final_output_type="video",
),
]
)
@pytest.mark.parametrize(
"stage_configs",
[
[SimpleNamespace(stage_type="diffusion")],
[
SimpleNamespace(
stage_type="diffusion",
final_output=True,
final_output_type="image",
)
],
[
{
"stage_type": "diffusion",
"final_output": False,
"final_output_type": "video",
}
],
],
)
def test_video_pipeline_rejects_non_video_final_outputs(stage_configs):
assert not is_video_generation_pipeline(stage_configs)
@pytest.mark.parametrize(
"model_class_name",
[
"LTX2TwoStagePipeline",
"LTX2DistilledOneStagePipeline",
"WanDMDPipeline",
"LingBotWorldCausalDMDPipeline",
"LongCatVideoAvatarPipeline",
],
)
def test_registered_video_aliases_declare_video_output(model_class_name):
assert get_diffusion_output_type(model_class_name) == "video"
@@ -64,7 +64,13 @@ class MockVideoResult:
class FakeAsyncOmni:
def __init__(self):
self.stage_configs = [SimpleNamespace(stage_type="diffusion")]
self.stage_configs = [
SimpleNamespace(
stage_type="diffusion",
final_output=True,
final_output_type="video",
)
]
self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
self.model_class_name = "WanPipeline"
self.captured_prompt = None
@@ -828,6 +834,18 @@ def test_multi_video_generation_preserves_uploaded_files_until_generation(
assert engine.captured_sampling_params_list[0].extra_args["duration"] == 15.0
def test_mixed_reference_capability_uses_model_metadata_when_config_defaults_false(test_client):
handler = test_client.app.state.openai_serving_video
handler._engine_client.model_class_name = None
handler._engine_client.stage_configs = [
SimpleNamespace(engine_args={"model_class_name": "MiniMaxH3TextEncoder"}),
SimpleNamespace(engine_args={"model_class_name": "MiniMaxH3Pipeline"}),
]
handler._stage_configs = handler._engine_client.stage_configs
assert handler.supports_mixed_reference_inputs
def test_decode_video_bytes_can_keep_first_frames():
from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes
@@ -33,6 +33,64 @@ def test_default_stage_config_includes_cache_backend():
assert engine_args["model_stage"] == "diffusion"
def test_default_stage_config_preserves_model_extras():
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"extras": {"ltx2_use_conv_vae": True}})[0]
assert stage_cfg["engine_args"]["extras"]["ltx2_use_conv_vae"] is True
def test_default_stage_config_preserves_and_overrides_promoted_extras():
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
{
"extras": {
"auxiliary_text_encoder": "/models/extras-llama",
"default_llama_model_id": "extras/default-llama",
},
"auxiliary_text_encoder": None,
}
)[0]
extras = stage_cfg["engine_args"]["extras"]
assert extras["auxiliary_text_encoder"] == "/models/extras-llama"
assert extras["default_llama_model_id"] == "extras/default-llama"
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
{
"extras": {
"auxiliary_text_encoder": "/models/extras-llama",
"default_llama_model_id": "extras/default-llama",
},
"auxiliary_text_encoder": "/models/top-level-llama",
"default_llama_model_id": "top-level/default-llama",
}
)[0]
extras = stage_cfg["engine_args"]["extras"]
assert extras["auxiliary_text_encoder"] == "/models/top-level-llama"
assert extras["default_llama_model_id"] == "top-level/default-llama"
def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocker):
"""Local/unregistered Diffusers checkpoints still honor stage-0 extras."""
def resolve_with_default(*_args, default_stage_cfg_factory, **_kwargs):
return None, default_stage_cfg_factory(), None
mocker.patch(
"vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
side_effect=resolve_with_default,
)
engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
_, stage_configs = engine._resolve_stage_configs(
"/models/LTX-2.5-Diffusers",
{"stage_overrides": '{"0":{"extras":{"ltx2_use_conv_vae":true}}}'},
trust_remote_code=False,
)
assert stage_configs[0]["engine_args"]["extras"]["ltx2_use_conv_vae"] is True
def test_default_cache_config_used_when_missing():
"""Ensure default cache_config is synthesized when only backend is given."""
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
@@ -257,6 +315,29 @@ def test_serve_cli_accepts_ulysses_mode():
assert parallel_config.ulysses_mode == "advanced_uaa"
def test_serve_cli_accepts_text_encoder_tp_size():
parser = TrackingArgumentParser()
subparsers = parser.add_subparsers(dest="command")
OmniServeCommand().subparser_init(subparsers)
args = parser.parse_args(
[
"serve",
"MiniMaxAI/MiniMax-H3",
"--omni",
"--text-encoder-tp-size",
"4",
]
)
explicit_kwargs = args.get_explicit_kwargs_dict()
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
parallel_config = stage_cfg["engine_args"]["parallel_config"]
assert args.text_encoder_tp_size == 4
assert parallel_config.text_encoder_tp_size == 4
def test_serve_cli_forwards_model_defined_task_type_to_diffusion_stage():
parser = TrackingArgumentParser()
subparsers = parser.add_subparsers(dest="command")
@@ -516,6 +597,36 @@ def test_serve_cli_accepts_additional_config():
assert engine_args["additional_config"] == {"torchair_graph_config": {"enabled": True}}
def test_default_stage_resolves_video_output_from_checkpoint(mocker):
captured = {}
def resolve_with_default(*args, default_stage_cfg_factory, **kwargs):
del args, kwargs
captured.update(default_stage_cfg_factory()[0])
stage = SimpleNamespace(stage_type="diffusion", engine_args=SimpleNamespace())
return ("", [stage], None)
resolver = mocker.patch(
"vllm_omni.engine.async_omni_engine.resolve_model_class_name",
return_value="MiniMaxH3Pipeline",
)
mocker.patch(
"vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
side_effect=resolve_with_default,
)
engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
engine._strip_single_engine_args = lambda kwargs: kwargs
engine._resolve_stage_configs(
"/models/MiniMax-H3/FL2VA",
{},
trust_remote_code=False,
)
resolver.assert_called_once_with("/models/MiniMax-H3/FL2VA", "default")
assert captured["final_output_type"] == "video"
def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mocker):
"""Ensure YAML/deploy stage resolution forwards top-level additional_config."""
fake_diffusion_stage = SimpleNamespace(
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import base64
import importlib.util
import sys
@@ -169,6 +172,64 @@ def test_open_streaming_response_requires_post_commit_drain():
)
def _asymmetric_image():
from PIL import Image
image = Image.new("RGB", (4, 2), (0, 0, 0))
image.putpixel((0, 0), (255, 0, 0))
return image
@pytest.mark.parametrize(
"rotation,expected_corner",
[
(0, (0, 0)),
(90, (0, 3)),
(180, (3, 1)),
(-90, (1, 0)),
(270, (1, 0)),
(-270, (0, 3)),
],
)
def test_upright_applies_the_display_matrix_quarter_turn(rotation, expected_corner):
demo = _load_demo_module()
rotated = demo._upright(_asymmetric_image(), rotation)
assert rotated.getpixel(expected_corner) == (255, 0, 0)
def test_decoded_frames_apply_the_container_display_matrix(monkeypatch, tmp_path):
"""A -90 phone clip must reach the model upright, not on its side."""
demo = _load_demo_module()
source = _asymmetric_image()
class _Frame:
rotation = -90
def to_image(self):
return source
class _Container:
streams = type("_Streams", (), {"video": [type("_Stream", (), {})()]})()
def decode(self, stream):
yield _Frame()
def __enter__(self):
return self
def __exit__(self, *exc):
return False
monkeypatch.setitem(sys.modules, "av", type("_Av", (), {"open": staticmethod(lambda path: _Container())}))
(image,) = demo._decode_video_frames_rgb(tmp_path / "clip.mp4", [0])
assert image.size == (2, 4)
assert image.getpixel((1, 0)) == (255, 0, 0)
def test_streaming_output_writer_persists_audio_deltas_as_they_arrive(tmp_path, capsys):
demo = _load_demo_module()
writer = demo._StreamingOutputWriter(tmp_path)
+10 -2
View File
@@ -150,8 +150,16 @@ def _print_device_processes() -> None:
elif current_omni_platform.is_npu():
_run_smi("Ascend NPU Information (npu-smi info)", ["npu-smi", "info"], 40)
elif current_omni_platform.is_rocm():
_run_smi("AMD GPU Information (amd-smi)", ["amd-smi"], 30)
_run_smi("Detailed AMD GPU Processes (amd-smi process)", ["amd-smi", "process"], 100, timeout=3)
# amd-smi can enter uninterruptible sleep in the amdgpu CPER ioctl
# (amdgpu_cper_ring_write). In that state subprocess.run's timeout
# cannot reap it, causing this diagnostic hook to hang the entire test.
# rocm-smi provides the information needed here without that ioctl.
_run_smi(
"AMD GPU Information (rocm-smi)",
["rocm-smi", "--showuse", "--showmeminfo", "vram"],
60,
)
_run_smi("Detailed AMD GPU Processes (rocm-smi)", ["rocm-smi", "--showpids"], 100, timeout=3)
elif current_omni_platform.is_xpu():
_run_smi("Intel XPU Information (xpu-smi discovery)", ["xpu-smi", "discovery"], 40)
elif current_omni_platform.is_musa():
+13
View File
@@ -11,6 +11,8 @@ from typing import Any
import yaml
from vllm_omni.config.stage_config import load_deploy_config
def modify_stage_config(
yaml_path: str,
@@ -669,6 +671,16 @@ def get_deploy_config_stage(rel_path: str, stage_id: int) -> dict[str, Any]:
raise KeyError(f"No stage_id={stage_id} in deploy config {rel_path!r}")
def get_deploy_duplex_max_sessions(rel_path: str) -> int:
"""Return the duplex session capacity a deploy yaml admits.
Loads through ``load_deploy_config`` so ``base_config`` merging and
``DuplexSessionRuntimeConfig`` defaults stay in lockstep with the server
instead of being re-implemented here.
"""
return load_deploy_config(get_deploy_config_path(rel_path)).duplex_session.max_sessions
def _stage_ids_from_deploy_yaml(stage_config_path: str) -> list[int]:
"""Return ``stage_id`` values from a new-schema deploy YAML (``stages``)."""
with open(stage_config_path, encoding="utf-8") as f:
@@ -734,6 +746,7 @@ def stage_config_path_for_run_level(stage_config_path: str | None, run_level: st
__all__ = [
"get_deploy_config_path",
"get_deploy_config_stage",
"get_deploy_duplex_max_sessions",
"get_stage_entries",
"load_stage_ids",
"load_stage_replica_counts",
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,119 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Unit tests for MammothModa2 config classes.
These tests verify the correct behavior of the Mammothmoda2Config class,
particularly the get_text_config() method which is called during parent
class initialization (PretrainedConfig.__init__ -> validation -> get_text_config).
Regression test for: AttributeError: 'Mammothmoda2Config' object has no attribute 'llm_config'
"""
import pytest
from vllm_omni.transformers_utils.configs.mammoth_moda2 import (
Mammothmoda2Config,
Mammothmoda2Qwen2_5_VLConfig,
Mammothmoda2Qwen2_5_VLTextConfig,
Mammothmoda2Qwen3VLConfig,
Mammothmoda2Qwen3VLTextConfig,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestMammothmoda2Config:
"""Tests for Mammothmoda2Config class."""
def test_get_text_config_returns_nested_text_config(self):
"""
Verify that get_text_config() returns the nested text_config from llm_config.
This is a regression test for the AttributeError that occurred when
get_text_config() was called during parent class initialization before
llm_config was set.
"""
config = Mammothmoda2Config(
llm_config={
"model_type": "mammothmoda2_qwen2_5_vl",
}
)
assert config.llm_config is not None
assert isinstance(config.llm_config, Mammothmoda2Qwen2_5_VLConfig)
text_config = config.get_text_config()
assert text_config is not None
assert isinstance(text_config, Mammothmoda2Qwen2_5_VLTextConfig)
# The text_config should be the same object as llm_config.text_config
assert text_config is config.llm_config.text_config
def test_get_text_config_returns_none_when_llm_config_is_none(self):
"""Verify that get_text_config() returns None when llm_config is None."""
config = Mammothmoda2Config(llm_config=None)
assert config.llm_config is None
assert config.get_text_config() is None
def test_init_sets_llm_config_before_parent_init(self):
"""
Verify that llm_config is set before parent __init__ is called.
This is critical because the parent PretrainedConfig.__init__ calls
validation which in turn calls get_text_config(). If llm_config is not
set before super().__init__(), an AttributeError will be raised.
"""
# This should not raise an AttributeError
config = Mammothmoda2Config(
llm_config={
"model_type": "mammothmoda2_qwen2_5_vl",
}
)
assert config.llm_config is not None
class TestMammothmoda2Qwen2_5_VLConfig:
"""Tests for Mammothmoda2Qwen2_5_VLConfig class."""
def test_has_text_config_attribute(self):
"""Verify that Mammothmoda2Qwen2_5_VLConfig has text_config attribute."""
config = Mammothmoda2Qwen2_5_VLConfig()
assert hasattr(config, "text_config")
assert isinstance(config.text_config, Mammothmoda2Qwen2_5_VLTextConfig)
def test_text_config_has_expected_attributes(self):
"""Verify that the nested text_config has expected attributes for token validation."""
config = Mammothmoda2Qwen2_5_VLConfig()
text_config = config.text_config
assert hasattr(text_config, "vocab_size")
assert hasattr(text_config, "image_token_id")
assert hasattr(text_config, "video_token_id")
class TestMammothmoda2Qwen3VLConfig:
def test_dev_config_uses_qwen3_vl_subconfigs(self):
config = Mammothmoda2Config(
llm_config={
"model_type": "mammothmoda2_qwen3_vl",
"text_config": {
"model_type": "mammothmoda2_qwen3_vl_text",
"vocab_size": 151936,
"gen_vocab_size": 32800,
"gen_vocab_start_index": 152064,
},
"vision_config": {
"model_type": "mammothmoda2_qwen3_vl_vision",
"deepstack_visual_indexes": [8, 16, 24],
},
}
)
assert isinstance(config.llm_config, Mammothmoda2Qwen3VLConfig)
assert isinstance(config.get_text_config(), Mammothmoda2Qwen3VLTextConfig)
# vLLM validates sampling against the complete base + visual vocabulary.
assert config.get_text_config().vocab_size == 152064 + 32800
assert config.get_text_config().gen_vocab_start_index == 152064
assert config.llm_config.vision_config.deepstack_visual_indexes == [8, 16, 24]
@@ -0,0 +1,92 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""What is allowed to keep the CFM wrapper from capturing a graph.
These run without a GPU: ``_capture`` is mocked, so nothing here touches a
CUDA context. The point is the decision made *before* capture, not the capture.
"""
from typing import NamedTuple
from unittest.mock import Mock
import pytest
import torch
from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import (
CFMGraphWrapper,
_format_memory_delta,
_memory_snapshot,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _InputSignature(NamedTuple):
"""The three properties ``replay`` reads from an input before capturing.
Everything ``replay`` does ahead of the capture decision -- the device
check and the cache key -- comes from these, and ``_capture`` is mocked
out, so the tests need no CUDA tensors and no GPU.
"""
shape: tuple[int, ...]
dtype: torch.dtype
device: torch.device
def _cfm_input_signatures(chunk_size: int) -> tuple[_InputSignature, ...]:
device = torch.device("cuda:0")
shapes = ((2, 16, chunk_size), (2, 1, 8), (2, 2, 8, 2), (2, 2, 1, 4, 8), (2, 2, 8, 2), (2, 2, 1, 12, 8))
return tuple(_InputSignature(shape, torch.float32, device) for shape in shapes)
def _wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 4) -> CFMGraphWrapper:
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
wrapper = object.__new__(CFMGraphWrapper)
wrapper.max_graphs = max_graphs
wrapper.enabled = True
wrapper.graph_fn = Mock(return_value=torch.tensor([42.0]))
wrapper.device = torch.device("cuda:0")
wrapper._cache = {}
wrapper._unsupported = set()
wrapper._stats = {"calls": 0, "hits": 0, "captures": 0, "flushes": 0, "eager": 0}
wrapper._capture = Mock(return_value=None)
return wrapper
def test_capture_is_not_gated_on_free_device_memory(monkeypatch: pytest.MonkeyPatch) -> None:
"""Whether to capture must not be decided from free device memory.
A capture is served by the caching allocator out of its private graph
pool, and the allocator falls back on memory it has already reserved. Free
device memory is therefore neither an upper nor a lower bound on what a
capture can draw, and a threshold on it decides the wrong question -- one
whose answer also silently varies with card size and with whatever else
shares the device. Anything consulting it here would fail this test.
"""
import vllm_omni.platforms as platforms_module
def _saturated(device: torch.device | None = None) -> int:
raise AssertionError("capture must not depend on device-level free memory")
# Setting the attribute on the module shadows the lazy ``__getattr__``, so
# this neither builds nor requires a real platform singleton.
monkeypatch.setattr(
platforms_module,
"current_omni_platform",
Mock(get_free_memory=_saturated),
raising=False,
)
wrapper = _wrapper(monkeypatch)
wrapper.replay(*_cfm_input_signatures(10))
wrapper._capture.assert_called_once()
def test_memory_reporting_is_inert_off_cuda() -> None:
"""Capture logging must not break the platforms that have no such counters."""
assert _memory_snapshot(torch.device("cpu")) is None
assert _format_memory_delta(None, (1, 2)) == ""
assert _format_memory_delta((1, 2), None) == ""
assert "allocated" in _format_memory_delta((0, 0), (1 << 20, 2 << 20))
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from functools import lru_cache
import sys
from types import SimpleNamespace
from unittest.mock import Mock
@@ -10,6 +10,7 @@ import torch
import torch.nn as nn
from vllm.platforms import current_platform
import vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper as wrapper_module
from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
HiFTGenerator,
)
@@ -248,50 +249,203 @@ def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes(
torch.testing.assert_close(graph_cnn, eager_inputs[4], rtol=1e-4, atol=1e-5)
torch.testing.assert_close(graph_att, eager_inputs[5], rtol=1e-4, atol=1e-5)
wrapper._capture_graph.cache_clear()
wrapper._flush()
def test_importing_wrapper_does_not_resolve_platform() -> None:
"""Importing this module must not build the OmniPlatform singleton.
The NPU platform's ``__init__`` patches Code2Wav, which imports
``batched_token2wav`` -> ``cuda_graph_wrapper``. Resolving
``current_omni_platform`` at module scope re-enters
``platforms.__getattr__`` while the singleton is still under construction
and the import graph deadlocks. The module reaches for nothing on
``vllm_omni.platforms`` today; this keeps it that way.
"""
import importlib
import vllm_omni.platforms as platforms_module
name = "vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper"
importlib.import_module(name)
saved = platforms_module._current_omni_platform
try:
del sys.modules[name]
platforms_module._current_omni_platform = None
importlib.import_module(name)
assert platforms_module._current_omni_platform is None
finally:
platforms_module._current_omni_platform = saved
def _cfm_mock_wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 1) -> CFMGraphWrapper:
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
wrapper = object.__new__(CFMGraphWrapper)
wrapper.max_graphs = max_graphs
wrapper.enabled = True
wrapper.graph_fn = Mock(return_value=torch.tensor([42.0]))
wrapper.device = torch.device("cuda")
@lru_cache(maxsize=max_graphs)
def _capture_graph(key: tuple):
return wrapper._do_capture(key)
wrapper._capture_graph = _capture_graph
wrapper._do_capture = Mock(return_value=None)
wrapper._cache = {}
wrapper._unsupported = set()
wrapper._stats = {"calls": 0, "hits": 0, "captures": 0, "flushes": 0, "eager": 0}
wrapper._capture = Mock(return_value=None)
return wrapper
def test_cfm_unseen_shape_is_lazily_captured(monkeypatch: pytest.MonkeyPatch) -> None:
wrapper = _cfm_mock_wrapper(monkeypatch)
wrapper._capture_graph = Mock(return_value=None)
inputs = _cfm_inputs(2, 10, 0)
wrapper.replay(*inputs)
wrapper._capture_graph.assert_called_once()
wrapper._capture.assert_called_once()
wrapper.graph_fn.assert_called_once()
def test_cfm_capture_failure_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None:
def test_cfm_returning_no_entry_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None:
"""A capture that yields no entry must still serve the request eagerly.
``_capture`` is mocked here, so this says nothing about ``_disable``; see
``test_cfm_capture_failure_disables_further_capture`` for that.
"""
wrapper = _cfm_mock_wrapper(monkeypatch)
wrapper._capture_graph = lru_cache(maxsize=1)(lambda key: None)
inputs = _cfm_inputs(2, 10, 0)
result = wrapper.replay(*inputs)
wrapper.graph_fn.assert_called_once()
assert result[0] is wrapper.graph_fn.return_value
assert wrapper._stats["eager"] == 1
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None:
"""Verify that LRU cache evicts oldest entries when full."""
def test_cfm_capture_failure_disables_further_capture(monkeypatch: pytest.MonkeyPatch) -> None:
"""A real capture failure must stop the wrapper capturing for good.
A failed capture can leave the capture stream current and the allocator
still routing into the graph pool, so the next shape must not try again.
"""
pool = torch.cuda.graph_pool_handle()
monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
torch.manual_seed(0)
estimator = _MiniDiT().eval().cuda()
wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=4)
def _explode(*args: object, **kwargs: object) -> None:
raise RuntimeError("capture failed")
monkeypatch.setattr(torch.cuda, "graph", _explode)
with torch.inference_mode():
wrapper.replay(*_cfm_inputs(2, 10, 0))
assert wrapper.enabled is False
assert wrapper._cache == {}
captures_after_failure = wrapper._stats["captures"]
wrapper.replay(*_cfm_inputs(2, 12, 0))
assert wrapper._stats["captures"] == captures_after_failure
assert wrapper._stats["eager"] == 2
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_cfm_unsupported_dtype_eagers_only_that_shape(monkeypatch: pytest.MonkeyPatch) -> None:
"""A key that cannot round-trip is a property of one shape, not of the GPU.
It must not disable the wrapper or retire the generation the way a capture
failure does.
"""
pool = torch.cuda.graph_pool_handle()
monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
torch.manual_seed(0)
estimator = _MiniDiT().eval().cuda()
wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=4)
unbuildable_width = 10
real_tensors_from_key = wrapper_module._tensors_from_key
def _reject_one_shape(key: tuple) -> tuple:
if key[1][0][2] == unbuildable_width:
raise KeyError("torch.int64")
return real_tensors_from_key(key)
monkeypatch.setattr(wrapper_module, "_tensors_from_key", _reject_one_shape)
with torch.inference_mode():
wrapper.replay(*_cfm_inputs(2, unbuildable_width, 0))
assert wrapper.enabled is True
assert wrapper._stats["captures"] == 0
# a capturable shape still gets a graph
wrapper.replay(*_cfm_inputs(2, 12, 0))
assert wrapper._stats["captures"] == 1
# and the rejected shape stays eager without retrying the capture
wrapper.replay(*_cfm_inputs(2, unbuildable_width, 0))
assert wrapper._stats["captures"] == 1
assert wrapper._stats["eager"] == 2
wrapper._flush()
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_hift_replay_survives_a_cfm_generation_flush(monkeypatch: pytest.MonkeyPatch) -> None:
"""Retiring a CFM generation must not disturb the vocoder.
Both wrappers capture into ``get_global_graph_pool()``, so the HiFT graphs
are exactly the live graphs a CFM flush could strand.
"""
pool = torch.cuda.graph_pool_handle()
monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
torch.manual_seed(0)
hift = _small_hift()
token2wav = SimpleNamespace(
hift=hift,
flow=SimpleNamespace(
encoder=SimpleNamespace(pre_lookahead_layer=SimpleNamespace(pre_lookahead_len=3)),
token_mel_ratio=2,
),
mel_cache_len=2,
source_cache_len=960,
)
hift_wrapper = HiFTGraphWrapper(
token2wav,
connector_config={"codec_chunk_frames": 2, "codec_left_context_frames": 3},
capture_batch_sizes=[1],
)
hift_wrapper.capture()
estimator = _MiniDiT().eval().cuda()
cfm = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=2)
speech_feat = torch.randn(1, 80, 4, device="cuda")
cache_source = torch.zeros(1, 1, 0, device="cuda")
with torch.inference_mode():
expected_speech, expected_source = hift.inference(speech_feat, cache_source)
for width in (10, 12, 14, 16):
cfm.replay(*_cfm_inputs(2, width, 0))
assert cfm._stats["flushes"] >= 1, "cache never overflowed; the test proves nothing"
actual_speech, actual_source = hift_wrapper.replay(speech_feat, cache_source)
torch.testing.assert_close(actual_speech, expected_speech, rtol=1e-4, atol=1e-5)
torch.testing.assert_close(actual_source, expected_source, rtol=1e-4, atol=1e-5)
cfm._flush()
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_cfm_cache_flushes_whole_generation(monkeypatch: pytest.MonkeyPatch) -> None:
"""A full cache is retired all at once, never one graph at a time.
Destroying a single graph while its peers stay live strands them, so the
cache must never hold a graph that outlived one of its generation-mates.
"""
pool = torch.cuda.graph_pool_handle()
monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
@@ -302,26 +456,32 @@ def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None:
with torch.inference_mode():
inputs_a = _cfm_inputs(2, 10, 0)
wrapper.replay(*inputs_a)
assert wrapper._capture_graph.cache_info().currsize == 1
assert len(wrapper._cache) == 1
inputs_b = _cfm_inputs(2, 12, 0)
wrapper.replay(*inputs_b)
assert wrapper._capture_graph.cache_info().currsize == 2
assert len(wrapper._cache) == 2
# A hit must not grow the cache or trigger a flush.
wrapper.replay(*inputs_a)
assert wrapper._capture_graph.cache_info().currsize == 2
assert len(wrapper._cache) == 2
assert wrapper._stats["flushes"] == 0
assert wrapper._stats["hits"] == 1
# The third distinct shape flushes the generation, then captures alone.
inputs_c = _cfm_inputs(2, 14, 0)
wrapper.replay(*inputs_c)
assert wrapper._capture_graph.cache_info().currsize == 2
assert wrapper._stats["flushes"] == 1
assert len(wrapper._cache) == 1
wrapper.replay(*inputs_b)
hits_before = wrapper._capture_graph.cache_info().hits
wrapper.replay(*inputs_b)
hits_after = wrapper._capture_graph.cache_info().hits
assert hits_after == hits_before + 1
# The flushed shapes are gone, so they capture again rather than hit.
hits_before = wrapper._stats["hits"]
wrapper.replay(*inputs_a)
assert wrapper._stats["hits"] == hits_before
wrapper.replay(*inputs_a)
assert wrapper._stats["hits"] == hits_before + 1
wrapper._capture_graph.cache_clear()
wrapper._flush()
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
@@ -14,6 +14,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
_REPO_ROOT = Path(__file__).resolve().parents[3]
_STALE_API_FILES = [
"vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py",
"vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py",
"vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py",
"vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py",
]
@@ -117,6 +118,10 @@ def test_auto_weights_loader_delegated_loaders_accept_mapped_cache_scale(
@pytest.mark.parametrize(
("module_name", "class_name"),
[
(
"vllm_omni.model_executor.models.mammoth_moda2.mammoth_moda2",
"MammothModa2Qwen2ForCausalLM",
),
(
"vllm_omni.model_executor.models.mimo_audio.mimo_audio_llm",
"MiMoAudioLLMForConditionalGeneration",
@@ -0,0 +1,62 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
from unittest.mock import patch
import pytest
import torch
from vllm.model_executor.models.qwen3_vl import Qwen3VLForConditionalGeneration
from vllm_omni.model_executor.models.minimax_h3.text_encoder import MiniMaxH3TextEncoder
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _model_without_weights() -> MiniMaxH3TextEncoder:
model = MiniMaxH3TextEncoder.__new__(MiniMaxH3TextEncoder)
model.config = SimpleNamespace(
vision_start_token_id=1,
vision_end_token_id=2,
image_token_id=3,
video_token_id=4,
)
model._token_tags = None
return model
def test_embed_failure_clears_token_tags():
model = _model_without_weights()
model._token_tags = torch.tensor([0])
with (
patch.object(Qwen3VLForConditionalGeneration, "embed_input_ids", side_effect=RuntimeError("failed")),
pytest.raises(RuntimeError, match="failed"),
):
model.embed_input_ids(torch.tensor([5, 3]))
assert model._token_tags is None
def test_invalid_output_length_consumes_token_tags():
model = _model_without_weights()
model._token_tags = torch.tensor([1, 0])
with pytest.raises(RuntimeError, match="token tags exceed"):
model.make_omni_output(torch.zeros(1, 8))
assert model._token_tags is None
def test_make_omni_output_uses_shared_stage_payload_fields():
model = _model_without_weights()
model._token_tags = torch.tensor([1, 0])
hidden = torch.zeros(2, 5120)
output = model.make_omni_output(hidden)
assert output.multimodal_outputs is not None
assert torch.equal(output.multimodal_outputs["hidden_states"]["output"], hidden)
assert torch.equal(output.multimodal_outputs["meta"]["token_role_ids"], torch.tensor([[1], [0]]))
assert "encoder_hidden_states" not in output.multimodal_outputs
assert "token_tags" not in output.multimodal_outputs
@@ -0,0 +1,387 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression tests for MiniMax H3's disaggregated text-encoder contract."""
import os
import shutil
from types import SimpleNamespace
from unittest.mock import Mock
import numpy as np
import pytest
import torch
from PIL import Image
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_load_audio,
resolve_minimax_h3_diffusion_model_path,
)
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
resolve_minimax_h3_model_root,
resolve_minimax_h3_partition,
)
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
MINIMAX_H3_CONDITION_LABELS_KEY,
MINIMAX_H3_PRESENTATION_TASK_KEY,
)
from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
minimax_h3_ref2va_presentation,
minimax_h3_ref2va_video_presentation,
)
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY,
deserialize_prepared_reference_videos,
)
from vllm_omni.model_executor.models.minimax_h3.text_encoder import (
MiniMaxH3MultiModalProcessor,
_build_minimax_h3_presentation,
)
from vllm_omni.model_executor.stage_input_processors.minimax_h3 import (
_audio_items,
_diffusion_sampling_params,
prepare_text_encoder_prompt,
text_encoder2diffusion,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _SegmentTokenizer:
_special_ids = {
"<|vision_start|>": 1,
"<|vision_end|>": 2,
"<|image_pad|>": 3,
"<|video_pad|>": 4,
}
def __call__(self, text, *, add_special_tokens):
assert add_special_tokens is False
return {"input_ids": [100 + len(text), 1000 + sum(text.encode())]}
def convert_tokens_to_ids(self, token):
return self._special_ids[token]
def test_h3_processor_reprocesses_media_instead_of_using_partial_sender_cache(mocker):
processor = object.__new__(MiniMaxH3MultiModalProcessor)
sentinel = ([1, 2, 3], object(), True)
apply_processor = mocker.patch.object(
MiniMaxH3MultiModalProcessor,
"_apply_hf_processor",
return_value=sentinel,
)
inputs = object()
timing_ctx = object()
result = processor._cached_apply_hf_processor(inputs, timing_ctx)
assert result is sentinel
apply_processor.assert_called_once_with(inputs, timing_ctx)
@pytest.mark.parametrize(
("value", "expected_count"),
[
((torch.zeros(16), 16_000), 1),
([np.zeros(16), 16_000], 1),
([(torch.zeros(16), 16_000), (torch.ones(16), 24_000)], 2),
(["first.wav", "second.wav"], 2),
],
)
def test_audio_items_preserves_waveform_pairs(value, expected_count):
assert len(_audio_items(value)) == expected_count
def test_h3_selects_the_single_explicit_diffusion_stage_params() -> None:
stage_zero = SimpleNamespace(extra_args={"task": "t2va"})
diffusion = OmniDiffusionSamplingParams(extra_args={"task": "ref2va"})
assert _diffusion_sampling_params([stage_zero, diffusion]) is diffusion
def test_h3_rejects_missing_or_ambiguous_diffusion_stage_params() -> None:
with pytest.raises(RuntimeError, match="exactly one OmniDiffusionSamplingParams"):
_diffusion_sampling_params([SimpleNamespace(extra_args={"task": "t2va"})])
with pytest.raises(RuntimeError, match="got 2"):
_diffusion_sampling_params([OmniDiffusionSamplingParams(), OmniDiffusionSamplingParams()])
def test_fused_audio_loader_accepts_list_waveform_pair():
waveform, sample_rate = _load_audio([[0.0, 0.5, -0.5], 16_000])
assert sample_rate == 16_000
torch.testing.assert_close(waveform, torch.tensor([0.0, 0.5, -0.5]))
def test_prepare_ref2va_keeps_original_text_and_exact_condition_order():
prompt = {
"prompt": "hello",
"multi_modal_data": {
"image": Image.new("RGB", (256, 256)),
"audio": [np.zeros(16), 16_000],
},
}
sampling = OmniDiffusionSamplingParams(
height=256,
width=448,
extra_args={"task": "ref2va"},
)
transformed = prepare_text_encoder_prompt(prompt, [sampling])
assert transformed["prompt"] == "hello"
assert len(transformed["multi_modal_data"]["image"]) == 1
assert "audio" not in transformed["multi_modal_data"]
assert transformed["mm_processor_kwargs"][MINIMAX_H3_PRESENTATION_TASK_KEY] == "ref2va"
assert transformed["mm_processor_kwargs"][MINIMAX_H3_CONDITION_LABELS_KEY] == [
("image", 1),
("audio", 1),
]
def test_text_encoder_prompt_rejects_injected_prepared_video_descriptor():
prompt = {
"prompt": "hello",
"additional_information": {"meta": {MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY: '{"artifact_dir":"/tmp/user"}'}},
}
sampling = OmniDiffusionSamplingParams(height=256, width=448, extra_args={"task": "t2va"})
transformed = prepare_text_encoder_prompt(prompt, [sampling])
assert MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY not in transformed["additional_information"]["meta"]
def test_prepare_ref2va_video_uses_shared_frame_sampler_once(monkeypatch):
prepare_videos = Mock()
def prepare(_videos, *, target_frame_count, workdir, start_time_seconds):
prepare_videos(_videos, target_frame_count, workdir, start_time_seconds)
prepared_path = os.path.join(workdir, "prepared.mp4")
open(prepared_path, "wb").close()
return [
{
"original_path": "/tmp/input.mp4",
"prepared_path": prepared_path,
"input_has_audio": True,
"width": 448,
"height": 256,
"start_time_seconds": 0.0,
"duration_seconds": 5.2,
"audio_duration_seconds": 5.2,
}
]
sample_frames = Mock(return_value={"frames": [np.zeros((4, 4, 3), dtype=np.uint8)]})
monkeypatch.setattr("vllm_omni.model_executor.stage_input_processors.minimax_h3.prepare_reference_videos", prepare)
monkeypatch.setattr(
"vllm_omni.model_executor.stage_input_processors.minimax_h3.sample_reference_video_frames",
sample_frames,
)
prompt = {
"prompt": "hello",
"multi_modal_data": {"video": "/tmp/input.mp4"},
}
sampling = OmniDiffusionSamplingParams(height=256, width=448, num_frames=124, extra_args={"task": "ref2va"})
transformed = prepare_text_encoder_prompt(prompt, [sampling])
prepare_videos.assert_called_once()
_, target_frame_count, artifact_dir, start_time_seconds = prepare_videos.call_args.args
assert target_frame_count == 124
assert start_time_seconds is None
descriptor = transformed["additional_information"]["meta"][MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY]
described_dir, described_videos = deserialize_prepared_reference_videos(descriptor)
assert described_dir == artifact_dir
assert os.path.isfile(described_videos[0]["prepared_path"])
sample_frames.assert_called_once_with(described_videos[0]["prepared_path"])
assert prompt["multi_modal_data"] == {"video": "/tmp/input.mp4"}
assert transformed["mm_processor_kwargs"][MINIMAX_H3_CONDITION_LABELS_KEY] == [
("audio", 1),
("video", 1),
]
shutil.rmtree(artifact_dir)
def test_stage_wire_rejects_multiple_text_encoder_sources():
with pytest.raises(RuntimeError, match="exactly one text-encoder source"):
text_encoder2diffusion([SimpleNamespace(), SimpleNamespace()], prompt={"prompt": "hello"})
def test_stage_wire_rejects_request_id_mismatch():
source = SimpleNamespace(request_id="other", outputs=[SimpleNamespace()])
prompt = {
"prompt": "hello",
"additional_information": {"global_request_id": ["req-1"]},
}
with pytest.raises(RuntimeError, match="request ID does not match"):
text_encoder2diffusion([source], prompt=prompt)
def test_stage_wire_rejects_multiple_completions():
source = SimpleNamespace(request_id="req-1", outputs=[SimpleNamespace(), SimpleNamespace()])
with pytest.raises(RuntimeError, match="exactly one completion"):
text_encoder2diffusion([source], prompt={"prompt": "hello"})
def _source_output(payload: dict) -> SimpleNamespace:
completion = SimpleNamespace(multimodal_output=payload)
return SimpleNamespace(request_id="request-1", outputs=[completion])
def _diffusion_prompt() -> dict:
return {
"request_id": "request-1",
"prompt": "test prompt",
"additional_information": {},
}
def test_text_encoder2diffusion_reads_hidden_states_output_and_token_role_ids() -> None:
hidden = torch.randn(4, 5120)
token_role_ids = torch.tensor([[1], [1], [0], [0]])
result = text_encoder2diffusion(
[
_source_output(
{
"hidden_states": {"output": hidden},
"meta": {"token_role_ids": token_role_ids},
}
)
],
_diffusion_prompt(),
)
conditioning = result["additional_information"]["text_encoder_output"]
torch.testing.assert_close(conditioning["hidden_states"], hidden)
torch.testing.assert_close(conditioning["token_tags"], token_role_ids.squeeze(-1))
@pytest.mark.parametrize(
("payload", "message"),
[
({"meta": {"token_role_ids": torch.ones(4, 1)}}, "no hidden_states payload"),
({"hidden_states": {}, "meta": {"token_role_ids": torch.ones(4, 1)}}, "no hidden_states.output tensor"),
({"hidden_states": {"output": torch.randn(4, 5120)}}, "no conditioning metadata"),
({"hidden_states": {"output": torch.randn(4, 5120)}, "meta": {}}, "no token_role_ids tensor"),
(
{
"hidden_states": {"output": torch.randn(4, 5120)},
"meta": {"token_role_ids": torch.ones(4)},
},
r"must have shape \[tokens, 1\]",
),
],
)
def test_text_encoder2diffusion_rejects_invalid_conditioning_payload(payload: dict, message: str) -> None:
with pytest.raises(RuntimeError, match=message):
text_encoder2diffusion([_source_output(payload)], _diffusion_prompt())
def test_ref2va_one_image_tokens_and_tags_match_fused_presentation():
tokenizer = _SegmentTokenizer()
labels = [("image", 1), ("audio", 1)]
image_grid = torch.tensor([[1, 4, 4]])
actual = _build_minimax_h3_presentation(
tokenizer,
prompt="hello",
task="ref2va",
condition_labels=labels,
image_grid_thw=image_grid,
video_grid_thw=None,
video_timestamps=None,
merge_size=2,
)
expected = minimax_h3_ref2va_presentation(
tokenizer,
prompt="hello",
condition_labels=labels,
image_token_count=[4],
)
torch.testing.assert_close(actual[0], expected[0])
torch.testing.assert_close(actual[1], expected[1])
def test_ref2va_video_tokens_and_tags_match_fused_without_outer_markers():
tokenizer = _SegmentTokenizer()
labels = [("audio", 1), ("video", 1)]
video_grid = torch.tensor([[2, 4, 4]])
timestamps = [[0.2, 0.4]]
actual = _build_minimax_h3_presentation(
tokenizer,
prompt="hello",
task="ref2va",
condition_labels=labels,
image_grid_thw=None,
video_grid_thw=video_grid,
video_timestamps=timestamps,
merge_size=2,
)
expected = minimax_h3_ref2va_video_presentation(
tokenizer,
prompt="hello",
condition_labels=labels,
image_token_count=None,
video_block_token_counts=[[4, 4]],
video_block_timestamps=timestamps,
)
torch.testing.assert_close(actual[0], expected[0])
torch.testing.assert_close(actual[1], expected[1])
assert int((actual[0] == tokenizer._special_ids["<|vision_start|>"]).sum()) == 2
assert int((actual[0] == tokenizer._special_ids["<|vision_end|>"]).sum()) == 2
def test_checkpoint_resolver_selects_local_partition(tmp_path):
root = tmp_path / "MiniMax-H3"
(root / "FL2VA" / "text_encoder").mkdir(parents=True)
(root / "Ref2VA" / "text_encoder").mkdir(parents=True)
assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder")
assert resolve_minimax_h3_model_root(str(root), None, "ref2va") == str(root / "Ref2VA" / "text_encoder")
assert resolve_minimax_h3_model_root(str(root / "Ref2VA"), None, None) == str(root / "Ref2VA" / "text_encoder")
assert resolve_minimax_h3_model_root(str(root / "FL2VA"), None, "ref2va") == str(root / "Ref2VA" / "text_encoder")
def test_checkpoint_resolver_rejects_unknown_task(tmp_path):
with pytest.raises(ValueError, match="task_type must be one of"):
resolve_minimax_h3_model_root(str(tmp_path), None, "unknown")
def test_partition_resolver_preserves_consumer_auto_default(tmp_path):
root = tmp_path / "MiniMax-H3"
ref2va = root / "Ref2VA"
ref2va.mkdir(parents=True)
assert resolve_minimax_h3_partition(str(root), "auto", auto_partition="fl2va") == "fl2va"
assert resolve_minimax_h3_partition(str(root), "auto", auto_partition="combined") == "combined"
assert resolve_minimax_h3_partition(str(ref2va), "auto", auto_partition="combined") == "ref2va"
def test_diffusion_resolver_selects_startup_partition(tmp_path):
root = tmp_path / "MiniMax-H3"
fl2va = root / "FL2VA"
ref2va = root / "Ref2VA"
fl2va.mkdir(parents=True)
ref2va.mkdir()
(fl2va / "model_index.json").write_text("{}")
(ref2va / "model_index.json").write_text("{}")
assert resolve_minimax_h3_diffusion_model_path(str(root), None, "fl2va") == str(fl2va)
assert resolve_minimax_h3_diffusion_model_path(str(root), None, "ref2va") == str(ref2va)
assert resolve_minimax_h3_diffusion_model_path(str(root), None, None) == str(fl2va)
assert resolve_minimax_h3_diffusion_model_path(str(root), None, "combined") == str(root)
assert resolve_minimax_h3_diffusion_model_path(str(ref2va), None, None) == str(ref2va)
def test_diffusion_resolver_normalizes_partial_partition_directory(tmp_path):
root = tmp_path / "MiniMax-H3"
ref2va = root / "Ref2VA"
(ref2va / "text_encoder").mkdir(parents=True)
assert resolve_minimax_h3_diffusion_model_path(str(ref2va), None, "ref2va") == str(ref2va)
@@ -414,8 +414,10 @@ def test_talker_drains_all_received_positions_per_wake() -> None:
session = talker._sessions["req-0"]
assert session["step"] == 1 and not session["sync_mode"]
# First decode wake: drains ALL prompt-region steps (t=1..P-1) at once.
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
# First resumed wake: drains ALL prompt-region steps (t=1..P-1) at once.
# vLLM 0.28 reports each resumed one-token segment as prefill; the model
# session must still advance instead of restarting its EAR-TTS warmup.
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=True, **info)
assert session["step"] == prompt_len # t advanced 1 -> P in one wake
assert update["codes"]["audio"].shape == (prompt_len - 1, 31)
assert float(embeds[0, 0]) == 0.0 # offline stream waits for upstream terminal
@@ -8,9 +8,12 @@ import pytest
import torch
from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import (
AdaptiveChunkController,
compute_adaptive_emit,
compute_dynamic_initial_chunk_size,
compute_ramp_emit,
max_ic_for_chunk_size,
parse_adaptive_config,
parse_chunk_ramp,
ramp_chunk_size,
ramp_cumulative,
@@ -41,7 +44,17 @@ def _req(rid, *, finished, initial_codec_chunk_frames=None):
)
def _tm(*, chunk_frames=25, left_context=25, max_num_seqs=1, initial_chunk_frames=0, chunk_ramp=None):
def _tm(
*,
chunk_frames=25,
left_context=25,
max_num_seqs=1,
initial_chunk_frames=0,
chunk_ramp=None,
adaptive=False,
adaptive_min=2,
adaptive_margin=200.0,
):
extra = {
"codec_chunk_frames": chunk_frames,
"codec_left_context_frames": left_context,
@@ -49,11 +62,16 @@ def _tm(*, chunk_frames=25, left_context=25, max_num_seqs=1, initial_chunk_frame
}
if chunk_ramp is not None:
extra["codec_chunk_ramp"] = chunk_ramp
if adaptive:
extra["codec_chunk_adaptive"] = True
extra["codec_chunk_min_frames"] = adaptive_min
extra["codec_chunk_safety_margin_ms"] = adaptive_margin
return SimpleNamespace(
code_prompt_token_ids=defaultdict(list),
scheduler_max_num_seqs=max_num_seqs,
put_req_chunk=defaultdict(int),
ramp_chunk_count=defaultdict(int),
_adaptive_states={},
request_payload={},
connector=SimpleNamespace(config={"extra": extra}),
)
@@ -605,18 +623,32 @@ class TestRampHelpers:
def test_parse_ramp_warns_on_tail_mismatch(self, caplog):
import logging
with caplog.at_level(logging.WARNING):
target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.chunk_size_utils")
target_logger.addHandler(caplog.handler)
prev_level = target_logger.level
target_logger.setLevel(logging.WARNING)
try:
result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8]}, steady=25)
finally:
target_logger.removeHandler(caplog.handler)
target_logger.setLevel(prev_level)
assert result == [4, 4, 8]
assert any("reintroduces" in r.message for r in caplog.records)
assert "reintroduces" in caplog.text
def test_parse_ramp_no_warn_on_tail_match(self, caplog):
import logging
with caplog.at_level(logging.WARNING):
target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.chunk_size_utils")
target_logger.addHandler(caplog.handler)
prev_level = target_logger.level
target_logger.setLevel(logging.WARNING)
try:
result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8, 16, 25]}, steady=25)
finally:
target_logger.removeHandler(caplog.handler)
target_logger.setLevel(prev_level)
assert result == [4, 4, 8, 16, 25]
assert not any("reintroduces" in r.message for r in caplog.records)
assert "reintroduces" not in caplog.text
@pytest.mark.parametrize(
"index,ramp,steady,expected",
@@ -941,3 +973,488 @@ class TestChunkRampEmission:
p_seg2_0 = self._emit(tm, rid, 4)
assert p_seg2_0 is not None
assert p_seg2_0.meta.left_context_size == 0
class TestAdaptiveController:
def test_compute_next_chunk_size_basic(self):
ctrl = AdaptiveChunkController(
frame_time_ewma_ms=20.0,
first_emit_time=0.0,
last_emit_time=0.0,
frames_emitted_this_segment=10,
)
now = 0.5
size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=100.0)
assert 2 <= size <= 25
def test_compute_next_chunk_size_clamp_min(self):
"""min_frames floor holds when ramp_floor (last_target + delta) is
below min_frames. Uses a small EWMA so delta=2 and last_target=0,
i.e. ramp_floor=2 < min_frames=4 -> greedy=min_frames wins."""
ctrl = AdaptiveChunkController(
frame_time_ewma_ms=10.0,
first_emit_time=0.0,
last_emit_time=0.0,
frames_emitted_this_segment=1,
)
# buffer = 80 - 70 = 10ms (positive but < margin=200)
# greedy: available=0, max_n=0, greedy=min_frames=4
# ramp_floor: last_target=0 + delta=max(2,int(10/15))=2 -> 2 < 4
now = 0.07
size = ctrl.compute_next_chunk_size(now, min_frames=4, max_frames=25, safety_margin_ms=200.0)
assert size == 4 # min_frames floor dominates ramp_floor=2
def test_compute_next_chunk_size_clamp_max(self):
ctrl = AdaptiveChunkController(
frame_time_ewma_ms=10.0,
first_emit_time=0.0,
last_emit_time=0.0,
frames_emitted_this_segment=100,
)
now = 1.0
size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0)
assert size == 25
def test_hysteresis_locks_steady(self):
ctrl = AdaptiveChunkController(
frame_time_ewma_ms=10.0,
first_emit_time=0.0,
last_emit_time=0.0,
frames_emitted_this_segment=50,
)
now = 0.1
size1 = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0)
assert ctrl.steady_state_reached
assert size1 == 25
size2 = ctrl.compute_next_chunk_size(now + 10.0, min_frames=2, max_frames=25, safety_margin_ms=50.0)
assert size2 == 25
def test_ramp_floor_overrides_greedy_when_buffer_low(self):
"""When buffer is negative (overloaded), ramp_floor still forces
gradual growth to avoid stagnating at min_frames (which would cause
too many Code2Wav decode passes and RTF explosion)."""
ctrl = AdaptiveChunkController(
frame_time_ewma_ms=50.0,
first_emit_time=0.0,
last_emit_time=0.0,
frames_emitted_this_segment=2,
last_target=4,
)
# buffer = 160 - 200 = -40ms (negative)
# greedy: available=0, max_n=0, greedy=2
# ramp_floor: last_target=4 + delta=max(2,int(50/15))=3 -> 7
now = 0.2
size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0)
assert not ctrl.steady_state_reached
assert size == 7 # ramp_floor (unconditional) dominates greedy=2
def test_ramp_floor_dominates_when_last_target_set(self):
"""Once last_target is set (post-emit), the unconditional ramp_floor
(last_target + ewma-scaled delta) dominates greedy whenever greedy is
smaller. This is intended: chunk size grows monotonically toward max
to minimize Code2Wav decode passes (pre-W5; see W5 coupling)."""
ctrl = AdaptiveChunkController(
frame_time_ewma_ms=50.0,
first_emit_time=0.0,
last_emit_time=0.0,
frames_emitted_this_segment=4,
last_target=4,
)
# buffer = 320 - 100 = 220ms (positive)
# greedy: available=170, max_n=int(170/50)=3, greedy=3
# ramp_floor: 4 + delta=max(2,int(50/15))=3 -> 7
now = 0.1
size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0)
assert size == 7 # ramp_floor (7) dominates greedy (3)
def test_greedy_dominates_when_ramp_floor_low(self):
"""When last_target=0 (fresh, no prior emit) and buffer is positive
but below the hysteresis threshold, greedy selects an intermediate
size larger than ramp_floor and greedy dominates. This is the only
regime where greedy (not ramp_floor or hysteresis) sets the size."""
ctrl = AdaptiveChunkController(
frame_time_ewma_ms=50.0,
first_emit_time=0.0,
last_emit_time=0.0,
frames_emitted_this_segment=8,
last_target=0,
)
# buffer = 640 - 100 = 540ms (positive, < hysteresis 1300ms)
# greedy: available=490, max_n=int(490/50)=9, greedy=9
# ramp_floor: 0 + delta=max(2,int(50/15))=3 -> 3
now = 0.1
size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0)
assert not ctrl.steady_state_reached
assert size == 9 # greedy (9) dominates ramp_floor (3)
def test_record_emit_ewma_update(self):
ctrl = AdaptiveChunkController(
frame_time_ewma_ms=20.0,
first_emit_time=0.0,
last_emit_time=0.0,
frames_emitted_this_segment=4,
)
ctrl.record_emit(now=0.5, n_frames=8, accumulated=12, target_size=8)
assert ctrl.frames_emitted_this_segment == 12
assert ctrl.accumulated_at_last_emit == 12
assert ctrl.last_emit_time == 0.5
assert ctrl.last_target == 8 # D6: stores target, not actual
measured = (0.5 - 0.0) * 1000.0 / 8
expected = 0.3 * measured + 0.7 * 20.0
assert abs(ctrl.frame_time_ewma_ms - expected) < 0.01
def test_record_emit_first_then_second(self):
ctrl = AdaptiveChunkController(
frame_time_ewma_ms=0.0,
first_emit_time=0.0,
last_emit_time=0.0,
frames_emitted_this_segment=0,
)
ctrl.record_emit(now=0.4, n_frames=4, accumulated=4, target_size=4)
assert ctrl.frame_time_ewma_ms == 0.0
assert ctrl.frames_emitted_this_segment == 4
assert ctrl.last_target == 4
ctrl.record_emit(now=0.8, n_frames=4, accumulated=8, target_size=4)
expected = (0.8 - 0.4) * 1000.0 / 4
assert abs(ctrl.frame_time_ewma_ms - expected) < 0.01
def test_record_emit_collects_underrun(self):
"""record_emit mirrors stream_client.py's playback-cursor underrun
model: per_chunk_gap_ms records each chunk's lateness, total_gap_s
accumulates gaps > 5ms. Deterministic given explicit `now`."""
ctrl = AdaptiveChunkController(
first_emit_time=0.0,
last_emit_time=0.0,
frames_emitted_this_segment=0,
)
# chunk0: now=0.0 (= first_emit_time), n=2 -> 0.16s audio
ctrl.record_emit(now=0.0, n_frames=2, accumulated=2, target_size=2)
assert ctrl.chunk_sizes == [2]
assert ctrl.per_chunk_gap_ms == [0.0]
assert ctrl.total_gap_s == 0.0
assert abs(ctrl.playback_cursor_s - 0.16) < 1e-9
# chunk1: now=0.2 (200ms wall) -> arrival 0.2, late by 0.2-0.16=0.04s
ctrl.record_emit(now=0.2, n_frames=2, accumulated=4, target_size=2)
assert ctrl.chunk_sizes == [2, 2]
assert abs(ctrl.per_chunk_gap_ms[1] - 40.0) < 1e-6
assert abs(ctrl.total_gap_s - 0.04) < 1e-9
assert abs(ctrl.playback_cursor_s - 0.36) < 1e-9
def test_ewma_zero_skips_hysteresis(self):
"""When EWMA=0, hysteresis must not trigger (steady_gen_ms=0
would make any positive buffer satisfy the condition)."""
ctrl = AdaptiveChunkController(
frame_time_ewma_ms=0.0,
first_emit_time=0.0,
last_emit_time=0.0,
frames_emitted_this_segment=10,
)
now = 0.001
size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0)
assert not ctrl.steady_state_reached
assert size == 2 # min_frames (EWMA=0 → max_n=min_frames)
def test_ramp_divisor_configurable(self):
"""ramp_divisor/ramp_delta_min are controller fields wired from config.
Defaults (15/2) preserve the original formula; changing them changes
delta. Lets a deploy tune the ramp via yaml without code changes, while
the unit tests (which use defaults) stay green."""
# default divisor=15: ewma=50 -> delta=max(2,int(50/15))=3 -> floor=4+3=7
c1 = AdaptiveChunkController(
frame_time_ewma_ms=50.0,
first_emit_time=0.0,
last_emit_time=0.0,
frames_emitted_this_segment=2,
last_target=4,
)
assert c1.compute_next_chunk_size(0.2, 2, 25, 50.0) == 7
# divisor=10: delta=max(2,int(50/10))=5 -> floor=4+5=9
c2 = AdaptiveChunkController(
frame_time_ewma_ms=50.0,
first_emit_time=0.0,
last_emit_time=0.0,
frames_emitted_this_segment=2,
last_target=4,
ramp_divisor=10,
)
assert c2.compute_next_chunk_size(0.2, 2, 25, 50.0) == 9
# delta_min=4 + ewma=10 (int(10/15)=0): delta=max(4,0)=4 -> floor=4+4=8
c3 = AdaptiveChunkController(
frame_time_ewma_ms=10.0,
first_emit_time=0.0,
last_emit_time=0.0,
frames_emitted_this_segment=2,
last_target=4,
ramp_delta_min=4,
)
assert c3.compute_next_chunk_size(0.2, 2, 25, 50.0) == 8
class TestAdaptiveEmitHelper:
@pytest.mark.parametrize(
"length,accumulated,target,finished,expected_emit,expected_ctx",
[
(4, 0, 4, False, True, 4),
(3, 0, 4, False, False, 0),
(12, 4, 8, False, True, 8),
(10, 4, 8, False, False, 0),
# D3: finished flushes ALL remaining frames, not just target_size
(6, 4, 8, True, True, 2),
(4, 4, 8, True, True, 0),
(20, 4, 8, True, True, 16), # new_frames=16 > target=8 → flush 16
(30, 0, 25, True, True, 30), # finished, new_frames=30 > target=25 → flush 30
# B1: non-finished path also ships ALL accumulated frames (new_frames),
# not just target_size. When new_frames > target_size (e.g. scheduler
# stall dropped target below queued frames), returning target_size
# would cause the [-context_length:] slice to skip surplus frames.
(10, 0, 4, False, True, 10), # new_frames=10 > target=4 → ship 10
(8, 2, 4, False, True, 6), # new_frames=6 > target=4 → ship 6
(100, 0, 25, False, True, 100), # new_frames=100 > target=25 → ship 100
],
)
def test_compute_adaptive_emit(self, length, accumulated, target, finished, expected_emit, expected_ctx):
emit, ctx = compute_adaptive_emit(length, accumulated, target, finished)
assert emit == expected_emit
assert ctx == expected_ctx
class TestAdaptiveConfigParsing:
def test_disabled_by_default(self):
enabled, min_f, margin, divisor, delta_min = parse_adaptive_config({})
assert enabled is False
assert min_f == 2
assert margin == 50.0
assert divisor == 15
assert delta_min == 2
def test_enabled(self):
cfg = {"codec_chunk_adaptive": True}
enabled, *_ = parse_adaptive_config(cfg)
assert enabled is True
def test_custom_values(self):
cfg = {
"codec_chunk_adaptive": True,
"codec_chunk_min_frames": 4,
"codec_chunk_safety_margin_ms": 300.0,
"codec_chunk_ramp_divisor": 10,
"codec_chunk_ramp_delta_min": 3,
}
enabled, min_f, margin, divisor, delta_min = parse_adaptive_config(cfg)
assert enabled is True
assert min_f == 4
assert margin == 300.0
assert divisor == 10
assert delta_min == 3
def test_min_frames_clamped(self):
cfg = {"codec_chunk_adaptive": True, "codec_chunk_min_frames": 0}
_, min_f, _, _, _ = parse_adaptive_config(cfg)
assert min_f == 1
def test_ramp_divisor_clamped(self):
cfg = {"codec_chunk_adaptive": True, "codec_chunk_ramp_divisor": 0}
_, _, _, divisor, _ = parse_adaptive_config(cfg)
assert divisor == 1
def test_ramp_delta_min_clamped(self):
cfg = {"codec_chunk_adaptive": True, "codec_chunk_ramp_delta_min": 0}
_, _, _, _, delta_min = parse_adaptive_config(cfg)
assert delta_min == 1
class TestAdaptiveEmission:
def _emit(self, tm, rid, n_frames, finished=False):
tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(n_frames)]
return talker2code2wav_async_chunk(
transfer_manager=tm,
multimodal_output={"codes": {"audio": torch.zeros((0,))}},
request=_req(rid, finished=finished),
is_finished=finished,
)
def test_adaptive_chunk0_uses_dynamic_ic(self):
tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8)
rid = "adaptive-ic"
p0 = self._emit(tm, rid, 2)
assert p0 is not None
assert len(p0.codes.audio) == _Q * 2
def test_adaptive_chunk0_holds_until_ic(self):
tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8)
rid = "adaptive-hold"
p = self._emit(tm, rid, 1)
assert p is None
p0 = self._emit(tm, rid, 2)
assert p0 is not None
assert len(p0.codes.audio) == _Q * 2
def test_adaptive_finished_flush(self):
tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8)
rid = "adaptive-flush"
p0 = self._emit(tm, rid, 2)
assert p0 is not None
tm.ramp_chunk_count[rid] = 1
p_fin = self._emit(tm, rid, 5, finished=True)
assert p_fin is not None
assert p_fin.meta.finished.item() is True
def test_adaptive_finished_no_new_frames(self):
tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8)
rid = "adaptive-no-new"
p0 = self._emit(tm, rid, 2)
assert p0 is not None
tm.ramp_chunk_count[rid] = 1
p_fin = self._emit(tm, rid, 2, finished=True)
assert p_fin is not None
assert p_fin.meta.finished.item() is True
assert p_fin.codes.audio.numel() == 0
def test_adaptive_precedence_over_fixed_ramp(self):
tm = _tm(chunk_ramp=[4, 4, 8, 16, 25], adaptive=True, adaptive_min=2, max_num_seqs=8)
rid = "adaptive-precedence"
p = self._emit(tm, rid, 1)
assert p is None
p0 = self._emit(tm, rid, 2)
assert p0 is not None
assert len(p0.codes.audio) == _Q * 2
def test_adaptive_backward_compat(self):
tm = _tm(initial_chunk_frames=1)
rid = "no-adaptive"
p0 = self._emit(tm, rid, 1)
assert p0 is not None
assert len(p0.codes.audio) == _Q * 1
def test_adaptive_segment_reset(self):
tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8)
rid = "adaptive-segment"
p0 = self._emit(tm, rid, 2)
assert p0 is not None
tm.ramp_chunk_count[rid] = 1
tm.put_req_chunk[rid] = 1
tm.code_prompt_token_ids[rid] = []
tm.ramp_chunk_count.pop(rid, None)
tm._adaptive_states.pop(rid, None)
p_seg2 = self._emit(tm, rid, 2)
assert p_seg2 is not None
assert p_seg2.meta.left_context_size == 0
def test_adaptive_with_ref_code_first_chunk(self):
tm = _tm(adaptive=True, adaptive_min=2, left_context=25, max_num_seqs=8)
rid = "adaptive-ref"
tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(2)]
ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
payload = talker2code2wav_async_chunk(
transfer_manager=tm,
multimodal_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
request=_req(rid, finished=False),
is_finished=False,
)
assert payload is not None
assert payload.meta.ref_context_size == 2
assert payload.meta.ref_context_included is True
assert payload.meta.left_context_size == 2
class TestAdaptiveAccumulationPath:
"""Regression test for the frame-skipping bug (PR #6001 review B1).
Appends one frame per call with a controlled clock, asserting that
the concatenation of every emitted chunk equals the accumulated
frame sequence, in order, with no gaps. Each frame has a distinct
value so the assertion can detect a skipped frame rather than just
a short total.
"""
def test_no_frames_skipped_when_target_drops(self, monkeypatch):
"""When greedy target_size is high (large buffer) and then drops
(clock jumps), new_frames > target_size at the emit point.
compute_adaptive_emit now returns new_frames (not target_size),
so all accumulated frames are shipped and the [-context_length:]
slice stays contiguous with the previous emit. last_target is set
to target_size (not n_frames) to prevent the ramp floor from
ratcheting up on transient overshoot.
Clock schedule (max_num_seqs=4 gives IC=2 via dynamic IC):
t=0.0 chunk 0 (IC=2), emit 2. ewma=0.
t=0.05 emit target=4 (ramp_floor). ewma=12.5 (direct assign).
t=0.05 greedy dominates (target=18), 17 frames accumulate (holds).
t=100.0 clock jumps -> greedy collapses, target drops to
ramp_floor=6. new_frames=18 > 6 -> emit 18 (all new frames).
last_target=6 (not 18, preventing floor ratchet).
"""
tm = _tm(
adaptive=True,
adaptive_min=2,
max_num_seqs=4,
)
rid = "adaptive-accum"
tm.code_prompt_token_ids[rid] = []
clock = [0.0]
monkeypatch.setattr(
"vllm_omni.model_executor.stage_input_processors.qwen3_tts.time.monotonic",
lambda: clock[0],
)
emitted_indices: list[int] = []
total_frames = 40
def append_and_emit(frame_idx, finished=False):
tm.code_prompt_token_ids[rid].append([frame_idx] * _Q)
p = talker2code2wav_async_chunk(
transfer_manager=tm,
multimodal_output={"codes": {"audio": torch.zeros((0,))}},
request=_req(rid, finished=finished),
is_finished=finished,
)
if p is not None:
if p.codes.audio.numel() > 0:
n = p.codes.audio.numel() // _Q
emitted_indices.extend(p.codes.audio[:n].tolist())
tm.ramp_chunk_count[rid] += 1
return p
# Phase 1 (t=0.0): chunk 0, IC=2
for i in range(2):
append_and_emit(i)
# Phase 2 (t=0.05): build EWMA, emit target=4 at length=6
clock[0] = 0.05
for i in range(2, 6):
append_and_emit(i)
# Phase 3 (t=0.05): greedy target=18 (large buffer), accumulate
# 17 frames while holding (new_frames 1..17 < 18)
for i in range(6, 23):
append_and_emit(i)
# Phase 4 (t=100.0): clock jumps, greedy collapses, target drops
# to ramp_floor=6. new_frames=18 > 6 -> bug trigger point.
clock[0] = 100.0
for i in range(23, total_frames):
append_and_emit(i, finished=(i == total_frames - 1))
# Assert: all frames 0..39 emitted, in order, no gaps.
expected = list(range(total_frames))
assert emitted_indices == expected, (
f"Frame loss detected! "
f"Missing: {sorted(set(expected) - set(emitted_indices))}, "
f"Extra: {sorted(set(emitted_indices) - set(expected))}"
)
+78
View File
@@ -31,6 +31,7 @@ from vllm_omni.model_extras import (
[
({"model_type": "bagel"}, "bagel"),
({"architectures": ["HunyuanImage3ForConditionalGeneration"]}, "hunyuan_image3"),
({"model_type": "mammothmoda2"}, "mammoth_moda2"),
({"model_type": "qwen2_vl"}, "generic"),
],
)
@@ -50,6 +51,22 @@ def test_bagel_x_to_text_prompt_builder() -> None:
assert stop_token_ids is None
@pytest.mark.core_model
@pytest.mark.cpu
def test_mammoth_x_to_text_prompt_builder() -> None:
prompt, stop_token_ids = build_x_to_text_prompt("mammoth_moda2", "unused", "Hello.", has_image=False)
assert prompt == {
"prompt": (
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
"<|im_start|>user\nHello.<|im_end|>\n"
"<|im_start|>assistant\n"
),
"modalities": ["text"],
"additional_information": {"omni_task": ["chat"]},
}
assert stop_token_ids is None
@pytest.mark.core_model
@pytest.mark.cpu
def test_generic_x_to_text_prompt_builder() -> None:
@@ -555,3 +572,64 @@ def test_declared_extra_args_apply_to_existing_sampling_params() -> None:
"cfg_text_scale": 4.0,
"think": False,
}
@pytest.mark.core_model
@pytest.mark.cpu
def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> None:
for model_class_name in (
"MammothModa2DiTPipeline",
"MammothModa2ForConditionalGeneration",
"Mammothmoda2Model",
):
assert get_extra_body_params(model_class_name) == frozenset(
{
"text_guidance_scale",
"cfg_range",
"num_inference_steps",
}
)
assert get_extra_output_params(model_class_name) == frozenset()
assert should_init_extra_args_for_non_diffusion_stages(model_class_name) is True
wrapper_prompt = build_text_to_image_prompt(
"MammothModa2ForConditionalGeneration",
{"prompt": "a cat", "modalities": ["image"]},
height=256,
width=256,
)
assert wrapper_prompt["additional_information"]["omni_task"] == ["t2i"]
@pytest.mark.core_model
@pytest.mark.cpu
def test_mammothmoda2_text_to_image_prompt_builder() -> None:
# Image dims are converted to the AR grid (width/16 x height/16); the negative
# prompt is ignored (MammothModa2 t2i uses CFG, not an explicit negative path).
assert build_text_to_image_prompt(
"MammothModa2DiTPipeline",
{
"prompt": "a cat",
"modalities": ["image"],
"negative_prompt": "blurry",
},
height=512,
width=768,
) == {
"prompt": (
"<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n"
"<|im_start|>user\na cat<|im_end|>\n"
"<|im_start|>assistant\n"
"<|image start|>48*32<|image token|>"
),
"additional_information": {
"omni_task": ["t2i"],
"ar_width": [48],
"ar_height": [32],
"image_height": [512],
"image_width": [768],
"eol_token_id": [152064],
"visual_token_start_id": [152072],
"visual_token_end_id": [168456],
},
}
@@ -0,0 +1,217 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the NPU diffusion attention backend selector.
Regression coverage for the eager ``mindiesd`` import in
``NPUOmniPlatform.get_diffusion_attn_backend_cls``: the import must run only
for backends that reach mindiesd kernels (FLASH_ATTN, and RAINFUSION_ATTN via
its dense FlashAttention fallback), so that a broken optional mindiesd
install cannot block backends that never use it (e.g. TORCH_SDPA).
The platform module is loaded from source with fake vllm/vllm-ascend
dependencies, so the tests run on CPU without NPU or mindiesd installed.
"""
from __future__ import annotations
import builtins
import importlib.util
import sys
import types
from enum import Enum
from pathlib import Path
from types import SimpleNamespace
import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
_BACKEND_NAMES = (
"FLASH_ATTN",
"FLASH_ATTN_HUB",
"FLASH_ATTN_3_HUB",
"TORCH_SDPA",
"RAINFUSION_ATTN",
)
def _repo_root() -> Path:
marker = Path("vllm_omni") / "platforms" / "npu" / "platform.py"
for parent in Path(__file__).resolve().parents:
if (parent / marker).is_file():
return parent
raise FileNotFoundError(f"could not locate repo root containing {marker}")
def _install_fake_module(monkeypatch: pytest.MonkeyPatch, name: str, **attrs):
module = types.ModuleType(name)
for key, value in attrs.items():
setattr(module, key, value)
monkeypatch.setitem(sys.modules, name, module)
return module
def _make_fake_backend_enum():
members = {
name: SimpleNamespace(get_path=lambda n=name: f"fake.backends.{n.lower()}.Backend") for name in _BACKEND_NAMES
}
class _FakeEnumMeta(type):
def __getitem__(cls, name):
return members[name]
def __getattr__(cls, name):
try:
return members[name]
except KeyError:
raise AttributeError(name) from None
class FakeDiffusionAttentionBackendEnum(metaclass=_FakeEnumMeta):
pass
return FakeDiffusionAttentionBackendEnum
class _FakeOmniPlatformEnum(Enum):
CUDA = "cuda"
NPU = "npu"
class _FakeNPUPlatform:
pass
class _FakeOmniPlatform:
_omni_enum = _FakeOmniPlatformEnum.NPU
@classmethod
def validate_diffusion_attn_backend(cls, selected_backend: str) -> None:
pass
def _load_platform_module(monkeypatch: pytest.MonkeyPatch):
"""Load vllm_omni/platforms/npu/platform.py with fake dependencies."""
_install_fake_module(monkeypatch, "vllm")
_install_fake_module(
monkeypatch,
"vllm.logger",
init_logger=lambda name: SimpleNamespace(
debug=lambda *a, **k: None,
info=lambda *a, **k: None,
warning=lambda *a, **k: None,
error=lambda *a, **k: None,
),
)
_install_fake_module(monkeypatch, "vllm_ascend")
_install_fake_module(monkeypatch, "vllm_ascend.platform", NPUPlatform=_FakeNPUPlatform)
for pkg in (
"vllm_omni",
"vllm_omni.diffusion",
"vllm_omni.diffusion.attention",
"vllm_omni.diffusion.attention.backends",
"vllm_omni.platforms",
):
_install_fake_module(monkeypatch, pkg)
_install_fake_module(
monkeypatch,
"vllm_omni.diffusion.attention.backends.registry",
DiffusionAttentionBackendEnum=_make_fake_backend_enum(),
)
_install_fake_module(
monkeypatch,
"vllm_omni.platforms.interface",
OmniPlatform=_FakeOmniPlatform,
OmniPlatformEnum=_FakeOmniPlatformEnum,
)
path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "platform.py"
spec = importlib.util.spec_from_file_location("npu_platform_under_test", path)
if spec is None or spec.loader is None:
raise RuntimeError(f"cannot load import spec for {path}")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module.NPUOmniPlatform
@pytest.fixture
def mindiesd_tracker(monkeypatch: pytest.MonkeyPatch):
"""Make mindiesd importable and record eager import attempts."""
monkeypatch.delitem(sys.modules, "mindiesd", raising=False)
attempts = []
real_import = builtins.__import__
def fake_import(name, *args, **kwargs):
if name == "mindiesd":
attempts.append(name)
return types.ModuleType("mindiesd")
return real_import(name, *args, **kwargs)
monkeypatch.setattr(builtins, "__import__", fake_import)
monkeypatch.setattr(
"importlib.util.find_spec",
lambda name, *a, **k: object() if name == "mindiesd" else None,
)
return attempts
@pytest.mark.parametrize(
("selected_backend", "expect_mindiesd_import"),
[
("FLASH_ATTN", True),
# HF-hub variants fall back to local FLASH_ATTN on NPU.
("FLASH_ATTN_HUB", True),
("FLASH_ATTN_3_HUB", True),
("RAINFUSION_ATTN", True),
# Backends that never reach mindiesd must not trigger the import.
("TORCH_SDPA", False),
],
)
def test_eager_mindiesd_import_scoped_to_mindiesd_backends(
monkeypatch: pytest.MonkeyPatch,
mindiesd_tracker: list[str],
selected_backend: str,
expect_mindiesd_import: bool,
):
platform_cls = _load_platform_module(monkeypatch)
path = platform_cls.get_diffusion_attn_backend_cls(selected_backend, head_size=64)
assert (len(mindiesd_tracker) > 0) is expect_mindiesd_import
expected = "FLASH_ATTN" if selected_backend in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB") else selected_backend
assert path == f"fake.backends.{expected.lower()}.Backend"
def test_explicit_flash_attn_without_mindiesd_still_resolves(
monkeypatch: pytest.MonkeyPatch,
):
platform_cls = _load_platform_module(monkeypatch)
monkeypatch.delitem(sys.modules, "mindiesd", raising=False)
monkeypatch.setattr("importlib.util.find_spec", lambda name, *a, **k: None)
path = platform_cls.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64)
assert path == "fake.backends.flash_attn.Backend"
def test_default_backend_prefers_flash_attn_when_mindiesd_available(
monkeypatch: pytest.MonkeyPatch,
mindiesd_tracker: list[str],
):
platform_cls = _load_platform_module(monkeypatch)
path = platform_cls.get_diffusion_attn_backend_cls(None, head_size=64)
assert path == "fake.backends.flash_attn.Backend"
assert len(mindiesd_tracker) > 0
def test_default_backend_falls_back_to_sdpa_without_mindiesd(
monkeypatch: pytest.MonkeyPatch,
):
platform_cls = _load_platform_module(monkeypatch)
monkeypatch.delitem(sys.modules, "mindiesd", raising=False)
monkeypatch.setattr("importlib.util.find_spec", lambda name, *a, **k: None)
path = platform_cls.get_diffusion_attn_backend_cls(None, head_size=64)
assert path == "fake.backends.torch_sdpa.Backend"
+159 -1
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import ast
@@ -849,6 +852,161 @@ def test_minicpmo_stage0_data_plane_prefill_matches_official_unit_format():
assert result["prompt_suffix_len"] == 0
def _stage0_vision_runtime():
import torch
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
MiniCPMO45Stage0DuplexRuntime,
)
class _StageModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.embed = torch.nn.Embedding(256, 2)
def get_input_embeddings(self):
return self.embed
def get_audio_hidden_states(self, _data):
return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)]
runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
runtime.stage_model = _StageModel()
runtime.thinker = runtime.stage_model
runtime.tokenizer = SimpleNamespace(
unk_token_id=0,
convert_tokens_to_ids=lambda token: {
"<unit>": 1,
"</unit>": 2,
"<|listen|>": 3,
"<|speak|>": 4,
"<|tts_bos|>": 5,
"<|tts_eos|>": 6,
"<|tts_pad|>": 7,
"<|chunk_eos|>": 8,
"<|chunk_tts_eos|>": 9,
"<|turn_eos|>": 10,
"<|audio|>": 11,
"<image>": 12,
"</image>": 13,
"<slice>": 14,
"</slice>": 15,
}.get(token, 0),
encode=lambda text, add_special_tokens=False: [201, 5],
)
runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4)
runtime.device = "cpu"
runtime._init_token_ids()
# One 2-row block per frame keeps the expected token_ids readable; the real
# resampler emits VISION_EMBEDS_PER_FRAME rows per slice.
runtime._stage_vision_embeddings = lambda frames: [
[torch.full((2, 2), float(index), dtype=torch.float32)] for index in range(len(frames))
]
return runtime
def test_minicpmo_stage0_puts_every_frame_of_an_append_in_one_unit():
"""Official streaming_prefill feeds the whole frame_list into one unit.
Stacked pairs used to be spread one frame per unit, so a two-frame append
that closed a single unit left a frame over, failed the prefill, and still
advanced the session state.
"""
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
_MiniCPMO45Stage0SessionState,
)
runtime = _stage0_vision_runtime()
state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-stacked-frames")
result = runtime._stage_prefill_embeddings_only(
state,
np.zeros(4, dtype=np.float32),
video_frames=["frame-prev", "frame-cur"],
seq=1,
)
assert result["success"] is True
# <unit>, then each frame as <image> + embeds + </image>, then the audio.
assert result["input_token_ids"] == [1, 12, 0, 0, 13, 12, 0, 0, 13, 11]
assert state.audio_chunk_idx == 1
def test_minicpmo_stage0_wraps_hd_slices_after_the_source_image():
"""Official stacked pair is max_slice_nums=[2, 1]: source + slices, then audio."""
import torch
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
_MiniCPMO45Stage0SessionState,
)
runtime = _stage0_vision_runtime()
state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-hd-slices")
runtime._stage_vision_embeddings = lambda frames: [
[torch.ones((2, 2)), torch.full((2, 2), 2.0), torch.full((2, 2), 3.0)],
[torch.full((2, 2), 4.0)],
]
result = runtime._stage_prefill_embeddings_only(
state,
np.zeros(4, dtype=np.float32),
video_frames=["base", "stack"],
seq=1,
)
assert result["success"] is True
# <unit> <image> src </image> <slice> p0 </slice> <slice> p1 </slice>
# <image> stack </image> audio
assert result["input_token_ids"] == [
1,
12,
0,
0,
13,
14,
0,
0,
15,
14,
0,
0,
15,
12,
0,
0,
13,
11,
]
def test_minicpmo_stage0_reports_frames_that_close_no_unit_instead_of_dropping_them():
"""A frame only binds to the unit its append closes.
Frames are not carried across appends, so one arriving early used to vanish
under a generic "audio not enough" buffering result.
"""
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
_MiniCPMO45Stage0SessionState,
)
runtime = _stage0_vision_runtime()
state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-early-frame")
assert runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1)["success"] is True
early = runtime._stage_prefill_embeddings_only(
state,
np.zeros(2, dtype=np.float32),
video_frames=["frame-cur"],
seq=2,
)
assert early["success"] is False
assert "closes no model unit" in early["reason"]
# The half-filled append must not advance the unit counter.
assert state.audio_chunk_idx == 1
def test_minicpmo_stage0_speech_append_sets_pending_context_once():
import torch
@@ -2514,7 +2672,7 @@ def _ar_runner_classdefs():
def _method_calls(method):
calls = {}
calls: dict[str, int] = {}
for node in ast.walk(method):
if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
calls.setdefault(node.func.attr, node.lineno)
@@ -379,6 +379,7 @@ class TestLoadCustomFuncSelection(unittest.TestCase):
def test_skips_non_payload_stage_input_processors_for_full_payload_mode(self):
incompatible_paths = [
"vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav",
"vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit",
"vllm_omni.model_executor.stage_input_processors.cosyvoice3.text2flow",
"vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion",
]
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Daily-Omni Dataset loader for benchmark.
Daily-Omni is an audio-visual reasoning benchmark with 684 videos
@@ -627,10 +630,13 @@ class DailyOmniDataset(BenchmarkDataset):
def _load_from_local_json(self) -> None:
"""Load QA data from local JSON file."""
if not self.qa_json_path.exists():
raise FileNotFoundError(f"QA JSON file not found: {self.qa_json_path}")
qa_json_path = self.qa_json_path
if qa_json_path is None:
raise RuntimeError("qa_json_path is required for local JSON loading")
if not qa_json_path.exists():
raise FileNotFoundError(f"QA JSON file not found: {qa_json_path}")
with open(self.qa_json_path, encoding="utf-8") as f:
with open(qa_json_path, encoding="utf-8") as f:
data = json.load(f)
# Support both list format and dict with "train"/"test" splits
@@ -682,8 +688,11 @@ class DailyOmniDataset(BenchmarkDataset):
one ``sample()`` / request-building implementation for both.
"""
snapshot_error: Exception | None = None
dataset_path = self.dataset_path
if dataset_path is None:
raise ValueError("dataset_path is required for HuggingFace loading")
try:
root = ensure_daily_omni_hub_root(self.dataset_path)
root = ensure_daily_omni_hub_root(dataset_path)
except Exception as e: # noqa: BLE001 - any failure just falls through to `datasets`
snapshot_error = e
else:
@@ -1256,7 +1265,7 @@ class DailyOmniDataset(BenchmarkDataset):
num_video_frames, avg_fps = _probe_video_frames_and_fps(video_path)
duration = num_video_frames / avg_fps
num_seconds = max(1, int(math.ceil(duration)))
second_timestamps = list(range(num_seconds))
second_timestamps = [float(i) for i in range(num_seconds)]
if duration > max_num_frames:
timestamps = [round(i * 0.1, 1) for i in range(int(duration / 0.1))]
+9
View File
@@ -0,0 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU-friendly Omni-DuplexEval generation and scoring helpers."""
from .omni_duplex_eval_dataset import DuplexSample, load_samples
from .omni_duplex_eval_metrics import PROTOCOL_PIN
__all__ = ["DuplexSample", "PROTOCOL_PIN", "load_samples"]
@@ -0,0 +1,106 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Media-clock sentence extraction and response JSON normalization."""
from __future__ import annotations
from collections.abc import Iterable
from dataclasses import dataclass
from typing import Any
import regex as re
_TERMINAL = re.compile(r"(.+?[.!?。!?]+)(?:\s+|$)", re.S)
@dataclass
class TimedSentence:
sentence: str
start: float
end: float
def as_dict(self) -> dict[str, Any]:
return {"sentence": self.sentence, "start": round(self.start, 3), "end": round(self.end, 3)}
def split_text(text: str) -> list[str]:
text = str(text or "").strip()
if not text:
return []
parts: list[str] = []
position = 0
for match in _TERMINAL.finditer(text):
parts.append(match.group(1).strip())
position = match.end()
tail = text[position:].strip()
if tail:
parts.append(tail)
return [part for part in parts if part]
def extract_timed_sentences(events: Iterable[dict[str, Any]], *, clock: str = "media") -> list[TimedSentence]:
"""Split output deltas while anchoring text to input ``audio_end_ms``."""
if clock != "media":
raise ValueError("only clock=media is supported")
sentences: list[TimedSentence] = []
pending = ""
start: float | None = None
last: float | None = None
for event in events:
event_type = event.get("type")
if event_type not in {"response.output_text.delta", "response.audio_transcript.delta", "response.text.delta"}:
continue
delta = event.get("delta")
if not isinstance(delta, str) or not delta.strip():
continue
timestamp = event.get("audio_end_ms")
if timestamp is None:
timestamp = event.get("_media_clock_ms", event.get("current_time_ms", 0.0))
try:
# Every accepted event field is explicitly millisecond-valued.
now = float(timestamp) / 1000.0
except (TypeError, ValueError):
now = last or 0.0
pending += delta
start = now if start is None else start
last = now
pieces = split_text(pending)
complete = pieces if pieces and re.search(r"[.!?。!?]$", pending.strip()) else pieces[:-1]
if complete:
for piece in complete:
sentences.append(TimedSentence(piece, start, max(start, now)))
consumed = " ".join(complete)
pending = pending[len(consumed) :].strip()
start = now if pending else None
if pending.strip():
sentences.append(TimedSentence(pending.strip(), start or last or 0.0, last or start or 0.0))
return sentences
def normalize_response_items(raw: Any) -> list[dict[str, Any]]:
if isinstance(raw, dict):
items = raw.get("sentences", raw.get("chunks"))
else:
items = raw
if not isinstance(items, list):
raise ValueError("response JSON must be a list or contain sentences/chunks")
normalized = []
for item in items:
if not isinstance(item, dict):
continue
text = item.get("sentence", item.get("text", item.get("content", "")))
if not isinstance(text, str) or not text.strip():
continue
try:
start = float(item.get("start", item.get("current_time", item.get("time", 0.0))) or 0.0)
end = float(item.get("end", item.get("current_time", item.get("time", start))) or start)
except (TypeError, ValueError):
continue
normalized.append({"sentence": text.strip(), "start": start, "end": max(start, end)})
return normalized
def validate_clock(meta: dict[str, Any], *, allow_invalid: bool = False) -> None:
if meta.get("clock") == "invalid" and not allow_invalid:
raise ValueError("generate artifact has clock=invalid; pass --allow-invalid-clock to score it")
@@ -0,0 +1,193 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Dataset and split normalization for Omni-DuplexEval.
The loader accepts a Hugging Face dataset id, a JSON/JSONL manifest, or an
already materialized iterable. Media is deliberately resolved at use time so
the benchmark remains usable in air-gapped environments.
"""
from __future__ import annotations
import json
from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path
from typing import Any
DEFAULT_DATASET = "Hothan/Omni-DuplexEval"
RTD_SPLITS = frozenset(
{
"RTD_world_knowledge",
"RTD_counting",
"RTD_fine_grained_movement",
"RTD_interaction_relation",
"RTD_OCR",
"RTD_Omni",
}
)
PR_SPLITS = frozenset({"PR_correction", "PR_event_reminder", "PR_post_event_reminder"})
_TASK_ALIASES = {
"correction": "correction",
"pr_correction": "correction",
"event_reminder": "proactive_reminder",
"pr_event_reminder": "proactive_reminder",
"proactive_reminder": "proactive_reminder",
"post_event_reminder": "post_event_reminder",
"pr_post_event_reminder": "post_event_reminder",
}
def canonical_task_type(value: str) -> str:
key = str(value).strip().lower().replace("-", "_").replace(" ", "_")
if key not in _TASK_ALIASES:
raise ValueError(f"unsupported Omni-DuplexEval task type: {value!r}")
return _TASK_ALIASES[key]
def task_type_for_split(split: str) -> str:
"""Match the task routing used by the official HF batch evaluator."""
key = str(split).strip().lower()
if "correction" in key:
return "correction"
if "post_event" in key:
return "post_event_reminder"
if key.startswith("pr"):
return "proactive_reminder"
raise ValueError(f"cannot infer proactive-reminder task from split {split!r}")
def family_for_split(split: str, task_type: str | None = None) -> str:
if split in RTD_SPLITS or str(split).upper().startswith("RTD"):
return "rtd"
if split in PR_SPLITS or str(split).upper().startswith("PR"):
return "pr"
if task_type:
return "rtd" if str(task_type).lower().startswith("rtd") else "pr"
raise ValueError(f"cannot infer benchmark family from split {split!r}")
def _value(row: dict[str, Any], *keys: str, default: Any = None) -> Any:
for key in keys:
if key in row and row[key] is not None:
return row[key]
return default
def _float(value: Any, default: float | None = None) -> float | None:
try:
return None if value is None else float(value)
except (TypeError, ValueError):
return default
@dataclass(frozen=True)
class DuplexSample:
id: str
split: str
family: str
task_type: str | None
video: Any
question_audio: Any = None
question_text: str = ""
answer1: str = ""
answer2: str = ""
reminder1: Any = None
reminder2: Any = None
video_duration: float | None = None
video_type: str = ""
raw: dict[str, Any] | None = None
@classmethod
def from_row(cls, row: dict[str, Any], *, split: str | None = None, media_root: Path | None = None) -> DuplexSample:
chosen_split = str(split or _value(row, "split", "subset", "config", default=""))
task_value = _value(row, "task_type", "task", "type")
family = str(_value(row, "family", default="") or family_for_split(chosen_split, task_value))
task = None
if family == "pr":
task = canonical_task_type(str(task_value)) if task_value else task_type_for_split(chosen_split)
sample_id = str(_value(row, "id", "sample_id", "uid", "name", default=""))
video = _value(row, "video", "video_path", "video_file")
audio = _value(row, "question_audio", "audio", "question_wav")
if media_root:
def resolve(value: Any) -> Any:
if not isinstance(value, str):
return value
candidate = Path(value).expanduser()
return str(candidate if candidate.is_absolute() else media_root / candidate)
video, audio = resolve(video), resolve(audio)
return cls(
id=sample_id,
split=chosen_split,
family=family,
task_type=task,
video=video,
question_audio=audio,
question_text=str(_value(row, "question_text", "question", "instruction", default="") or ""),
answer1=str(_value(row, "answer1", "answer", "ground_answer", default="") or ""),
answer2=str(_value(row, "answer2", default="") or ""),
reminder1=_value(row, "reminder1", "reminder_1"),
reminder2=_value(row, "reminder2", "reminder_2"),
video_duration=_float(_value(row, "video_duration", "duration")),
video_type=str(_value(row, "video_type", default="") or ""),
raw=dict(row),
)
def _read_manifest(path: Path) -> list[dict[str, Any]]:
if path.suffix.lower() == ".jsonl":
return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
payload = json.loads(path.read_text(encoding="utf-8"))
if isinstance(payload, dict):
payload = payload.get("data", payload.get("samples", [payload]))
if not isinstance(payload, list):
raise ValueError("manifest must contain a JSON list")
return payload
def load_samples(
dataset: str | Path | Iterable[dict[str, Any]] = DEFAULT_DATASET,
*,
split: str | None = None,
family: str = "all",
media_root: str | Path | None = None,
limit: int | None = None,
ids: Iterable[str] | None = None,
) -> list[DuplexSample]:
if isinstance(dataset, (str, Path)) and Path(str(dataset)).exists():
rows = _read_manifest(Path(str(dataset)))
elif not isinstance(dataset, (str, Path)):
rows = list(dataset)
else:
try:
from datasets import load_dataset
except ImportError as exc:
raise RuntimeError("datasets is required for Hugging Face loading; use a local manifest instead") from exc
kwargs: dict[str, Any] = {}
if split and split != "all":
kwargs["split"] = split
loaded = load_dataset(str(dataset), **kwargs)
if isinstance(loaded, dict):
rows = []
for name, table in loaded.items():
rows.extend({**dict(row), "split": name} for row in table)
else:
rows = [dict(row) for row in loaded]
wanted = set(str(item) for item in ids) if ids else None
root = Path(media_root).expanduser() if media_root else None
result = []
for row in rows:
sample = DuplexSample.from_row(row, split=None if split in (None, "all") else split, media_root=root)
if split and split != "all" and sample.split != split:
continue
if family != "all" and sample.family != family:
continue
if wanted is not None and sample.id not in wanted:
continue
result.append(sample)
if limit is not None:
result = result[: max(0, limit)]
return result
Loaded 100 of 211 files, more files were not shown because too many files have changed in this diff. Show more