Compare commits
25
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
60165f24be | ||
|
|
6d84409a68 | ||
|
|
ebb3b20576 | ||
|
|
be335a86f9 | ||
|
|
d5aced9296 | ||
|
|
4ad455e61e | ||
|
|
da94c3ebe1 | ||
|
|
13173d6bbb | ||
|
|
009b992c5b | ||
|
|
04b97a5922 | ||
|
|
15c6ab8fa7 | ||
|
|
c201fd77c6 | ||
|
|
b6a4baa659 | ||
|
|
ed837e350e | ||
|
|
ab8d45d3cc | ||
|
|
f8f8aa8eb7 | ||
|
|
6ec5ab11bd | ||
|
|
4245fd6955 | ||
|
|
5e4daaf197 | ||
|
|
262ff14693 | ||
|
|
069e6c5fbc | ||
|
|
54c1fc64b7 | ||
|
|
c6df39e9f7 | ||
|
|
94b1546ce3 | ||
|
|
19da238246 |
No files matched your search
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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][]
|
||||
|
||||
|
||||
@@ -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. |
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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 |
|
||||
|
||||
@@ -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.
|
||||
@@ -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)
|
||||
|
||||
@@ -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)))
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -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:
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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():
|
||||
|
||||
@@ -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))}"
|
||||
)
|
||||
@@ -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],
|
||||
},
|
||||
}
|
||||
@@ -242,19 +242,18 @@ def tiny_flux_kontext_builder() -> str:
|
||||
|
||||
def tiny_flux2_builder() -> str:
|
||||
def shrink_text_encoder(config: dict) -> dict:
|
||||
# The real checkpoint ships language_model.lm_head.weight as its own
|
||||
# tensor even though tie_word_embeddings defaults to True; keep it
|
||||
# untied here too so transformers doesn't drop it on save.
|
||||
config["tie_word_embeddings"] = False
|
||||
config["text_config"]["num_hidden_layers"] = 31
|
||||
config["text_config"]["num_hidden_layers"] = 3
|
||||
config["text_config"]["hidden_size"] = 32
|
||||
config["text_config"]["intermediate_size"] = 64
|
||||
config["text_config"]["num_attention_heads"] = 4
|
||||
config["text_config"]["head_dim"] = 32
|
||||
config["text_config"]["num_attention_heads"] = 2
|
||||
config["text_config"]["head_dim"] = 16
|
||||
config["text_config"]["num_key_value_heads"] = 2
|
||||
config["text_config"]["text_encoder_out_layers"] = [0, 1, 2]
|
||||
config["vision_config"]["num_hidden_layers"] = 2
|
||||
config["vision_config"]["intermediate_size"] = 64
|
||||
config["vision_config"]["num_attention_heads"] = 4
|
||||
config["vision_config"]["head_dim"] = 32
|
||||
config["vision_config"]["num_attention_heads"] = 2
|
||||
config["vision_config"]["head_dim"] = 16
|
||||
return config
|
||||
|
||||
return build_tiny_from_configs(
|
||||
@@ -262,6 +261,6 @@ def tiny_flux2_builder() -> str:
|
||||
"black-forest-labs/FLUX.2-dev",
|
||||
transform={
|
||||
"text_encoder": shrink_text_encoder,
|
||||
"transformer": partial(_shrink_dit_rope_config, num_single_layers=2),
|
||||
"transformer": partial(_shrink_dit_rope_config, num_single_layers=2, joint_attention_dim=96),
|
||||
},
|
||||
)
|
||||
@@ -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"
|
||||
@@ -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))]
|
||||
|
||||
@@ -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
|
||||
@@ -0,0 +1,245 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Evaluate and summarize generated Omni-DuplexEval artifacts."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .omni_duplex_eval_clock import normalize_response_items, validate_clock
|
||||
from .omni_duplex_eval_dataset import DuplexSample
|
||||
from .omni_duplex_eval_judge import DuplexJudge
|
||||
from .omni_duplex_eval_media import extract_jpeg, materialize_media, video_duration
|
||||
from .omni_duplex_eval_metrics import (
|
||||
PROTOCOL_PIN,
|
||||
build_content_prompt,
|
||||
build_reminder_prompt,
|
||||
build_temporal_prompt,
|
||||
parse_judge_json,
|
||||
reminder_window,
|
||||
summarize_pr_results,
|
||||
summarize_temporal_results,
|
||||
temporal_window,
|
||||
)
|
||||
|
||||
|
||||
def _read(path: Path) -> Any:
|
||||
return json.loads(path.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def _write(path: Path, value: Any) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
|
||||
|
||||
def _text(items: list[dict[str, Any]]) -> str:
|
||||
return " ".join(str(item.get("sentence", item.get("text", ""))).strip() for item in items).strip()
|
||||
|
||||
|
||||
def evaluate_sample(
|
||||
sample: DuplexSample,
|
||||
response_path: Path,
|
||||
score_path: Path,
|
||||
judge: DuplexJudge,
|
||||
*,
|
||||
judge_fps: int = 2,
|
||||
judge_video_mode: str = "video_url",
|
||||
window_size: float = 10.0,
|
||||
allow_invalid_clock: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
raw = _read(response_path)
|
||||
meta = (
|
||||
_read(response_path.with_name(response_path.stem + ".meta.json"))
|
||||
if response_path.with_name(response_path.stem + ".meta.json").exists()
|
||||
else {}
|
||||
)
|
||||
validate_clock(meta, allow_invalid=allow_invalid_clock)
|
||||
items = normalize_response_items(raw)
|
||||
result: dict[str, Any] = {
|
||||
"id": sample.id,
|
||||
"family": sample.family,
|
||||
"task_type": sample.task_type,
|
||||
"protocol_pin": PROTOCOL_PIN,
|
||||
"response_meta": meta,
|
||||
"judge_model": getattr(judge, "model", None),
|
||||
"judge_video_mode": judge_video_mode,
|
||||
}
|
||||
if sample.family == "rtd":
|
||||
video_path = materialize_media(sample.video, score_path.parent / ".media", sample.id, ".mp4")
|
||||
duration = sample.video_duration or video_duration(video_path)
|
||||
temporal_rows = []
|
||||
for item in items:
|
||||
window = temporal_window(item["start"], item["end"], duration)
|
||||
if not window:
|
||||
temporal_rows.append(
|
||||
{
|
||||
**item,
|
||||
"sentence_start": item["start"],
|
||||
"sentence_end": item["end"],
|
||||
"sentence_duration": max(0.0, item["end"] - item["start"]),
|
||||
"error": "No valid temporal window",
|
||||
"temporal_score": 0,
|
||||
"is_relevant": 0,
|
||||
}
|
||||
)
|
||||
continue
|
||||
frames = _extract_frames(video_path, _times(*window, fps=judge_fps))
|
||||
parsed = parse_judge_json(
|
||||
judge.temporal(build_temporal_prompt(*window, item["sentence"], sample.question_text), frames)
|
||||
)
|
||||
temporal_rows.append(
|
||||
{
|
||||
**item,
|
||||
"sentence_start": item["start"],
|
||||
"sentence_end": item["end"],
|
||||
"sentence_duration": max(0.0, item["end"] - item["start"]),
|
||||
"window_start": window[0],
|
||||
"window_end": window[1],
|
||||
"error": None,
|
||||
**parsed,
|
||||
}
|
||||
)
|
||||
content_frames = None
|
||||
if judge_video_mode == "frame-sample":
|
||||
content_frames = _extract_frames(video_path, _content_frame_times(duration))
|
||||
content = parse_judge_json(
|
||||
judge.content(
|
||||
build_content_prompt(_text(items), sample.question_text, [sample.answer1, sample.answer2]),
|
||||
video_path,
|
||||
frames=content_frames,
|
||||
mode=judge_video_mode,
|
||||
)
|
||||
)
|
||||
if content_frames is not None:
|
||||
content["frame_count"] = len(content_frames)
|
||||
result.update(
|
||||
{
|
||||
"temporal": {"sentences": temporal_rows, "summary": summarize_temporal_results(temporal_rows)},
|
||||
"content": content,
|
||||
}
|
||||
)
|
||||
else:
|
||||
if sample.task_type is None:
|
||||
raise ValueError(f"missing proactive-reminder task type for split {sample.split!r}")
|
||||
task = sample.task_type
|
||||
times = _reminder_times(sample)
|
||||
event_rows = []
|
||||
if task == "correction" or not times:
|
||||
response = _text(items)
|
||||
parsed = parse_judge_json(
|
||||
judge.chat(build_reminder_prompt(sample.question_text, response, task, sample.answer1))
|
||||
)
|
||||
event_rows.append({"response_segment": response, **parsed})
|
||||
else:
|
||||
for start in times:
|
||||
lo, hi = reminder_window(start, window_size)
|
||||
response = " ".join(item["sentence"] for item in items if lo <= item["start"] <= hi)
|
||||
if not response:
|
||||
event_rows.append({"start_time": start, "response_segment": "", "success_score": 0})
|
||||
else:
|
||||
parsed = parse_judge_json(
|
||||
judge.chat(build_reminder_prompt(sample.question_text, response, task, sample.answer1))
|
||||
)
|
||||
event_rows.append({"start_time": start, "response_segment": response, **parsed})
|
||||
scores = [int(row.get("success_score", 0)) for row in event_rows]
|
||||
result["pr"] = {
|
||||
"events": event_rows,
|
||||
"all_success": bool(scores) and all(score == 1 for score in scores),
|
||||
"total_score": int(bool(scores) and all(score == 1 for score in scores)),
|
||||
}
|
||||
_write(score_path, result)
|
||||
return result
|
||||
|
||||
|
||||
def _times(start: float, end: float, *, fps: int) -> list[float]:
|
||||
count = max(1, int((end - start) * fps))
|
||||
return (
|
||||
[start + (end - start) / 2] if count == 1 else [start + i * (end - start) / (count - 1) for i in range(count)]
|
||||
)
|
||||
|
||||
|
||||
def _extract_frames(path: str | Path, timestamps: list[float]) -> list[bytes]:
|
||||
"""Extract only decodable JPEG frames; clips can end on non-keyframes."""
|
||||
frames = []
|
||||
for timestamp in timestamps:
|
||||
try:
|
||||
frame = extract_jpeg(path, timestamp=timestamp)
|
||||
except (OSError, subprocess.CalledProcessError, ValueError):
|
||||
continue
|
||||
if frame.startswith(b"\xff\xd8"):
|
||||
frames.append(frame)
|
||||
return frames
|
||||
|
||||
|
||||
def _content_frame_times(duration: float) -> list[float]:
|
||||
"""Sample the full video at approximately one frame every three seconds."""
|
||||
if duration <= 0:
|
||||
return []
|
||||
times = []
|
||||
timestamp = 0.0
|
||||
while timestamp < duration:
|
||||
times.append(min(timestamp, max(0.0, duration - 0.01)))
|
||||
timestamp += 3.0
|
||||
# Qwen2.5-VL has a finite per-request image-token budget. Preserve
|
||||
# coverage while bounding long clips to a safely portable frame count.
|
||||
if len(times) > 16:
|
||||
stride = (len(times) - 1) / 15
|
||||
times = [times[round(i * stride)] for i in range(16)]
|
||||
return times
|
||||
|
||||
|
||||
def _reminder_times(sample: DuplexSample) -> list[float]:
|
||||
values = []
|
||||
for value in (sample.reminder1, sample.reminder2):
|
||||
if isinstance(value, dict):
|
||||
value = value.get("start", value.get("time"))
|
||||
try:
|
||||
if value is not None:
|
||||
values.append(float(value))
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
return values
|
||||
|
||||
|
||||
def summarize_scores(score_root: str | Path) -> dict[str, Any]:
|
||||
rows = []
|
||||
for path in sorted(Path(score_root).rglob("*.json")):
|
||||
if not path.name.endswith("_summary.json"):
|
||||
row = _read(path)
|
||||
row.setdefault("split", path.parent.name)
|
||||
rows.append(row)
|
||||
rtd_temporal = [
|
||||
row["temporal"]["summary"] | {"content_score": row.get("content", {}).get("content_score", 0.0)}
|
||||
for row in rows
|
||||
if "temporal" in row
|
||||
]
|
||||
pr = [{"task_type": row.get("task_type"), **row.get("pr", {})} for row in rows if "pr" in row]
|
||||
result = {"protocol_pin": PROTOCOL_PIN, "samples": len(rows)}
|
||||
if rtd_temporal:
|
||||
by_task: dict[str, list[dict[str, float]]] = {}
|
||||
for row in rows:
|
||||
if "temporal" in row:
|
||||
by_task.setdefault(str(row.get("split", row.get("task_type", "unknown"))), []).append(
|
||||
{
|
||||
"content": float(row.get("content", {}).get("content_score", 0.0)),
|
||||
"temporal": float(row["temporal"]["summary"].get("avg_temporal_score", 0.0)),
|
||||
}
|
||||
)
|
||||
result["rtd"] = {
|
||||
"mean_content_score": sum(item["content_score"] for item in rtd_temporal) / len(rtd_temporal),
|
||||
"mean_avg_temporal_score": sum(item["avg_temporal_score"] for item in rtd_temporal) / len(rtd_temporal),
|
||||
"by_task": {
|
||||
task: {
|
||||
"mean_content_score": sum(item["content"] for item in values) / len(values),
|
||||
"mean_avg_temporal_score": sum(item["temporal"] for item in values) / len(values),
|
||||
}
|
||||
for task, values in by_task.items()
|
||||
},
|
||||
}
|
||||
if pr:
|
||||
result["pr"] = summarize_pr_results(pr)
|
||||
return result
|
||||
Loaded 100 of 210 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user