Compare commits
41
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6ffdf1a8f1 | ||
|
|
55a226dcf1 | ||
|
|
f66f05633c | ||
|
|
457c6e2793 | ||
|
|
bfe8967ca2 | ||
|
|
fbc93a0e08 | ||
|
|
d934a5f7b8 | ||
|
|
bbb3ae01a1 | ||
|
|
46cfa8435b | ||
|
|
0dc8f45906 | ||
|
|
d94061201d | ||
|
|
9558fd2ea8 | ||
|
|
55a8bd6b80 | ||
|
|
f584729f24 | ||
|
|
034af1f657 | ||
|
|
8bfad23f98 | ||
|
|
6b9ce05feb | ||
|
|
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
|
||||
@@ -10,7 +10,7 @@ steps:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
# Buildkite will automatically replace this with the actual commit hash
|
||||
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
|
||||
VLLM_VERSION: "v0.27.0"
|
||||
VLLM_VERSION: "v0.28.0"
|
||||
priority: 100
|
||||
timeout_in_minutes: 60
|
||||
soft_fail: false
|
||||
@@ -5,7 +5,12 @@ set -ex
|
||||
|
||||
omni_source_dir=$(git rev-parse --show-toplevel)
|
||||
|
||||
base_image_name="xpu/vllm-omni-ci-base:${VLLM_VERSION:?VLLM_VERSION must be set}"
|
||||
: "${VLLM_VERSION:?VLLM_VERSION must be set}"
|
||||
# Official upstream XPU image. It is published from vllm's own
|
||||
# docker/Dockerfile.xpu, which is exactly what build_vllm_base() clones and
|
||||
# builds below, so pulling it is equivalent to building the fallback base.
|
||||
upstream_base_image="vllm/vllm-openai-xpu:${VLLM_VERSION}"
|
||||
local_base_image="xpu/vllm-omni-ci-base:${VLLM_VERSION}"
|
||||
image_name="xpu/vllm-omni-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set}"
|
||||
container_name="xpu_${BUILDKITE_COMMIT}_$(
|
||||
tr -dc A-Za-z0-9 </dev/urandom | head -c 10
|
||||
@@ -49,9 +54,36 @@ build_vllm_base() (
|
||||
"${vllm_source_dir}"
|
||||
)
|
||||
|
||||
if [ -z "$(docker images -q "${base_image_name}")" ]; then
|
||||
build_vllm_base
|
||||
fi
|
||||
# Resolve the vLLM base image, in priority order:
|
||||
# 1. the published upstream image for VLLM_VERSION, pulled fresh (so moving
|
||||
# tags like nightly are refreshed rather than served stale from disk);
|
||||
# 2. a local copy of that same image, when the registry is unreachable;
|
||||
# 3. VLLM_BASE, if that image is already on disk;
|
||||
# 4. otherwise build VLLM_BASE from the matching vLLM source tag (~30min).
|
||||
# VLLM_BASE names the fallback base and defaults to ${local_base_image}; set it
|
||||
# to reuse or produce a base under a different tag.
|
||||
resolve_vllm_base() {
|
||||
if docker pull "${upstream_base_image}"; then
|
||||
base_image_name="${upstream_base_image}"
|
||||
return
|
||||
fi
|
||||
|
||||
# Registry unreachable, but a previous run may have left the image behind.
|
||||
if [ -n "$(docker images -q "${upstream_base_image}")" ]; then
|
||||
echo "WARNING: could not pull ${upstream_base_image}; using local copy" >&2
|
||||
base_image_name="${upstream_base_image}"
|
||||
return
|
||||
fi
|
||||
|
||||
echo "WARNING: ${upstream_base_image} unavailable remotely and locally" >&2
|
||||
base_image_name="${VLLM_BASE:-${local_base_image}}"
|
||||
if [ -z "$(docker images -q "${base_image_name}")" ]; then
|
||||
echo "WARNING: building ${base_image_name} from vLLM ${VLLM_VERSION} source" >&2
|
||||
build_vllm_base
|
||||
fi
|
||||
}
|
||||
|
||||
resolve_vllm_base
|
||||
|
||||
# Try building the docker image
|
||||
docker_build "${image_name}" --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}"
|
||||
|
||||
+27
-134
@@ -1,140 +1,33 @@
|
||||
# Argument to configure vllm base image if pre-built
|
||||
ARG VLLM_BASE=vllm-base
|
||||
|
||||
FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base
|
||||
|
||||
WORKDIR /workspace/
|
||||
|
||||
ARG PYTHON_VERSION=3.12
|
||||
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
|
||||
|
||||
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
|
||||
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
|
||||
|
||||
RUN apt clean && apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends --fix-missing \
|
||||
curl \
|
||||
ffmpeg \
|
||||
git \
|
||||
libsndfile1 \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libgl1 \
|
||||
lsb-release \
|
||||
libaio-dev \
|
||||
numactl \
|
||||
wget \
|
||||
vim \
|
||||
python3.12 \
|
||||
python3.12-dev \
|
||||
python3-pip
|
||||
|
||||
# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer.
|
||||
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
|
||||
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
|
||||
printf '%s\n' \
|
||||
'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \
|
||||
'Pin: version 2025.3*' \
|
||||
'Pin-Priority: 1001' \
|
||||
> /etc/apt/preferences.d/oneapi-2025.3.pref && \
|
||||
apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
intel-oneapi-compiler-dpcpp-cpp-2025.3 \
|
||||
intel-oneapi-mkl-devel-2025.3 \
|
||||
intel-oneapi-dnnl-devel-2025.3 && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install UMD with fixed version
|
||||
RUN mkdir neo && \
|
||||
cd neo && \
|
||||
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
|
||||
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
|
||||
dpkg -i *.deb && \
|
||||
cd .. && \
|
||||
rm -rf neo
|
||||
|
||||
ENV PATH="/root/.local/bin:$PATH"
|
||||
ENV VIRTUAL_ENV="/opt/venv"
|
||||
ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
|
||||
RUN curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
|
||||
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
|
||||
|
||||
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
|
||||
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
|
||||
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
|
||||
bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
|
||||
rm "${ONECCL_INSTALLER}" && \
|
||||
echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
|
||||
echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc
|
||||
RUN rm -f /opt/intel/oneapi/ccl/latest && \
|
||||
ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest
|
||||
|
||||
SHELL ["bash", "-c"]
|
||||
CMD ["bash", "-c", "source /root/.bashrc && exec bash"]
|
||||
|
||||
WORKDIR /workspace/
|
||||
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
# Configure package index for XPU
|
||||
ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
ENV UV_LINK_MODE="copy"
|
||||
|
||||
ARG VLLM_VERSION=v0.27.0
|
||||
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --upgrade pip && \
|
||||
uv pip install -r requirements/xpu.txt && \
|
||||
uv pip install setuptools_scm &&\
|
||||
uv pip install grpcio-tools protobuf nanobind && \
|
||||
source /opt/intel/oneapi/setvars.sh --force && \
|
||||
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
|
||||
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
|
||||
|
||||
ENV LD_LIBRARY_PATH="/opt/intel/oneapi/ccl/2021.15/lib:$LD_LIBRARY_PATH:/usr/local/lib/"
|
||||
|
||||
ENV VLLM_TARGET_DEVICE=xpu
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip uninstall triton triton-xpu && \
|
||||
uv pip install triton-xpu==3.7.2 && \
|
||||
uv pip install --no-build-isolation --no-deps .
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
FROM vllm-base AS vllm-openai
|
||||
|
||||
# install nixl from source code
|
||||
RUN uv pip install nixl
|
||||
|
||||
# install additional dependencies for openai api server
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install accelerate hf_transfer pytest pytest_asyncio lm_eval[api] modelscope
|
||||
|
||||
# install development dependencies (for testing)
|
||||
RUN uv pip install -e tests/vllm_test_utils
|
||||
|
||||
# ensure vllm is properly installed
|
||||
RUN python -c "import vllm, inspect; print(vllm.__file__)"
|
||||
RUN uv pip show vllm
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
ENTRYPOINT []
|
||||
# Target vLLM version: v0.28.0. Keep VLLM_VERSION aligned with the target vLLM
|
||||
# release on every rebase.
|
||||
#
|
||||
# This file layers vLLM-Omni on top of the published upstream XPU image, which
|
||||
# already ships a self-consistent XPU runtime (torch-xpu, oneCCL/UCX+NIXL, the
|
||||
# pinned UMD/compute-runtime stack, and vLLM itself) built from
|
||||
# https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu. Docker
|
||||
# pulls that base automatically, so there is nothing to build here beyond the
|
||||
# Omni layers. Mirrors how docker/Dockerfile.ci consumes vllm/vllm-openai.
|
||||
#
|
||||
# If the tag does not exist yet (targeting an unreleased vLLM commit), build the
|
||||
# base from upstream's own Dockerfile and point VLLM_BASE at the result:
|
||||
#
|
||||
# git clone --depth 1 --branch $VLLM_VERSION https://github.com/vllm-project/vllm /tmp/vllm
|
||||
# docker build -t vllm-openai-xpu:local --target vllm-openai \
|
||||
# -f /tmp/vllm/docker/Dockerfile.xpu /tmp/vllm
|
||||
# docker build -f docker/Dockerfile.xpu --build-arg VLLM_BASE=vllm-openai-xpu:local .
|
||||
#
|
||||
# .buildkite/intel/scripts/run-xpu-test.sh automates that pull-then-build fallback.
|
||||
ARG VLLM_BASE_IMAGE=vllm/vllm-openai-xpu
|
||||
ARG VLLM_VERSION=v0.28.0
|
||||
ARG VLLM_BASE=${VLLM_BASE_IMAGE}:${VLLM_VERSION}
|
||||
|
||||
FROM ${VLLM_BASE} AS vllm-omni
|
||||
|
||||
# The base image ships ENTRYPOINT ["vllm", "serve"]; clear it so this stage stays
|
||||
# usable as a plain shell/test container regardless of which VLLM_BASE was used.
|
||||
ENTRYPOINT []
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
RUN apt-get update && \
|
||||
apt-get install -y espeak-ng ffmpeg git jq && \
|
||||
apt-get clean && \
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -23,6 +23,19 @@ vLLM-Omni currently recommends using the Docker image setup steps below.
|
||||
DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g .
|
||||
```
|
||||
|
||||
This layers vLLM-Omni on top of the published `vllm/vllm-openai-xpu:<VLLM_VERSION>`
|
||||
base image, which Docker pulls automatically. To target a different vLLM release,
|
||||
pass `--build-arg VLLM_VERSION=<tag>`. If that tag has not been published yet, build
|
||||
the base from upstream's own Dockerfile first and point the build at it:
|
||||
|
||||
```bash
|
||||
git clone --depth 1 --branch "$VLLM_VERSION" https://github.com/vllm-project/vllm /tmp/vllm
|
||||
DOCKER_BUILDKIT=1 docker build -t vllm-openai-xpu:local --target vllm-openai \
|
||||
-f /tmp/vllm/docker/Dockerfile.xpu /tmp/vllm
|
||||
DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g \
|
||||
--build-arg VLLM_BASE=vllm-openai-xpu:local .
|
||||
```
|
||||
|
||||
#### Launch the docker image
|
||||
|
||||
##### Launch with OpenAI API Server
|
||||
|
||||
@@ -26,15 +26,15 @@ th {
|
||||
| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | — |
|
||||
| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | | — |
|
||||
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) |
|
||||
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) |
|
||||
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) |
|
||||
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) |
|
||||
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — |
|
||||
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — |
|
||||
| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | ✅︎ | — |
|
||||
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
|
||||
@@ -42,7 +42,7 @@ th {
|
||||
| `SanaWmPipeline` | SANA-WM | `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-WM.md) |
|
||||
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — |
|
||||
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | ✅︎ | — |
|
||||
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
|
||||
| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
|
||||
| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — |
|
||||
@@ -62,13 +62,15 @@ th {
|
||||
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `LongCatVideoAvatarPipeline` | LongCat-Video-Avatar-1.5 A2V/AI2V (native single-speaker and multi-speaker AI2V/AVC) | `meituan-longcat/LongCat-Video-Avatar-1.5` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md) |
|
||||
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | | — |
|
||||
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | ✅︎ | — |
|
||||
| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
|
||||
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
|
||||
| `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` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
@@ -76,10 +78,10 @@ th {
|
||||
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | | — |
|
||||
| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | ✅︎ | — |
|
||||
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | — |
|
||||
| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | | — |
|
||||
| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | — |
|
||||
@@ -99,10 +101,10 @@ th {
|
||||
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
|
||||
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — |
|
||||
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — |
|
||||
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | — |
|
||||
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | ✅︎ | — |
|
||||
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
|
||||
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | | — |
|
||||
| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | ✅︎ | — |
|
||||
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
|
||||
| `HiDreamO1ImagePipeline` | HiDream-O1-Image | `HiDream-ai/HiDream-O1-Image` | ✅︎ | | | | — |
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
"""
|
||||
Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2/LTX-2.3,
|
||||
@@ -111,6 +111,13 @@ def build_image_to_video_prompt(
|
||||
return result
|
||||
|
||||
|
||||
def _validate_video_output_type(output_type: str) -> None:
|
||||
if output_type not in {"image", "video"}:
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{output_type}', expected 'video' or legacy 'image' for video generation."
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=(
|
||||
@@ -665,10 +672,7 @@ def main():
|
||||
frames = frames[0] if frames else None
|
||||
|
||||
if isinstance(frames, OmniRequestOutput):
|
||||
if frames.final_output_type != "image":
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
|
||||
)
|
||||
_validate_video_output_type(frames.final_output_type)
|
||||
if frames.multimodal_output and "audio" in frames.multimodal_output:
|
||||
audio = frames.multimodal_output["audio"]
|
||||
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
|
||||
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
import argparse
|
||||
import json
|
||||
@@ -174,6 +174,13 @@ def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict
|
||||
return result
|
||||
|
||||
|
||||
def _validate_video_output_type(output_type: str) -> None:
|
||||
if output_type not in {"image", "video"}:
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{output_type}', expected 'video' or legacy 'image' for video generation."
|
||||
)
|
||||
|
||||
|
||||
def _normalize_float_tensor(tensor: torch.Tensor, source_range: str) -> torch.Tensor:
|
||||
"""Normalize decoded frames according to the pipeline's output contract."""
|
||||
if not tensor.is_floating_point():
|
||||
@@ -620,10 +627,7 @@ def main():
|
||||
frames = frames[0] if frames else None
|
||||
|
||||
if isinstance(frames, OmniRequestOutput):
|
||||
if frames.final_output_type != "image":
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
|
||||
)
|
||||
_validate_video_output_type(frames.final_output_type)
|
||||
if frames.multimodal_output and "audio" in frames.multimodal_output:
|
||||
audio = frames.multimodal_output["audio"]
|
||||
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
|
||||
|
||||
@@ -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.
|
||||
@@ -18,11 +18,11 @@ examples in this repository.
|
||||
|
||||
Qwen3-TTS supports three task types, each backed by a dedicated model checkpoint:
|
||||
|
||||
| Task Type | Model | Description |
|
||||
| ------------- | ---------------------------------------- | ------------------------------------------------------------ |
|
||||
| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speaker voices with optional style/emotion control |
|
||||
| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Generate speech from a natural language voice description |
|
||||
| `Base` | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from reference audio + transcript |
|
||||
| Task Type | Model | Description |
|
||||
| ------------- | ---------------------------------------- | ------------------------------------------------------------- |
|
||||
| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speaker voices with optional style/emotion control |
|
||||
| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Generate speech from a natural language voice description |
|
||||
| `Base` | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from reference audio + transcript |
|
||||
|
||||
Smaller 0.6B variants are also available for `CustomVoice` and `Base`.
|
||||
|
||||
@@ -166,6 +166,15 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty
|
||||
- Key flags: `--omni` is required. `--deploy-config` points to the bundled two-stage pipeline config.
|
||||
- Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Raw audio streaming requires `stream=true`, `stream_format="audio"`, and `response_format="pcm"`.
|
||||
- Task/model matching: Each task type requires its matching model checkpoint. Using a CustomVoice model for a Base (voice clone) request will fail.
|
||||
- Base codec termination: Base requests without an explicit `max_new_tokens`
|
||||
use a text-scaled safety ceiling (at least 192 codec frames and no more than
|
||||
the configured model limit). If the Talker reaches that ceiling without
|
||||
codec EOS, non-streaming serving discards the incomplete audio and retries
|
||||
once with a fresh seed; an explicit `seed` or `max_new_tokens` disables the
|
||||
retry. SSE and WebSocket clients receive structured errors and must discard
|
||||
previously emitted audio when the error contains `"action":"discard"`;
|
||||
raw PCM streams terminate with a connection error after any partial bytes
|
||||
already sent.
|
||||
- Known limitations: The server serves one model variant at a time. To switch task types (e.g., CustomVoice to Base), restart the server with the corresponding model.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -46,9 +46,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
3646.797
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1098
|
||||
]
|
||||
@@ -82,9 +79,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
4059.0878
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1335
|
||||
]
|
||||
@@ -120,9 +114,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
5639.5769
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1798
|
||||
]
|
||||
|
||||
@@ -45,13 +45,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
15793.1209,
|
||||
30658.6632,
|
||||
37463.4943,
|
||||
47990.5839,
|
||||
64980.2199
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.169,
|
||||
0.2041,
|
||||
@@ -87,9 +80,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
4556.4347
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.138
|
||||
]
|
||||
@@ -123,9 +113,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
4768.0972
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1609
|
||||
]
|
||||
@@ -161,9 +148,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
6395.3517
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.2048
|
||||
]
|
||||
|
||||
@@ -11,13 +11,16 @@ to the always-kept prefix segment before generating the block mask.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout, VideoTokenSpan
|
||||
from vllm_omni.diffusion.attention.backends.rainfusion_attn import (
|
||||
_BLOCK_SIZE,
|
||||
RainFusionAttentionBackend,
|
||||
RainFusionAttentionImpl,
|
||||
RainFusionPlan,
|
||||
)
|
||||
@@ -97,6 +100,62 @@ def test_video_segment_must_be_the_tail_of_packed_document_zero():
|
||||
assert make_impl()._resolve_plan(metadata) is None
|
||||
|
||||
|
||||
def test_ref2va_multi_video_spans_resolve():
|
||||
layout = VideoTokenLayout(
|
||||
used_len=12000,
|
||||
video_spans=(
|
||||
VideoTokenSpan(start=128, latent_grid=(4, 16, 64), role="reference"),
|
||||
VideoTokenSpan(start=5000, latent_grid=(4, 16, 64), role="target"),
|
||||
),
|
||||
)
|
||||
plan = make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout))
|
||||
|
||||
assert plan is not None
|
||||
assert plan.used_len == 12000
|
||||
assert plan.video_spans == [
|
||||
{"start": 128, "latent_shape": [4, 16, 64]},
|
||||
{"start": 5000, "latent_shape": [4, 16, 64]},
|
||||
]
|
||||
|
||||
|
||||
def test_invalid_ref2va_video_spans_fall_back_to_dense():
|
||||
layout = VideoTokenLayout(
|
||||
used_len=12000,
|
||||
video_spans=(
|
||||
VideoTokenSpan(start=128, latent_grid=(4, 16, 64), role="reference"),
|
||||
VideoTokenSpan(start=4000, latent_grid=(4, 16, 64), role="target"),
|
||||
),
|
||||
)
|
||||
assert make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout)) is None
|
||||
|
||||
|
||||
def test_validate_available_rejects_legacy_mindiesd(monkeypatch):
|
||||
mindiesd = types.ModuleType("mindiesd")
|
||||
|
||||
def sparse_attention(query, key, value, **kwargs):
|
||||
return query
|
||||
|
||||
mindiesd.sparse_attention = sparse_attention
|
||||
monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
|
||||
monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
|
||||
|
||||
with pytest.raises(ValueError, match="video_spans"):
|
||||
RainFusionAttentionBackend.validate_available()
|
||||
|
||||
|
||||
def test_validate_available_accepts_new_mindiesd(monkeypatch):
|
||||
mindiesd = types.ModuleType("mindiesd")
|
||||
|
||||
def sparse_attention(query, key, value, *, video_spans=None, **kwargs):
|
||||
return query
|
||||
|
||||
mindiesd.sparse_attention = sparse_attention
|
||||
monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
|
||||
monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
|
||||
|
||||
RainFusionAttentionBackend.validate_available()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("grid", [(4, 24, 40), (1, 24, 40)])
|
||||
def test_short_video_stays_dense(grid):
|
||||
assert make_impl()._resolve_plan(make_metadata(grid)) is None
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
@@ -194,6 +194,47 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
|
||||
assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize(
|
||||
("fully_shard_degree", "error_message"),
|
||||
[(0, "fully_shard_degree must be positive"), (3, "must be divisible by fully_shard_degree")],
|
||||
)
|
||||
def test_invalid_hsdp_shard_size_fails_before_group_creation(monkeypatch, fully_shard_degree: int, error_message: str):
|
||||
"""Reject invalid HSDP shard sizes before allocating any process groups."""
|
||||
world_size = 4
|
||||
created_groups: list[object] = []
|
||||
|
||||
def fail_if_group_created(*args, **kwargs):
|
||||
del args, kwargs
|
||||
created_groups.append(object())
|
||||
raise AssertionError("a process group was created before HSDP validation")
|
||||
|
||||
fake_world_group = SimpleNamespace(
|
||||
rank_in_group=0,
|
||||
local_rank=0,
|
||||
device_group=object(),
|
||||
)
|
||||
monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
|
||||
monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
|
||||
monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
|
||||
monkeypatch.setattr(parallel_state, "init_model_parallel_group", fail_if_group_created)
|
||||
|
||||
for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_HSDP_REPLICATE", "_EXPERT_PARALLEL_GROUP_RANKS"):
|
||||
monkeypatch.setattr(parallel_state, name, None)
|
||||
for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
|
||||
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
|
||||
|
||||
with pytest.raises(ValueError, match=error_message):
|
||||
parallel_state.initialize_model_parallel(
|
||||
fully_shard_degree=fully_shard_degree,
|
||||
use_hsdp=True,
|
||||
backend="gloo",
|
||||
)
|
||||
|
||||
assert created_groups == []
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.core_model
|
||||
def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
|
||||
|
||||
@@ -9,6 +9,7 @@ Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
@@ -75,8 +76,13 @@ def _run_generate(
|
||||
if num_requests < 2:
|
||||
raise ValueError("num_requests must be >= 2 (1 warmup + >=1 timed)")
|
||||
|
||||
gc.collect()
|
||||
current_omni_platform.empty_cache()
|
||||
device_index = current_omni_platform.current_device()
|
||||
with current_omni_platform.device(device_index):
|
||||
free_bytes, total_bytes = current_omni_platform.mem_get_info()
|
||||
initial_used_mb = (total_bytes - free_bytes) / (1024**2)
|
||||
|
||||
monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
|
||||
monitor.start()
|
||||
try:
|
||||
@@ -113,7 +119,10 @@ def _run_generate(
|
||||
pass
|
||||
|
||||
median_time_s = float(np.median(per_request_times_s))
|
||||
peak_memory_mb = monitor.peak_used_mb
|
||||
# DeviceMemoryMonitor samples absolute device usage. Normalize by
|
||||
# this run's starting usage so compiler and allocator caches retained
|
||||
# by the preceding TP=1 run are not charged to the TP=2 run.
|
||||
peak_memory_mb = max(0.0, monitor.peak_used_mb - initial_used_mb)
|
||||
return _extract_single_image([last_output]), median_time_s, peak_memory_mb
|
||||
finally:
|
||||
monitor.stop()
|
||||
|
||||
@@ -1,11 +1,12 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Focused LTX-2.5 pipeline correctness tests."""
|
||||
|
||||
import json
|
||||
from dataclasses import replace
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -170,7 +171,10 @@ def _resolve_request_inputs_for_test(
|
||||
|
||||
|
||||
def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch):
|
||||
pipe = SimpleNamespace(component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None))
|
||||
pipe = SimpleNamespace(
|
||||
component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None),
|
||||
use_diffusion_decoder=True,
|
||||
)
|
||||
od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16)
|
||||
|
||||
monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
|
||||
@@ -187,7 +191,19 @@ def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatc
|
||||
|
||||
def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
|
||||
revision = "pinned-revision"
|
||||
calls = {"components": []}
|
||||
calls: dict[str, Any] = {"components": []}
|
||||
natten_processor = object()
|
||||
|
||||
class FakeDiffusionDecoder:
|
||||
def set_attn_processor(self, processor):
|
||||
calls["diffusion_decoder_processor"] = processor
|
||||
|
||||
def enable_tiling(self):
|
||||
calls["diffusion_decoder_tiling"] = True
|
||||
|
||||
def set_parallel_size(self, parallel_size, mode):
|
||||
calls["diffusion_decoder_parallel"] = (parallel_size, mode)
|
||||
|
||||
profile = replace(
|
||||
LTX25_DISTILLED_COMPONENT_PROFILE,
|
||||
text_encoder_cls=object,
|
||||
@@ -196,12 +212,19 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
|
||||
vocoder_fallback_cls=None,
|
||||
scheduler_use_dynamic_shifting=False,
|
||||
)
|
||||
pipeline = SimpleNamespace(component_profile=profile, pipeline_kind="distilled_two_stage")
|
||||
pipeline = SimpleNamespace(
|
||||
component_profile=profile,
|
||||
pipeline_kind="distilled_two_stage",
|
||||
use_diffusion_decoder=True,
|
||||
)
|
||||
od_config = SimpleNamespace(
|
||||
model="org/converted-ltx25",
|
||||
revision=revision,
|
||||
dtype=torch.bfloat16,
|
||||
quantization_config=None,
|
||||
extras={},
|
||||
vae_use_tiling=False,
|
||||
parallel_config=SimpleNamespace(vae_patch_parallel_size=2, vae_parallel_mode="tile"),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
|
||||
@@ -227,6 +250,10 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
|
||||
)
|
||||
return object()
|
||||
|
||||
def fake_native_diffusion_decoder(model, **kwargs):
|
||||
calls["native_diffusion_decoder"] = (model, kwargs)
|
||||
return FakeDiffusionDecoder()
|
||||
|
||||
def fake_transformer_config(model, subfolder, local_files_only, *, revision):
|
||||
calls["transformer_config"] = (model, subfolder, local_files_only, revision)
|
||||
return {"component": "transformer"}
|
||||
@@ -234,7 +261,13 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
|
||||
monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch)
|
||||
monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer)
|
||||
monkeypatch.setattr(ltx2_components, "_load_component", fake_component)
|
||||
monkeypatch.setattr(ltx2_components, "_load_ltx25_native_diffusion_decoder", fake_native_diffusion_decoder)
|
||||
monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config)
|
||||
monkeypatch.setattr(
|
||||
ltx2_components,
|
||||
"LTX2VideoVaeNeighborhoodNattenProcessor",
|
||||
lambda: natten_processor,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ltx2_components,
|
||||
"create_transformer_from_config",
|
||||
@@ -261,6 +294,14 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
|
||||
"latent_upsampler",
|
||||
}
|
||||
assert all(call[3]["revision"] == revision for call in calls["components"])
|
||||
assert "diffusion_decoder" not in calls["prefetch"][1]
|
||||
assert calls["native_diffusion_decoder"] == (
|
||||
od_config.model,
|
||||
{"local_files_only": False, "dtype": torch.bfloat16, "revision": revision},
|
||||
)
|
||||
assert calls["diffusion_decoder_processor"] is natten_processor
|
||||
assert calls["diffusion_decoder_tiling"] is True
|
||||
assert calls["diffusion_decoder_parallel"] == (2, "tile")
|
||||
assert calls["transformer_config"] == (
|
||||
od_config.model,
|
||||
profile.transformer_subfolder,
|
||||
@@ -270,6 +311,85 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
|
||||
assert calls["scheduler"][1]["revision"] == revision
|
||||
|
||||
|
||||
def test_ltx25_native_diffusion_decoder_uses_diffusers_config_and_canonical_artifact(monkeypatch):
|
||||
revision = "diffusers-revision"
|
||||
config = {"decoder_stage_channels": [8, 8, 8, 8, 8]}
|
||||
calls: dict[str, Any] = {}
|
||||
|
||||
class FakeDecoder:
|
||||
def init_distributed(self):
|
||||
calls["init_distributed"] = True
|
||||
|
||||
expected = FakeDecoder()
|
||||
|
||||
def fake_load_config(model, **kwargs):
|
||||
calls["config"] = (model, kwargs)
|
||||
return config
|
||||
|
||||
def fake_resolve(model, repo_id, filename, **kwargs):
|
||||
calls["artifact"] = (model, repo_id, filename, kwargs)
|
||||
return "/models/native-diffvae.safetensors"
|
||||
|
||||
def fake_load(cls, path, passed_config, dtype):
|
||||
calls["load"] = (path, passed_config, dtype)
|
||||
return expected
|
||||
|
||||
monkeypatch.setattr(
|
||||
ltx2_components.DistributedLTX2VideoDiffusionDecoderModel,
|
||||
"load_config",
|
||||
staticmethod(fake_load_config),
|
||||
)
|
||||
monkeypatch.setattr(ltx2_components, "resolve_ltx_artifact", fake_resolve)
|
||||
monkeypatch.setattr(
|
||||
ltx2_components.DistributedLTX2VideoDiffusionDecoderModel,
|
||||
"from_ltx25_native_checkpoint",
|
||||
classmethod(fake_load),
|
||||
)
|
||||
|
||||
actual = ltx2_components._load_ltx25_native_diffusion_decoder(
|
||||
"Lightricks/LTX-2.5-Diffusers",
|
||||
local_files_only=False,
|
||||
dtype=torch.bfloat16,
|
||||
revision=revision,
|
||||
)
|
||||
|
||||
assert actual is expected
|
||||
assert calls["config"] == (
|
||||
"Lightricks/LTX-2.5-Diffusers",
|
||||
{
|
||||
"subfolder": "diffusion_decoder",
|
||||
"local_files_only": False,
|
||||
"revision": revision,
|
||||
},
|
||||
)
|
||||
assert calls["artifact"] == (
|
||||
"Lightricks/LTX-2.5-Diffusers",
|
||||
"Lightricks/LTX-2.5",
|
||||
"vae/ltx-2.5-video-vae-bf16.safetensors",
|
||||
{
|
||||
"model_revision": revision,
|
||||
"artifact_revision": ltx2_components.LTX25_NATIVE_ARTIFACT_REVISION,
|
||||
},
|
||||
)
|
||||
assert calls["load"] == ("/models/native-diffvae.safetensors", config, torch.bfloat16)
|
||||
assert calls["init_distributed"] is True
|
||||
|
||||
|
||||
def test_ltx25_natten_failure_has_actionable_omni_remedy(monkeypatch):
|
||||
class MissingNattenProcessor:
|
||||
def __init__(self):
|
||||
raise ImportError("kernels is unavailable")
|
||||
|
||||
monkeypatch.setattr(
|
||||
ltx2_components,
|
||||
"LTX2VideoVaeNeighborhoodNattenProcessor",
|
||||
MissingNattenProcessor,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match=r"kernels==0\.14\.1.*supported GPU.*allow Hub access"):
|
||||
ltx2_components._create_ltx25_natten_processor()
|
||||
|
||||
|
||||
def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path):
|
||||
(tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"}))
|
||||
(tmp_path / "transformer").mkdir()
|
||||
@@ -411,6 +531,46 @@ def test_ltx25_four_public_pipeline_semantics_are_disjoint():
|
||||
assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pipeline_cls",
|
||||
[
|
||||
LTX2Pipeline,
|
||||
LTX2TwoStagePipeline,
|
||||
LTX2DistilledOneStagePipeline,
|
||||
LTX2DistilledTwoStagePipeline,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
("extras", "use_diffusion_decoder"),
|
||||
[({}, True), ({"ltx2_use_conv_vae": True}, False)],
|
||||
)
|
||||
def test_ltx25_all_public_pipelines_default_to_diffvae_with_conv_opt_in(
|
||||
tmp_path, monkeypatch, pipeline_cls, extras, use_diffusion_decoder
|
||||
):
|
||||
(tmp_path / "model_index.json").write_text(
|
||||
json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
|
||||
)
|
||||
|
||||
def stub_components(pipe, od_config):
|
||||
pipe.od_config = od_config
|
||||
pipe.vae_spatial_compression_ratio = 32
|
||||
|
||||
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
|
||||
monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda _pipe: None)
|
||||
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
|
||||
|
||||
pipe = pipeline_cls(
|
||||
od_config=SimpleNamespace(
|
||||
model=str(tmp_path),
|
||||
extras=extras,
|
||||
enable_diffusion_pipeline_profiler=False,
|
||||
)
|
||||
)
|
||||
|
||||
assert pipe.use_diffusion_decoder is use_diffusion_decoder
|
||||
assert pipe._vae_modules.count("diffusion_decoder") == int(use_diffusion_decoder)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("pipeline_kind", "expected"),
|
||||
[
|
||||
@@ -673,7 +833,7 @@ def test_ltx25_distilled_two_stage_executes_custom_phase_schedules():
|
||||
max_sequence_length=16,
|
||||
)
|
||||
prompt_context = object()
|
||||
phase_calls = []
|
||||
phase_calls: list[Any] = []
|
||||
|
||||
def resolve_request_inputs(req, **kwargs):
|
||||
return request_inputs
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Unit tests for LTX image-to-video input and conditioning behavior."""
|
||||
|
||||
@@ -258,6 +258,39 @@ class TestLTXImageToVideoConditioning:
|
||||
|
||||
assert LTX2Pipeline.support_image_input is True
|
||||
|
||||
def test_ltx25_i2v_encode_uses_diffusion_decoder_statistics(self, monkeypatch):
|
||||
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
|
||||
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
||||
|
||||
pipe = object.__new__(LTX2Pipeline)
|
||||
torch.nn.Module.__init__(pipe)
|
||||
pipe.use_diffusion_decoder = True
|
||||
pipe.vae = SimpleNamespace(
|
||||
encode=lambda image: image,
|
||||
latents_mean=torch.full((2,), 100.0),
|
||||
latents_std=torch.full((2,), 10.0),
|
||||
config=SimpleNamespace(scaling_factor=1.0),
|
||||
)
|
||||
pipe.diffusion_decoder = SimpleNamespace(
|
||||
latents_mean=torch.tensor([1.0, 3.0]),
|
||||
latents_std=torch.tensor([2.0, 4.0]),
|
||||
config=SimpleNamespace(scaling_factor=2.0),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ltx2_conditioning,
|
||||
"retrieve_latents",
|
||||
lambda *_args, **_kwargs: torch.tensor([[[[[3.0]]], [[[7.0]]]]]),
|
||||
)
|
||||
|
||||
actual = pipe._encode_i2v_image_latents(
|
||||
torch.zeros(1, 3, 1, 1),
|
||||
batch_size=1,
|
||||
generator=None,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual, torch.full((1, 2, 1, 1, 1), 2.0))
|
||||
|
||||
def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
|
||||
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
||||
|
||||
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
@@ -47,6 +50,39 @@ def test_prepare_video_latents_matches_official_values_and_token_major_layout():
|
||||
assert actual.stride()[1:] == (1, actual.shape[1])
|
||||
|
||||
|
||||
def test_prepare_video_latents_uses_diffusion_decoder_statistics():
|
||||
pipeline = _make_pipeline(LTX2Pipeline)
|
||||
pipeline.use_diffusion_decoder = True
|
||||
pipeline.vae_spatial_compression_ratio = 1
|
||||
pipeline.vae_temporal_compression_ratio = 1
|
||||
pipeline.transformer_spatial_patch_size = 1
|
||||
pipeline.transformer_temporal_patch_size = 1
|
||||
pipeline.vae = SimpleNamespace(
|
||||
latents_mean=torch.full((2,), 100.0),
|
||||
latents_std=torch.full((2,), 10.0),
|
||||
config=SimpleNamespace(scaling_factor=1.0),
|
||||
)
|
||||
pipeline.diffusion_decoder = SimpleNamespace(
|
||||
latents_mean=torch.tensor([1.0, 3.0]),
|
||||
latents_std=torch.tensor([2.0, 4.0]),
|
||||
config=SimpleNamespace(scaling_factor=2.0),
|
||||
)
|
||||
|
||||
actual = pipeline.prepare_latents(
|
||||
batch_size=1,
|
||||
num_channels_latents=2,
|
||||
height=1,
|
||||
width=1,
|
||||
num_frames=1,
|
||||
noise_scale=0.0,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
latents=torch.tensor([[[[[3.0]]], [[[7.0]]]]]),
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual, torch.full((1, 1, 2), 2.0))
|
||||
|
||||
|
||||
def test_prepare_audio_latents_samples_directly_in_packed_token_space():
|
||||
pipeline = _make_pipeline(LTX2Pipeline)
|
||||
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Unit tests for the fixed LTX refinement-phase adapter."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -148,10 +149,13 @@ def test_ltx_phase_adapter_selects_execution_from_quantization(
|
||||
):
|
||||
pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config)
|
||||
manifest = object()
|
||||
captured = {}
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
def resolve_artifact(*_args, **kwargs):
|
||||
assert not kwargs
|
||||
assert kwargs == {
|
||||
"model_revision": None,
|
||||
"artifact_revision": None,
|
||||
}
|
||||
return "adapter"
|
||||
|
||||
def parse(transformer, path):
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Unit tests for the shared LTX pipeline runtime and public contracts."""
|
||||
|
||||
@@ -8,6 +8,7 @@ import os
|
||||
import tempfile
|
||||
from dataclasses import replace
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -165,7 +166,19 @@ def test_ltx_checkpoint_version_detection_uses_metadata(tmp_path):
|
||||
assert detect_ltx_model_version(str(tmp_path)) == "2"
|
||||
|
||||
|
||||
def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch):
|
||||
@pytest.mark.parametrize(
|
||||
("model", "repo_id", "expected_revision"),
|
||||
[
|
||||
("Lightricks/test", "Lightricks/test", "model-revision"),
|
||||
("Lightricks/test-Diffusers", "Lightricks/test", "artifact-revision"),
|
||||
],
|
||||
)
|
||||
def test_ltx_artifact_uses_source_revision_and_hub_fallback(
|
||||
monkeypatch,
|
||||
model,
|
||||
repo_id,
|
||||
expected_revision,
|
||||
):
|
||||
filename = "ltx-sidecar.safetensors"
|
||||
calls = []
|
||||
|
||||
@@ -176,9 +189,22 @@ def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch)
|
||||
monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
|
||||
|
||||
assert (
|
||||
resolve_ltx_artifact(str(tmp_path / "model"), "Lightricks/test", filename) == "/cache/ltx-sidecar.safetensors"
|
||||
resolve_ltx_artifact(
|
||||
model,
|
||||
repo_id,
|
||||
filename,
|
||||
model_revision="model-revision",
|
||||
artifact_revision="artifact-revision",
|
||||
)
|
||||
== "/cache/ltx-sidecar.safetensors"
|
||||
)
|
||||
assert calls == [{"repo_id": "Lightricks/test", "filename": filename}]
|
||||
assert calls == [
|
||||
{
|
||||
"repo_id": repo_id,
|
||||
"filename": filename,
|
||||
"revision": expected_revision,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
|
||||
@@ -187,7 +213,41 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
|
||||
expected.write_bytes(b"sidecar")
|
||||
monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup"))
|
||||
|
||||
assert resolve_ltx_artifact(str(tmp_path), "Lightricks/test", filename) == str(expected)
|
||||
assert resolve_ltx_artifact(
|
||||
str(tmp_path),
|
||||
"Lightricks/test",
|
||||
filename,
|
||||
model_revision="unused-model-revision",
|
||||
artifact_revision="unused-artifact-revision",
|
||||
) == str(expected)
|
||||
|
||||
|
||||
def test_ltx_artifact_local_model_missing_sidecar_falls_back_to_hub(tmp_path, monkeypatch):
|
||||
calls = []
|
||||
|
||||
def fake_download(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return "/cache/ltx-sidecar.safetensors"
|
||||
|
||||
monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
|
||||
|
||||
assert (
|
||||
resolve_ltx_artifact(
|
||||
str(tmp_path),
|
||||
"Lightricks/test",
|
||||
"ltx-sidecar.safetensors",
|
||||
model_revision="local-model-revision",
|
||||
artifact_revision="artifact-revision",
|
||||
)
|
||||
== "/cache/ltx-sidecar.safetensors"
|
||||
)
|
||||
assert calls == [
|
||||
{
|
||||
"repo_id": "Lightricks/test",
|
||||
"filename": "ltx-sidecar.safetensors",
|
||||
"revision": "artifact-revision",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -900,7 +960,7 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
|
||||
max_sequence_length=16,
|
||||
)
|
||||
prompt_context = object()
|
||||
phase_calls = []
|
||||
phase_calls: list[Any] = []
|
||||
|
||||
def resolve_request_inputs(req, **kwargs):
|
||||
return request_inputs
|
||||
@@ -959,7 +1019,7 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
|
||||
pipeline.vae_spatial_compression_ratio = 32
|
||||
pipeline.vae_temporal_compression_ratio = 8
|
||||
pipeline.latent_upsampler = FakeUpsampler()
|
||||
activated_slots = []
|
||||
activated_slots: list[Any] = []
|
||||
if pipeline_cls is LTX2TwoStagePipeline:
|
||||
pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append)
|
||||
object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
|
||||
@@ -1133,6 +1193,59 @@ def test_ltx_two_stage_recipe_rejects_sigmas():
|
||||
)
|
||||
|
||||
|
||||
def test_ltx_request_applies_diffvae_minimum_spatial_size_only_when_requested():
|
||||
request = LTXRequestInputs(
|
||||
prompt="prompt",
|
||||
negative_prompt="",
|
||||
height=128,
|
||||
width=128,
|
||||
num_frames=9,
|
||||
frame_rate=24.0,
|
||||
num_inference_steps=30,
|
||||
guidance=LTX2_ONE_STAGE_RECIPE.request_guidance,
|
||||
num_videos_per_prompt=1,
|
||||
generator=None,
|
||||
latents=None,
|
||||
audio_latents=None,
|
||||
prompt_embeds=None,
|
||||
negative_prompt_embeds=None,
|
||||
prompt_attention_mask=None,
|
||||
negative_prompt_attention_mask=None,
|
||||
decode_timestep=0.0,
|
||||
decode_noise_scale=None,
|
||||
output_type="np",
|
||||
max_sequence_length=16,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="DiffVAE requires video latent spatial dimensions of at least 7x7"):
|
||||
validate_pipeline_request(
|
||||
request,
|
||||
pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
|
||||
vae_spatial_compression_ratio=32,
|
||||
vae_temporal_compression_ratio=8,
|
||||
pipeline_name="LTX2Pipeline",
|
||||
min_video_latent_spatial_size=(7, 7),
|
||||
)
|
||||
|
||||
# ConvVAE does not use NATTEN, so its existing 128x128 boundary remains valid.
|
||||
validate_pipeline_request(
|
||||
request,
|
||||
pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
|
||||
vae_spatial_compression_ratio=32,
|
||||
vae_temporal_compression_ratio=8,
|
||||
pipeline_name="LTX2Pipeline",
|
||||
)
|
||||
|
||||
validate_pipeline_request(
|
||||
replace(request, height=256, width=256),
|
||||
pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
|
||||
vae_spatial_compression_ratio=32,
|
||||
vae_temporal_compression_ratio=8,
|
||||
pipeline_name="LTX2Pipeline",
|
||||
min_video_latent_spatial_size=(7, 7),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("direct_kwargs", "sampling_kwargs", "error"),
|
||||
[
|
||||
@@ -1284,7 +1397,7 @@ def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, step
|
||||
|
||||
|
||||
def test_ltx_two_stage_entries_select_the_official_adapter_slots():
|
||||
phase_switches = []
|
||||
phase_switches: list[Any] = []
|
||||
ordinary_pipeline = object.__new__(LTX2TwoStagePipeline)
|
||||
ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append)
|
||||
for phase in LTX2_TWO_STAGE_RECIPE.phases:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from collections.abc import Iterator
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLEncoder
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def restore_cudnn_sdp_state() -> Iterator[None]:
|
||||
original = torch.backends.cuda.cudnn_sdp_enabled()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch.backends.cuda.enable_cudnn_sdp(original)
|
||||
|
||||
|
||||
def _encoder_without_weights() -> MiniMaxH3Qwen3VLEncoder:
|
||||
encoder = MiniMaxH3Qwen3VLEncoder("unused", device=torch.device("cpu"), load_model=False)
|
||||
encoder.text_model = nn.Linear(1, 1)
|
||||
return encoder
|
||||
|
||||
|
||||
@pytest.mark.parametrize("initial_state", [False, True])
|
||||
def test_encode_ids_restores_cudnn_sdp_after_success(monkeypatch: pytest.MonkeyPatch, initial_state: bool) -> None:
|
||||
encoder = _encoder_without_weights()
|
||||
expected = torch.tensor([1.0])
|
||||
|
||||
def encode(*args, **kwargs):
|
||||
assert torch.backends.cuda.cudnn_sdp_enabled()
|
||||
return expected
|
||||
|
||||
monkeypatch.setattr(encoder, "_encode", encode)
|
||||
torch.backends.cuda.enable_cudnn_sdp(initial_state)
|
||||
|
||||
actual = encoder.encode_ids(torch.tensor([1]))
|
||||
|
||||
torch.testing.assert_close(actual, expected)
|
||||
assert torch.backends.cuda.cudnn_sdp_enabled() is initial_state
|
||||
|
||||
|
||||
@pytest.mark.parametrize("initial_state", [False, True])
|
||||
def test_encode_ids_restores_cudnn_sdp_after_failure(monkeypatch: pytest.MonkeyPatch, initial_state: bool) -> None:
|
||||
encoder = _encoder_without_weights()
|
||||
|
||||
def fail(*args, **kwargs):
|
||||
assert torch.backends.cuda.cudnn_sdp_enabled()
|
||||
raise RuntimeError("encode failed")
|
||||
|
||||
monkeypatch.setattr(encoder, "_encode", fail)
|
||||
torch.backends.cuda.enable_cudnn_sdp(initial_state)
|
||||
|
||||
with pytest.raises(RuntimeError, match="encode failed"):
|
||||
encoder.encode_ids(torch.tensor([1]))
|
||||
|
||||
assert torch.backends.cuda.cudnn_sdp_enabled() is initial_state
|
||||
@@ -74,6 +74,8 @@ def test_t2va_and_fl2va_packing_keep_update_rows_separate():
|
||||
assert fl2va["update_mask"].sum().item() == 12
|
||||
assert (~fl2va["update_mask"]).sum().item() == 6
|
||||
assert t2va["audio_pos"].numel() == fl2va["audio_pos"].numel() == 6
|
||||
assert t2va["video_spans"] == ({"start": 10, "latent_grid": (2, 2, 3), "role": "target"},)
|
||||
assert fl2va["video_spans"] == ({"start": 16, "latent_grid": (2, 2, 3), "role": "target"},)
|
||||
|
||||
|
||||
def test_ref2va_packing_tracks_video_and_audio_update_masks():
|
||||
@@ -101,6 +103,33 @@ def test_ref2va_packing_tracks_video_and_audio_update_masks():
|
||||
assert packed["cu_seqlens"].tolist() == [0, 30, 64]
|
||||
|
||||
|
||||
def test_ref2va_packing_publishes_physical_video_spans_only():
|
||||
from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
|
||||
minimax_h3_packed_sequence_ref2va_blocks,
|
||||
)
|
||||
|
||||
packed = minimax_h3_packed_sequence_ref2va_blocks(
|
||||
text_len=4,
|
||||
latent_t=2,
|
||||
latent_h=4,
|
||||
latent_w=6,
|
||||
audio_t=3,
|
||||
ref_blocks=[
|
||||
{"kind": "image", "latent_h": 4, "latent_w": 4},
|
||||
{"kind": "video_audio", "ref_audio_t": 2, "latent_t": 2, "latent_h": 4, "latent_w": 6},
|
||||
{"kind": "audio", "ref_audio_t": 1},
|
||||
{"kind": "video", "ref_audio_t": 0, "latent_t": 1, "latent_h": 4, "latent_w": 4},
|
||||
],
|
||||
)
|
||||
|
||||
# Image and all audio rows intentionally stay out of the sparse spans.
|
||||
assert packed["video_spans"] == (
|
||||
{"start": 12, "latent_grid": (2, 2, 3), "role": "reference"},
|
||||
{"start": 26, "latent_grid": (1, 2, 2), "role": "reference"},
|
||||
{"start": 36, "latent_grid": (2, 2, 3), "role": "target"},
|
||||
)
|
||||
|
||||
|
||||
def test_condition_noise_is_seeded_and_keeps_clean_anchor_at_timestep_one():
|
||||
from vllm_omni.diffusion.models.minimax_h3.condition_noise import (
|
||||
minimax_h3_audio_cond_noise_aug_rows,
|
||||
|
||||
@@ -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
|
||||
@@ -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 pathlib import Path
|
||||
|
||||
@@ -187,6 +187,7 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
|
||||
"model": model_ref,
|
||||
"enforce_eager": True,
|
||||
"tensor_parallel_size": 2,
|
||||
# The fused BF16 baseline only fits on H100-80GB with encoder TP.
|
||||
"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)"
|
||||
)
|
||||
@@ -81,7 +81,7 @@ class QualityTestConfig:
|
||||
num_frames: int = 5 # only for t2v
|
||||
seed: int = 42
|
||||
gpu: str = "H100" # minimum GPU requirement
|
||||
negative_prompt: str = ""
|
||||
negative_prompt: str | None = ""
|
||||
guidance_scale: float | None = None
|
||||
sigmas: list[float] | None = None
|
||||
enable_cpu_offload: bool = False
|
||||
@@ -212,24 +212,14 @@ QUALITY_CONFIGS = [
|
||||
quantization="fp8",
|
||||
task="t2v",
|
||||
prompt="A serene lakeside sunrise with mist over the water",
|
||||
max_lpips=0.10,
|
||||
height=256,
|
||||
width=256,
|
||||
num_frames=25,
|
||||
num_inference_steps=8,
|
||||
# Preserve the final scheduler trajectory used when this FP8 quality
|
||||
# gate was established. Explicit LTX sigmas bypass dynamic shifting.
|
||||
sigmas=[
|
||||
1.0,
|
||||
0.92185378074646,
|
||||
0.8327768445014954,
|
||||
0.7303033471107483,
|
||||
0.6111654043197632,
|
||||
0.4709382653236389,
|
||||
0.30347853899002075,
|
||||
0.10000002384185791,
|
||||
0.0,
|
||||
],
|
||||
max_lpips=0.20,
|
||||
height=384,
|
||||
width=512,
|
||||
num_frames=73,
|
||||
num_inference_steps=10,
|
||||
# Do not override the recipe's negative conditioning, guidance, or
|
||||
# scheduler trajectory: this gate follows the supported LTX default.
|
||||
negative_prompt=None,
|
||||
),
|
||||
]
|
||||
|
||||
@@ -296,8 +286,11 @@ def _generate_image(omni, config: QualityTestConfig):
|
||||
).manual_seed(config.seed)
|
||||
torch.accelerator.reset_peak_memory_stats()
|
||||
|
||||
prompt = {"prompt": config.prompt}
|
||||
if config.negative_prompt is not None:
|
||||
prompt["negative_prompt"] = config.negative_prompt
|
||||
outputs = omni.generate(
|
||||
{"prompt": config.prompt, "negative_prompt": config.negative_prompt},
|
||||
prompt,
|
||||
OmniDiffusionSamplingParams(
|
||||
height=config.height,
|
||||
width=config.width,
|
||||
@@ -328,8 +321,11 @@ def _generate_video(omni, config: QualityTestConfig):
|
||||
).manual_seed(config.seed)
|
||||
torch.accelerator.reset_peak_memory_stats()
|
||||
|
||||
prompt = {"prompt": config.prompt}
|
||||
if config.negative_prompt is not None:
|
||||
prompt["negative_prompt"] = config.negative_prompt
|
||||
outputs = omni.generate(
|
||||
{"prompt": config.prompt, "negative_prompt": config.negative_prompt},
|
||||
prompt,
|
||||
OmniDiffusionSamplingParams(
|
||||
height=config.height,
|
||||
width=config.width,
|
||||
@@ -465,6 +461,7 @@ def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch
|
||||
|
||||
|
||||
def test_generate_video_forwards_sigmas(monkeypatch):
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
monkeypatch.setattr(current_omni_platform, "device_type", "cpu", raising=False)
|
||||
@@ -475,7 +472,12 @@ def test_generate_video_forwards_sigmas(monkeypatch):
|
||||
class DummyOmni:
|
||||
def generate(self, _prompt, sampling):
|
||||
captured.sampling = sampling
|
||||
return [SimpleNamespace(images=[np.zeros((1, 2, 2, 3), dtype=np.float32)])]
|
||||
return [
|
||||
OmniRequestOutput.from_diffusion(
|
||||
request_id="req",
|
||||
images=[np.zeros((1, 2, 2, 3), dtype=np.float32)],
|
||||
)
|
||||
]
|
||||
|
||||
config = QualityTestConfig(
|
||||
id="ltx-sigmas",
|
||||
@@ -491,6 +493,41 @@ def test_generate_video_forwards_sigmas(monkeypatch):
|
||||
assert captured.sampling.sigmas == [1.0, 0.5]
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_ltx_quality_gate_uses_official_eager_defaults(monkeypatch):
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
monkeypatch.setattr(current_omni_platform, "device_type", "cpu", raising=False)
|
||||
monkeypatch.setattr(torch.accelerator, "reset_peak_memory_stats", lambda: None, raising=False)
|
||||
monkeypatch.setattr(torch.accelerator, "max_memory_allocated", lambda: 0, raising=False)
|
||||
captured = SimpleNamespace(prompt=None, sampling=None)
|
||||
|
||||
class DummyOmni:
|
||||
def generate(self, prompt, sampling):
|
||||
captured.prompt = prompt
|
||||
captured.sampling = sampling
|
||||
return [
|
||||
OmniRequestOutput.from_diffusion(
|
||||
request_id="req",
|
||||
images=[np.zeros((1, 2, 2, 3), dtype=np.float32)],
|
||||
)
|
||||
]
|
||||
|
||||
config = next(config for config in QUALITY_CONFIGS if config.id == "fp8_ltx2")
|
||||
_generate_video(DummyOmni(), config)
|
||||
|
||||
assert (config.width, config.height, config.num_frames, config.num_inference_steps) == (512, 384, 73, 10)
|
||||
assert config.max_lpips == 0.20
|
||||
assert config.sigmas is None
|
||||
assert config.guidance_scale is None
|
||||
assert config.negative_prompt is None
|
||||
assert "negative_prompt" not in captured.prompt
|
||||
assert captured.sampling.sigmas is None
|
||||
assert captured.sampling.guidance_scale is None
|
||||
|
||||
|
||||
_marks = hardware_marks(res={"cuda": "H100"})
|
||||
_OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUALITY_OUTPUT_DIR" in os.environ else None
|
||||
|
||||
@@ -501,8 +538,6 @@ def _quality_param(c: QualityTestConfig):
|
||||
marks.append(
|
||||
pytest.mark.skip(reason="Qwen-Image FP8 quality gate temporarily disabled (see CI / issue tracker).")
|
||||
)
|
||||
if c.id == "fp8_ltx2":
|
||||
marks.append(pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/6245"))
|
||||
return pytest.param(c, id=c.id, marks=marks)
|
||||
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -76,7 +76,7 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata(
|
||||
result = results[0]
|
||||
assert result.images == ["frame-0"]
|
||||
assert result.prompt == "prompt-0"
|
||||
assert result.final_output_type == "image"
|
||||
assert result.final_output_type == "video"
|
||||
assert result.multimodal_output == {
|
||||
"metadata": {
|
||||
"audio": {"sample_rate": 48000},
|
||||
@@ -315,7 +315,7 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
|
||||
|
||||
assert result.images == ["frame-0"]
|
||||
assert result.prompt == "watch and listen"
|
||||
assert result.final_output_type == "image"
|
||||
assert result.final_output_type == "video"
|
||||
assert result.multimodal_output == {
|
||||
"metadata": {
|
||||
"audio": {"sample_rate": 16000},
|
||||
|
||||
@@ -311,8 +311,8 @@ class TestPipelineStreamingOutputToEntrypoint:
|
||||
lambda *, output_format, fps, video_codec_options=None: FakeStreamingVideoEncoder(),
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
|
||||
return_value="diffusion",
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.is_video_generation_pipeline",
|
||||
return_value=True,
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
|
||||
|
||||
@@ -0,0 +1,255 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import importlib.util
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
import vllm_omni.diffusion.worker.diffusion_worker as worker_module
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _FakeGroup:
|
||||
world_size: int = 1
|
||||
rank_in_group: int = 0
|
||||
|
||||
|
||||
_GROUP_GETTERS = (
|
||||
"get_dp_group",
|
||||
"get_pp_group",
|
||||
"get_sp_group",
|
||||
"get_cfg_group",
|
||||
"get_tp_group",
|
||||
"get_fs_group",
|
||||
"get_hsdp_replicate_group",
|
||||
"get_ep_group",
|
||||
)
|
||||
|
||||
|
||||
def _patch_groups(
|
||||
mocker: MockerFixture,
|
||||
groups: dict[str, _FakeGroup],
|
||||
) -> dict[str, Mock]:
|
||||
patched_getters = {}
|
||||
for getter_name in _GROUP_GETTERS:
|
||||
patched_getters[getter_name] = mocker.patch.object(
|
||||
worker_module,
|
||||
getter_name,
|
||||
return_value=groups.get(getter_name, _FakeGroup()),
|
||||
)
|
||||
return patched_getters
|
||||
|
||||
|
||||
def test_setup_uses_default_name_before_model_parallel_init(
|
||||
mocker: MockerFixture,
|
||||
) -> None:
|
||||
mocker.patch.object(
|
||||
worker_module,
|
||||
"model_parallel_is_initialized",
|
||||
return_value=False,
|
||||
)
|
||||
group_getters = _patch_groups(mocker, {})
|
||||
set_process_title = mocker.patch.object(worker_module, "set_process_title")
|
||||
decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
|
||||
|
||||
worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=True, use_hsdp=True)
|
||||
|
||||
set_process_title.assert_called_once_with(
|
||||
name="DiffusionWorker",
|
||||
prefix="vLLM-Omni",
|
||||
)
|
||||
decorate_logs.assert_called_once_with("DiffusionWorker")
|
||||
for group_getter in group_getters.values():
|
||||
group_getter.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("groups", "enable_ep", "use_hsdp", "hsdp_replicate_size", "expected_name"),
|
||||
[
|
||||
({}, False, False, 1, "DiffusionWorker"),
|
||||
({"get_tp_group": _FakeGroup(2, 1)}, False, False, 1, "DiffusionWorker_TP1"),
|
||||
({"get_dp_group": _FakeGroup(4, 2)}, False, False, 1, "DiffusionWorker_DP2"),
|
||||
({"get_pp_group": _FakeGroup(2, 1)}, False, False, 1, "DiffusionWorker_PP1"),
|
||||
(
|
||||
{
|
||||
"get_cfg_group": _FakeGroup(2, 1),
|
||||
"get_tp_group": _FakeGroup(2, 1),
|
||||
},
|
||||
False,
|
||||
False,
|
||||
1,
|
||||
"DiffusionWorker_CFG1_TP1",
|
||||
),
|
||||
(
|
||||
{
|
||||
"get_dp_group": _FakeGroup(2, 1),
|
||||
"get_sp_group": _FakeGroup(4, 2),
|
||||
"get_tp_group": _FakeGroup(2, 1),
|
||||
},
|
||||
False,
|
||||
False,
|
||||
1,
|
||||
"DiffusionWorker_DP1_SP2_TP1",
|
||||
),
|
||||
(
|
||||
{"get_fs_group": _FakeGroup(4, 3)},
|
||||
False,
|
||||
True,
|
||||
1,
|
||||
"DiffusionWorker_FS3",
|
||||
),
|
||||
(
|
||||
{
|
||||
"get_sp_group": _FakeGroup(2, 1),
|
||||
"get_fs_group": _FakeGroup(4, 3),
|
||||
},
|
||||
False,
|
||||
True,
|
||||
1,
|
||||
"DiffusionWorker_SP1_FS3",
|
||||
),
|
||||
(
|
||||
{
|
||||
"get_fs_group": _FakeGroup(1, 0),
|
||||
"get_hsdp_replicate_group": _FakeGroup(2, 1),
|
||||
},
|
||||
False,
|
||||
True,
|
||||
2,
|
||||
"DiffusionWorker_RP1",
|
||||
),
|
||||
({"get_ep_group": _FakeGroup(4, 2)}, True, False, 1, "DiffusionWorker_EP2"),
|
||||
({}, True, False, 1, "DiffusionWorker"),
|
||||
],
|
||||
ids=[
|
||||
"all-singleton",
|
||||
"tp-only",
|
||||
"dp-only",
|
||||
"pp-only",
|
||||
"tp-cfg",
|
||||
"dp-sp-tp",
|
||||
"hsdp-only",
|
||||
"hsdp-sp",
|
||||
"hsdp-replicated",
|
||||
"expert-parallel",
|
||||
"singleton-ep",
|
||||
],
|
||||
)
|
||||
def test_setup_uses_initialized_parallel_groups(
|
||||
mocker: MockerFixture,
|
||||
groups: dict[str, _FakeGroup],
|
||||
enable_ep: bool,
|
||||
use_hsdp: bool,
|
||||
hsdp_replicate_size: int,
|
||||
expected_name: str,
|
||||
) -> None:
|
||||
mocker.patch.object(
|
||||
worker_module,
|
||||
"model_parallel_is_initialized",
|
||||
return_value=True,
|
||||
)
|
||||
group_getters = _patch_groups(mocker, groups)
|
||||
set_process_title = mocker.patch.object(worker_module, "set_process_title")
|
||||
decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
|
||||
|
||||
worker_module._setup_diffusion_worker_proc_title_and_log_prefix(
|
||||
enable_ep=enable_ep,
|
||||
use_hsdp=use_hsdp,
|
||||
hsdp_replicate_size=hsdp_replicate_size,
|
||||
)
|
||||
|
||||
set_process_title.assert_called_once_with(
|
||||
name=expected_name,
|
||||
prefix="vLLM-Omni",
|
||||
)
|
||||
decorate_logs.assert_called_once_with(expected_name)
|
||||
if use_hsdp:
|
||||
group_getters["get_fs_group"].assert_called_once_with()
|
||||
if hsdp_replicate_size > 1:
|
||||
group_getters["get_hsdp_replicate_group"].assert_called_once_with()
|
||||
else:
|
||||
group_getters["get_hsdp_replicate_group"].assert_not_called()
|
||||
else:
|
||||
group_getters["get_fs_group"].assert_not_called()
|
||||
group_getters["get_hsdp_replicate_group"].assert_not_called()
|
||||
if enable_ep:
|
||||
group_getters["get_ep_group"].assert_called_once_with()
|
||||
else:
|
||||
group_getters["get_ep_group"].assert_not_called()
|
||||
|
||||
|
||||
def test_missing_setproctitle_is_non_fatal(mocker: MockerFixture) -> None:
|
||||
mocker.patch.object(
|
||||
worker_module,
|
||||
"model_parallel_is_initialized",
|
||||
return_value=False,
|
||||
)
|
||||
mocker.patch.dict(sys.modules, {"setproctitle": None})
|
||||
decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
|
||||
|
||||
worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=False, use_hsdp=False)
|
||||
|
||||
decorate_logs.assert_called_once_with("DiffusionWorker")
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.platform != "linux" or shutil.which("ps") is None or importlib.util.find_spec("setproctitle") is None,
|
||||
reason="requires Linux ps and setproctitle",
|
||||
)
|
||||
def test_process_title_is_visible_through_ps() -> None:
|
||||
expected_name = "vLLM-Omni::DiffusionWorker_TP1"
|
||||
child_script = textwrap.dedent(
|
||||
"""
|
||||
import time
|
||||
|
||||
from vllm.utils.system_utils import set_process_title
|
||||
|
||||
|
||||
set_process_title(
|
||||
name="DiffusionWorker_TP1",
|
||||
prefix="vLLM-Omni",
|
||||
)
|
||||
time.sleep(30)
|
||||
"""
|
||||
)
|
||||
process = subprocess.Popen(
|
||||
[sys.executable, "-c", child_script],
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
)
|
||||
|
||||
observed_title = ""
|
||||
try:
|
||||
for _ in range(100):
|
||||
if process.poll() is not None:
|
||||
stderr = process.stderr.read() if process.stderr is not None else ""
|
||||
pytest.fail(f"child process exited before title check: {stderr}")
|
||||
result = subprocess.run(
|
||||
["ps", "-o", "args=", "-p", str(process.pid)],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
observed_title = result.stdout.strip()
|
||||
if expected_name in observed_title:
|
||||
break
|
||||
time.sleep(0.1)
|
||||
else:
|
||||
pytest.fail(f"expected {expected_name!r} in process title, got {observed_title!r}")
|
||||
finally:
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
process.wait(timeout=5)
|
||||
@@ -52,6 +52,11 @@ def client(mock_engine):
|
||||
c.shutdown()
|
||||
|
||||
|
||||
def test_inline_client_implements_shared_prompt_hook_protocol(client):
|
||||
assert client.prompt_transform_func is None
|
||||
assert client.prompt_expand_func is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inline_dispatch_request_success(client, mock_engine):
|
||||
mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()])
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""E2E accuracy guards against pinned official Lightricks LTX references.
|
||||
|
||||
@@ -43,7 +43,7 @@ LTX25_OFFICIAL_MODEL_ID = "Lightricks/LTX-2.5"
|
||||
LTX25_OFFICIAL_MODEL_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1"
|
||||
LTX25_OFFICIAL_COMMON_FILES = {
|
||||
"text_encoder": "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors",
|
||||
"video_vae": "vae/ltx-2.5-video-vae-conv-bf16.safetensors",
|
||||
"video_vae": "vae/ltx-2.5-video-vae-bf16.safetensors",
|
||||
"audio_vae": "vae/ltx-2.5-audio-vae-bf16.safetensors",
|
||||
}
|
||||
LTX25_OFFICIAL_VARIANT_FILES = {
|
||||
@@ -208,6 +208,9 @@ def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") ->
|
||||
model = Path(configured_model).resolve()
|
||||
snapshot_revision = model.name if model.parent.name == "snapshots" else None
|
||||
assert (model / transformer_subfolder).is_dir(), f"LTX-2.5 component not found: {model / transformer_subfolder}"
|
||||
assert (model / "diffusion_decoder" / "config.json").is_file(), (
|
||||
f"LTX-2.5 DiffVAE config not found: {model / 'diffusion_decoder' / 'config.json'}"
|
||||
)
|
||||
return model, configured_model, configured_revision or snapshot_revision
|
||||
model_id = configured_model or LTX25_OMNI_MODEL_ID
|
||||
revision = configured_revision
|
||||
@@ -223,6 +226,7 @@ def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") ->
|
||||
"connectors/config.json",
|
||||
"connectors/diffusion_pytorch_model.safetensors.index.json",
|
||||
"connectors/diffusion_pytorch_model-*.safetensors",
|
||||
"diffusion_decoder/config.json",
|
||||
"scheduler/*",
|
||||
"text_encoder/config.json",
|
||||
"text_encoder/generation_config.json",
|
||||
@@ -470,7 +474,7 @@ def _assert_strict_similarity(
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
@pytest.mark.parametrize("task", ["t2v", "i2v"])
|
||||
def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path, task: str) -> None:
|
||||
"""Compare Omni with the official runtime using the same checkpoint connector weights."""
|
||||
"""Compare Omni with official canonical DiffVAE and matching connector weights."""
|
||||
_require_ltx25_model_access()
|
||||
artifact_parent = accuracy_artifact_root / "ltx_official"
|
||||
output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_distilled_{task}")
|
||||
@@ -545,6 +549,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
|
||||
omni_metadata = json.loads((omni_output / "metadata.json").read_text())
|
||||
assert official_metadata["official_revision"] == official_revision
|
||||
assert official_metadata["pipeline"] == "DistilledPipeline"
|
||||
assert official_metadata["video_vae_path"] == str(official_artifacts["video_vae"])
|
||||
assert official_metadata["attention_backend"] == ATTENTION_BACKEND
|
||||
assert official_metadata["connector_model"] == str(omni_model)
|
||||
assert omni_metadata["model_class_name"] == "LTX2DistilledTwoStagePipeline"
|
||||
@@ -594,7 +599,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
|
||||
@pytest.mark.parametrize("task", ["t2v", "i2v"])
|
||||
@pytest.mark.parametrize("pipeline_mode", ["full_one_stage", "full_two_stage"])
|
||||
def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pipeline_mode: str) -> None:
|
||||
"""Compare raw official Full/SFT weights with Omni's converted Full pipeline."""
|
||||
"""Compare raw official Full/SFT and DiffVAE weights with Omni."""
|
||||
_require_ltx25_model_access()
|
||||
artifact_parent = accuracy_artifact_root / "ltx_official"
|
||||
output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_{pipeline_mode}_{task}")
|
||||
@@ -697,6 +702,7 @@ def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pi
|
||||
"full_two_stage": "TI2VidTwoStagesPipeline",
|
||||
}[pipeline_mode]
|
||||
)
|
||||
assert official_metadata["video_vae_path"] == str(official_artifacts["video_vae"])
|
||||
assert official_metadata["attention_backend"] == ATTENTION_BACKEND
|
||||
assert official_metadata["connector_model"] == str(omni_model)
|
||||
assert omni_metadata["model_class_name"] == model_class_name
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -119,6 +119,9 @@ def _model_name(model_env_var: str = MODEL_ENV_VAR, component: str = "FL2VA") ->
|
||||
def _server_args() -> list[str]:
|
||||
return [
|
||||
"--trust-remote-code",
|
||||
# Both official-reference requests use fixed shapes. Keep regional
|
||||
# compilation static so cold Inductor codegen is stable across CI nodes.
|
||||
"--no-diffusion-compile-dynamic",
|
||||
"--num-gpus",
|
||||
"4",
|
||||
"--usp",
|
||||
|
||||
@@ -537,7 +537,7 @@ def api_server(server_case: ServerCase, mock_async_omni, tmp_path):
|
||||
cmd = OmniServeCommand()
|
||||
cmd.subparser_init(subparsers)
|
||||
|
||||
port = get_open_port()
|
||||
port = get_open_port(host="0.0.0.0")
|
||||
args = parser.parse_args(["serve", server_case.served_model, "--omni", "--port", str(port)])
|
||||
|
||||
def run_server():
|
||||
@@ -814,7 +814,7 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
|
||||
ServerCase(
|
||||
served_model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
stage_list=["diffusion"],
|
||||
stage_configs=[{"stage_type": "diffusion"}],
|
||||
stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
|
||||
outputs=[_build_diffusion_video_output()],
|
||||
),
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
@@ -825,7 +825,7 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
|
||||
ServerCase(
|
||||
served_model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
stage_list=["diffusion"],
|
||||
stage_configs=[{"stage_type": "diffusion"}],
|
||||
stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
|
||||
outputs=[_build_diffusion_video_output()],
|
||||
),
|
||||
"Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
@@ -886,7 +886,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b
|
||||
ServerCase(
|
||||
served_model="MiniMaxAI/MiniMax-H3",
|
||||
stage_list=["diffusion"],
|
||||
stage_configs=[{"stage_type": "diffusion"}],
|
||||
stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
|
||||
outputs=[_build_diffusion_video_output()],
|
||||
),
|
||||
SamplingCase(kind=SamplingKind.VIDEO_REF2VA_IMAGE_AUDIO, sampling_params=None),
|
||||
@@ -897,7 +897,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b
|
||||
ServerCase(
|
||||
served_model="MiniMaxAI/MiniMax-H3",
|
||||
stage_list=["diffusion"],
|
||||
stage_configs=[{"stage_type": "diffusion"}],
|
||||
stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
|
||||
outputs=[_build_diffusion_video_output()],
|
||||
),
|
||||
SamplingCase(kind=SamplingKind.VIDEO_REF2VA_MULTI_VIDEO, sampling_params=None),
|
||||
|
||||
@@ -8,6 +8,7 @@ from dataclasses import FrozenInstanceError
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.stage_config import get_deploy_duplex_max_sessions
|
||||
from vllm_omni.config.stage_config import (
|
||||
DuplexSessionRuntimeConfig,
|
||||
load_deploy_config,
|
||||
@@ -66,3 +67,18 @@ def test_duplex_session_runtime_rejects_non_positive_values(tmp_path, name: str,
|
||||
|
||||
with pytest.raises(ValueError, match=rf"duplex_session\.{name} must be positive"):
|
||||
load_deploy_config(deploy_path)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("deploy_yaml", "expected"),
|
||||
[
|
||||
("minicpmo_4_5.yaml", 4),
|
||||
("minicpmo_4_5_8x4090.yaml", 1),
|
||||
("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4),
|
||||
],
|
||||
)
|
||||
def test_deploy_duplex_max_sessions_tracks_the_deploy_config(deploy_yaml: str, expected: int) -> None:
|
||||
# Guards the admission probe. A capacity edit, a config that declares none,
|
||||
# or an overlay inheriting one from its base must surface here rather than
|
||||
# as another nightly duplex admission timeout.
|
||||
assert get_deploy_duplex_max_sessions(deploy_yaml) == expected
|
||||
@@ -18,6 +18,10 @@ from vllm_omni.experimental.fullduplex.client import (
|
||||
from vllm_omni.experimental.fullduplex.minicpmo45.policy import (
|
||||
MiniCPMO45DuplexPolicy,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.video_stacking import (
|
||||
concat_frames,
|
||||
unit_subframe_offsets,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -418,3 +422,62 @@ def test_realtime_client_pcm16_wav_round_trip(tmp_path):
|
||||
assert wav_file.getnchannels() == 1
|
||||
assert wav_file.getframerate() == 16_000
|
||||
assert read_pcm16_wav(path) == pcm16
|
||||
|
||||
|
||||
def _stack_test_image(color, size=(40, 30)):
|
||||
from PIL import Image
|
||||
|
||||
return Image.new("RGB", size, color)
|
||||
|
||||
|
||||
def test_video_stacking_tiles_a_units_subframes_into_one_image():
|
||||
frames = [_stack_test_image((255, 0, 0)), _stack_test_image((0, 255, 0))]
|
||||
|
||||
composite = concat_frames(frames)
|
||||
|
||||
# Landscape cells stack into a column because that canvas is squarer, and
|
||||
# the interior seam carries the separator band that lets the model tell the
|
||||
# sub-frames apart.
|
||||
assert composite.size == (40, 30 * 2 + 6)
|
||||
assert composite.getpixel((20, 30 + 3)) == (0, 0, 0)
|
||||
|
||||
|
||||
def test_video_stacking_picks_the_squarest_grid():
|
||||
portrait = [_stack_test_image((255, 0, 0), size=(20, 60)) for _ in range(2)]
|
||||
four = [_stack_test_image((0, 0, 255)) for _ in range(4)]
|
||||
|
||||
# A portrait clip tiles side by side, four frames always tile 2x2.
|
||||
assert concat_frames(portrait).size == (20 * 2 + 6, 60)
|
||||
assert concat_frames(four).size == (40 * 2 + 6, 30 * 2 + 6)
|
||||
|
||||
|
||||
def test_video_stacking_skips_the_subframe_that_duplicates_the_base_frame():
|
||||
# Official samples stack_frames=5 as 0.2/0.4/0.6/0.8 s: offset 0 is the base
|
||||
# frame, already sent as frame_list[0].
|
||||
assert unit_subframe_offsets(5) == pytest.approx([0.2, 0.4, 0.6, 0.8])
|
||||
assert unit_subframe_offsets(1) == []
|
||||
|
||||
|
||||
def test_realtime_client_sends_each_units_composite_beside_its_base_frame():
|
||||
sent: list[dict] = []
|
||||
|
||||
class _Client(RealtimeDuplexClient):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
async def send(self, event):
|
||||
sent.append(event)
|
||||
|
||||
asyncio.run(
|
||||
_Client().stream_pcm16(
|
||||
b"\x01\x00" * (16_000 * 3),
|
||||
chunk_ms=200,
|
||||
realtime=False,
|
||||
video_frames=["f0", "f1"],
|
||||
stacked_video_frames=["s0", None],
|
||||
)
|
||||
)
|
||||
|
||||
# A composite belongs to the unit it was captured in, so it rides the same
|
||||
# append as that unit's base frame; a unit without one sends the base alone.
|
||||
assert [event["video_frames"] for event in sent if "video_frames" in event] == [["f0", "s0"], ["f1"]]
|
||||
@@ -0,0 +1,194 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import threading
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.benchmarks.duplex import omni_duplex_eval_runner as runner
|
||||
from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import DuplexSample
|
||||
from vllm_omni.benchmarks.duplex.omni_duplex_eval_judge import DuplexJudge
|
||||
from vllm_omni.entrypoints.cli.benchmark import omni_duplex_eval as cli
|
||||
from vllm_omni.entrypoints.cli.benchmark.omni_duplex_eval import OmniDuplexEvalSubcommand
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
|
||||
|
||||
|
||||
def test_cli_generate_evaluate_summarize_flow(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys):
|
||||
manifest = tmp_path / "manifest.json"
|
||||
manifest.write_text(
|
||||
json.dumps(
|
||||
[
|
||||
{
|
||||
"id": "sample-1",
|
||||
"split": "PR_correction",
|
||||
"question_text": "What changed?",
|
||||
"answer1": "The object moved.",
|
||||
}
|
||||
]
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
response_root = tmp_path / "responses"
|
||||
score_root = tmp_path / "scores"
|
||||
|
||||
async def fake_generate(sample, *, output_root, **kwargs):
|
||||
output = Path(output_root) / sample.split / f"{sample.id}.json"
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.write_text(json.dumps([{"sentence": "It moved.", "start": 0, "end": 1}]), encoding="utf-8")
|
||||
output.with_name(output.stem + ".meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8")
|
||||
return output
|
||||
|
||||
class FakeJudge:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def chat(self, *args, **kwargs):
|
||||
return '{"success_score": 1, "is_relevant": 1}'
|
||||
|
||||
monkeypatch.setattr(cli, "generate_sample", fake_generate)
|
||||
monkeypatch.setattr(cli, "DuplexJudge", FakeJudge)
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
OmniDuplexEvalSubcommand.add_cli_args(parser)
|
||||
|
||||
common = ["--dataset", str(manifest), "--family", "pr"]
|
||||
generate = parser.parse_args(
|
||||
[
|
||||
"generate",
|
||||
*common,
|
||||
"--model",
|
||||
"mock",
|
||||
"--ref-audio",
|
||||
str(manifest),
|
||||
"--response-root",
|
||||
str(response_root),
|
||||
"--concurrency",
|
||||
"2",
|
||||
]
|
||||
)
|
||||
OmniDuplexEvalSubcommand.cmd(generate)
|
||||
evaluate = parser.parse_args(
|
||||
[
|
||||
"evaluate",
|
||||
*common,
|
||||
"--response-root",
|
||||
str(response_root),
|
||||
"--score-root",
|
||||
str(score_root),
|
||||
"--judge-model",
|
||||
"mock-judge",
|
||||
"--eval-workers",
|
||||
"2",
|
||||
]
|
||||
)
|
||||
OmniDuplexEvalSubcommand.cmd(evaluate)
|
||||
OmniDuplexEvalSubcommand.cmd(parser.parse_args(["summarize", "--score-root", str(score_root)]))
|
||||
|
||||
summary = json.loads(capsys.readouterr().out)
|
||||
assert summary["samples"] == 1
|
||||
assert summary["pr"]["mean_all_success"] == 1.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_exercises_realtime_socket_and_media_clock(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
import websockets
|
||||
|
||||
received = []
|
||||
|
||||
async def handler(websocket):
|
||||
response_started = False
|
||||
async for raw in websocket:
|
||||
event = json.loads(raw)
|
||||
received.append(event)
|
||||
if event["type"] == "session.update":
|
||||
await websocket.send(json.dumps({"type": "session.created"}))
|
||||
elif event["type"] == "input_audio_buffer.append" and not response_started:
|
||||
response_started = True
|
||||
await websocket.send(json.dumps({"type": "response.created", "response": {"id": "r1"}}))
|
||||
await websocket.send(
|
||||
json.dumps({"type": "response.output_text.delta", "response_id": "r1", "delta": "Done."})
|
||||
)
|
||||
await websocket.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "response.audio.delta",
|
||||
"response_id": "r1",
|
||||
"delta": base64.b64encode(b"\0\0" * 240).decode(),
|
||||
"sample_rate_hz": 24_000,
|
||||
}
|
||||
)
|
||||
)
|
||||
elif event["type"] == "input_audio_buffer.commit":
|
||||
await websocket.send(json.dumps({"type": "response.done", "response": {"id": "r1"}}))
|
||||
elif event["type"] == "session.close":
|
||||
await websocket.send(json.dumps({"type": "session.closed"}))
|
||||
return
|
||||
|
||||
monkeypatch.setattr(runner, "read_audio_pcm16", lambda path: b"\0\0" * (16_000 * 8 // 10))
|
||||
monkeypatch.setattr(runner, "video_duration", lambda path: 0.8)
|
||||
monkeypatch.setattr(runner, "iter_jpegs", lambda *args, **kwargs: iter([(0.0, b"jpeg")]))
|
||||
audio = tmp_path / "question.wav"
|
||||
video = tmp_path / "video.mp4"
|
||||
ref = tmp_path / "ref.wav"
|
||||
for path in (audio, video, ref):
|
||||
path.write_bytes(b"data")
|
||||
sample = DuplexSample("sample", "PR_correction", "pr", "correction", video, audio)
|
||||
|
||||
async with websockets.serve(handler, "127.0.0.1", 0) as server:
|
||||
port = server.sockets[0].getsockname()[1]
|
||||
output = await runner.generate_sample(
|
||||
sample,
|
||||
url=f"ws://127.0.0.1:{port}/v1/realtime?duplex=1",
|
||||
model="mock",
|
||||
ref_audio=ref,
|
||||
output_root=tmp_path / "responses",
|
||||
)
|
||||
|
||||
assert json.loads(output.read_text(encoding="utf-8")) == [{"sentence": "Done.", "start": 0.8, "end": 0.8}]
|
||||
meta = json.loads(output.with_name("sample.meta.json").read_text(encoding="utf-8"))
|
||||
assert meta["response_done"] is True
|
||||
assert meta["drain_timeout"] is None
|
||||
assert any(event["type"] == "playback.ack" for event in received)
|
||||
|
||||
|
||||
def test_judge_exercises_openai_http_schema():
|
||||
requests = []
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def do_POST(self):
|
||||
length = int(self.headers["Content-Length"])
|
||||
requests.append((self.path, self.headers["Authorization"], json.loads(self.rfile.read(length))))
|
||||
payload = json.dumps({"choices": [{"message": {"content": '{"success_score": 1}'}}]}).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(payload)))
|
||||
self.end_headers()
|
||||
self.wfile.write(payload)
|
||||
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
judge = DuplexJudge(f"http://127.0.0.1:{server.server_port}", "judge", api_key="token")
|
||||
assert judge.chat("prompt") == '{"success_score": 1}'
|
||||
finally:
|
||||
server.shutdown()
|
||||
thread.join()
|
||||
server.server_close()
|
||||
|
||||
path, authorization, payload = requests[0]
|
||||
assert path == "/v1/chat/completions"
|
||||
assert authorization == "Bearer token"
|
||||
assert payload["model"] == "judge"
|
||||
assert payload["messages"] == [{"role": "user", "content": "prompt"}]
|
||||
@@ -0,0 +1,204 @@
|
||||
"""
|
||||
End-to-end test for MammothModa2 text-to-image generation.
|
||||
|
||||
Verifies that the AR->DiT pipeline produces an image tensor whose pixel values
|
||||
match a golden reference.
|
||||
|
||||
Model Hub repo id: ``bytedance-research/MammothModa2-Preview``.
|
||||
Deploy config: ``get_deploy_config_path("mammoth_moda2.yaml")`` -> ``vllm_omni/deploy/mammoth_moda2.yaml``
|
||||
|
||||
Golden pixel file: ``tests/e2e/offline_inference/fixtures/mammoth_moda2_t2i_golden.json``
|
||||
Regenerate with: ``UPDATE_GOLDEN=1 pytest tests/e2e/offline_inference/test_mammoth_moda2_expansion.py``
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from huggingface_hub import snapshot_download
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Constants
|
||||
# ---------------------------------------------------------------------------
|
||||
_IMAGE_TOKEN_ID = 151655 # "<|image_pad|>"
|
||||
_VIDEO_TOKEN_ID = 151656 # "<|video_pad|>"
|
||||
_VISION_START_TOKEN_ID = 151652 # "<|vision_start|>"
|
||||
_VISION_END_TOKEN_ID = 151653 # "<|vision_end|>"
|
||||
_AR_PATCH_SIZE = 16
|
||||
|
||||
MODEL_PATH = "bytedance-research/MammothModa2-Preview"
|
||||
T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
|
||||
|
||||
_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG)
|
||||
|
||||
# Golden pixel reference file. Set UPDATE_GOLDEN=1 to regenerate.
|
||||
_GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"
|
||||
# Fixed sampling coordinates: (channel, row_fraction, col_fraction)
|
||||
# Covers corners, centre, and mid-edges across all 3 channels.
|
||||
_PIXEL_SAMPLE_COORDS = [
|
||||
(0, 0.0, 0.0),
|
||||
(0, 0.5, 0.5),
|
||||
(0, 1.0, 1.0),
|
||||
(0, 0.25, 0.75),
|
||||
(1, 0.0, 1.0),
|
||||
(1, 0.5, 0.0),
|
||||
(1, 0.75, 0.25),
|
||||
(1, 1.0, 0.5),
|
||||
(2, 0.0, 0.5),
|
||||
(2, 0.5, 1.0),
|
||||
(2, 0.75, 0.75),
|
||||
(2, 1.0, 0.0),
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
def _load_t2i_gen_config(repo_id: str) -> dict:
|
||||
weights_dir = Path(snapshot_download(repo_id))
|
||||
cfg_path = weights_dir / "t2i_generation_config.json"
|
||||
if not cfg_path.exists():
|
||||
pytest.skip(f"t2i_generation_config.json not found at {cfg_path}")
|
||||
with cfg_path.open() as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def _format_t2i_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str:
|
||||
return (
|
||||
"<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n"
|
||||
f"<|im_start|>user\n{user_prompt}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
f"<|image start|>{ar_width}*{ar_height}<|image token|>"
|
||||
)
|
||||
|
||||
|
||||
def _sample_pixels(img_tensor: torch.Tensor) -> list[float]:
|
||||
"""Sample pixel values at fixed fractional coordinates from a (C, H, W) tensor."""
|
||||
t = img_tensor.float().clamp(0.0, 1.0)
|
||||
if t.ndim == 4:
|
||||
t = t[0] # unbatch
|
||||
C, H, W = t.shape
|
||||
values = []
|
||||
for c, rh, rw in _PIXEL_SAMPLE_COORDS:
|
||||
ri = min(int(rh * (H - 1)), H - 1)
|
||||
ci = min(int(rw * (W - 1)), W - 1)
|
||||
values.append(round(float(t[c, ri, ci]), 6))
|
||||
return values
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# End-to-end test
|
||||
# ---------------------------------------------------------------------------
|
||||
pytestmark = [
|
||||
pytest.mark.slow,
|
||||
pytest.mark.diffusion,
|
||||
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3201")
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
def test_mammothmoda2_t2i_e2e(omni_runner: OmniRunner):
|
||||
"""
|
||||
End-to-end text-to-image generation with MammothModa2 (AR -> DiT).
|
||||
|
||||
Verifies:
|
||||
- Omni pipeline initialises with the two-stage YAML config.
|
||||
- DiT stage outputs an image tensor with the correct shape.
|
||||
- A fixed set of pixel values matches a golden reference
|
||||
(regenerate with ``UPDATE_GOLDEN=1``).
|
||||
"""
|
||||
gen_cfg = _load_t2i_gen_config(MODEL_PATH)
|
||||
eol_token_id = int(gen_cfg["eol_token_id"])
|
||||
visual_start = int(gen_cfg["visual_token_start_id"])
|
||||
visual_end = int(gen_cfg["visual_token_end_id"])
|
||||
|
||||
height, width = 256, 256 # small for CI speed
|
||||
ar_height, ar_width = height // _AR_PATCH_SIZE, width // _AR_PATCH_SIZE
|
||||
expected_grid_tokens = ar_height * (ar_width + 1)
|
||||
|
||||
prompt_text = "A cat sitting on a laptop keyboard"
|
||||
formatted_prompt = _format_t2i_prompt(prompt_text, ar_width, ar_height)
|
||||
|
||||
omni = omni_runner.omni
|
||||
ar_sampling = SamplingParams(
|
||||
temperature=0.0,
|
||||
top_k=1,
|
||||
max_tokens=max(1, expected_grid_tokens + 1),
|
||||
detokenize=False,
|
||||
)
|
||||
dit_sampling = SamplingParams(temperature=0.0, max_tokens=1, detokenize=False)
|
||||
|
||||
outputs = list(
|
||||
omni.generate(
|
||||
[
|
||||
{
|
||||
"prompt": formatted_prompt,
|
||||
"additional_information": {
|
||||
"omni_task": ["t2i"],
|
||||
"ar_width": [ar_width],
|
||||
"ar_height": [ar_height],
|
||||
"eol_token_id": [eol_token_id],
|
||||
"visual_token_start_id": [visual_start],
|
||||
"visual_token_end_id": [visual_end],
|
||||
"image_height": [height],
|
||||
"image_width": [width],
|
||||
"num_inference_steps": [2],
|
||||
"text_guidance_scale": [1.0],
|
||||
"cfg_range": [0.0, 1.0],
|
||||
"visual_ids": [
|
||||
_IMAGE_TOKEN_ID,
|
||||
_VIDEO_TOKEN_ID,
|
||||
_VISION_START_TOKEN_ID,
|
||||
_VISION_END_TOKEN_ID,
|
||||
],
|
||||
},
|
||||
}
|
||||
],
|
||||
[ar_sampling, dit_sampling],
|
||||
)
|
||||
)
|
||||
|
||||
assert len(outputs) > 0, "Pipeline produced no outputs"
|
||||
|
||||
found_image = False
|
||||
for out in outputs:
|
||||
ro_list = out
|
||||
if not isinstance(ro_list, list):
|
||||
ro_list = [ro_list]
|
||||
for ro in ro_list:
|
||||
completion_outputs = getattr(ro, "outputs", None)
|
||||
if not isinstance(completion_outputs, list):
|
||||
continue
|
||||
for completion in completion_outputs:
|
||||
mm = getattr(completion, "multimodal_output", None)
|
||||
if not (isinstance(mm, dict) and "image" in mm):
|
||||
continue
|
||||
img_list = mm["image"] if isinstance(mm["image"], list) else [mm["image"]]
|
||||
for img_tensor in img_list:
|
||||
assert isinstance(img_tensor, torch.Tensor), f"Expected image tensor, got {type(img_tensor)}"
|
||||
assert img_tensor.ndim in (3, 4), f"Expected 3D or 4D image tensor, got {img_tensor.ndim}D"
|
||||
|
||||
sampled = _sample_pixels(img_tensor)
|
||||
|
||||
if os.environ.get("UPDATE_GOLDEN"):
|
||||
_GOLDEN_T2I_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
_GOLDEN_T2I_PATH.write_text(json.dumps({"pixels": sampled}, indent=2))
|
||||
print(f"\nGolden file written to {_GOLDEN_T2I_PATH}")
|
||||
elif _GOLDEN_T2I_PATH.exists():
|
||||
golden = json.loads(_GOLDEN_T2I_PATH.read_text())["pixels"]
|
||||
for i, (got, exp) in enumerate(zip(sampled, golden)):
|
||||
assert abs(got - exp) < 1e-4, f"Pixel {i} mismatch: got {got}, expected {exp}"
|
||||
|
||||
found_image = True
|
||||
|
||||
assert found_image, "No image tensor found in pipeline output"
|
||||
@@ -15,10 +15,11 @@ import pytest
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
from tests.helpers.runtime import OmniServerParams, get_model_prefix
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from tests.helpers.stage_config import get_deploy_config_path, get_deploy_duplex_max_sessions
|
||||
|
||||
MODEL = "openbmb/MiniCPM-o-4_5"
|
||||
DEPLOY_CONFIG = get_deploy_config_path("minicpmo_4_5.yaml")
|
||||
DEPLOY_CONFIG_REL = "minicpmo_4_5.yaml"
|
||||
DEPLOY_CONFIG = get_deploy_config_path(DEPLOY_CONFIG_REL)
|
||||
ASSET_DIR = Path(__file__).resolve().parents[3] / "assets" / "minicpmo_4_5"
|
||||
RESPONSE_REQUIRED_WAV = ASSET_DIR / "response_required_16k.wav"
|
||||
RESPONSE_REQUIRED_SHA256 = "2e5fd4eb3ee434ce107ee3a0591fa624a33f7683c7462f45fe651c443c9af941"
|
||||
@@ -39,6 +40,16 @@ SERVER_PARAMS = [
|
||||
]
|
||||
|
||||
|
||||
def deploy_max_sessions() -> int:
|
||||
"""Concurrent duplex sessions the deploy config under test admits.
|
||||
|
||||
The admission probe has to expect the capacity the server is actually
|
||||
started with. Hardcoding it silently drifts the moment the deploy config
|
||||
changes its capacity, turning that change into an unrelated probe timeout.
|
||||
"""
|
||||
return get_deploy_duplex_max_sessions(DEPLOY_CONFIG_REL)
|
||||
|
||||
|
||||
def validated_wav(path: Path, expected_sha256: str) -> Path:
|
||||
actual_sha256 = hashlib.sha256(path.read_bytes()).hexdigest()
|
||||
if actual_sha256 != expected_sha256:
|
||||
@@ -80,6 +91,19 @@ def resolve_ref_audio() -> Path:
|
||||
return ref_audio
|
||||
|
||||
|
||||
def duplex_camera_frames(*, seconds: int, cache_dir: Path) -> list[str]:
|
||||
"""Base64 JPEG frames of a moving synthetic scene, one per second.
|
||||
|
||||
Mirrors what the realtime web client captures (1 fps, no client-side
|
||||
resize: the server normalizes at ``scale_resolution=448``).
|
||||
"""
|
||||
from tests.e2e.online_serving.helpers.minicpmo_realtime_duplex_scenarios import _video_frames_from_file
|
||||
from tests.helpers.media import generate_synthetic_video
|
||||
|
||||
video = generate_synthetic_video(448, 448, 30 * seconds, cache_dir=cache_dir)
|
||||
return _video_frames_from_file(Path(video["file_path"]))
|
||||
|
||||
|
||||
def realtime_url(omni_server) -> str:
|
||||
return f"ws://{omni_server.host}:{omni_server.port}/v1/realtime?duplex=1"
|
||||
|
||||
@@ -102,6 +126,8 @@ def demo_args(
|
||||
output_audio_format="pcm16",
|
||||
chunk_ms=200,
|
||||
realtime_input=True,
|
||||
input_video=None,
|
||||
stack_frames=1,
|
||||
first_turn_ms=1400,
|
||||
turn_duration_ms=[],
|
||||
first_turn_transcript="duplex CI speech",
|
||||
|
||||
@@ -37,8 +37,10 @@ import json
|
||||
import sys
|
||||
import time
|
||||
import wave
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
if str(_REPO_ROOT) not in sys.path:
|
||||
@@ -55,15 +57,100 @@ from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
|
||||
PCM16_SAMPLE_RATE,
|
||||
RealtimeEventCollector,
|
||||
build_realtime_url,
|
||||
duplex_unit_boundary_ms,
|
||||
read_pcm16_wav,
|
||||
summarize_session_request_metrics,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
|
||||
reference_audio_data_url as _ref_audio_data_url,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.video_stacking import ( # noqa: E402
|
||||
concat_frames_b64,
|
||||
unit_subframe_offsets,
|
||||
)
|
||||
|
||||
_url_with_model = build_realtime_url
|
||||
_read_wav_pcm16 = read_pcm16_wav
|
||||
DemoArgs = argparse.Namespace | SimpleNamespace
|
||||
|
||||
|
||||
def _video_frames_from_file(
|
||||
path: Path,
|
||||
*,
|
||||
fps: float = 1.0,
|
||||
jpeg_quality: int = 70,
|
||||
stack_frames: int = 1,
|
||||
) -> tuple[list[str], list[str | None]]:
|
||||
"""Decode ``path`` into base64 JPEG frames at the omni duplex cadence.
|
||||
|
||||
One frame per second matches what the realtime web client captures, so the
|
||||
stream stays representative of a live camera. ``stack_frames > 1`` also
|
||||
samples each unit's interior sub-frames and tiles them into one composite,
|
||||
the official way to raise visual refresh rate without changing the audio
|
||||
cadence. Returns ``(base_frames, stacked_frames)`` with the tracks parallel.
|
||||
"""
|
||||
import cv2
|
||||
|
||||
offsets = unit_subframe_offsets(stack_frames)
|
||||
capture = cv2.VideoCapture(str(path))
|
||||
if not capture.isOpened():
|
||||
raise ValueError(f"cannot open video: {path}")
|
||||
|
||||
def encode(frame_bgr, index: int) -> str:
|
||||
ok, jpeg = cv2.imencode(".jpg", frame_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality])
|
||||
if not ok:
|
||||
raise ValueError(f"cannot JPEG-encode frame {index} of {path}")
|
||||
return base64.b64encode(jpeg.tobytes()).decode("ascii")
|
||||
|
||||
try:
|
||||
source_fps = capture.get(cv2.CAP_PROP_FPS) or 0.0
|
||||
stride = max(1, round(source_fps / fps)) if source_fps > 0 else 1
|
||||
# Sub-frame k of a unit sits ``offset * stride`` frames past its base.
|
||||
sub_strides = {max(1, round(offset * stride)): position for position, offset in enumerate(offsets)}
|
||||
frames: list[str] = []
|
||||
interiors: list[list[str]] = []
|
||||
index = 0
|
||||
while True:
|
||||
ok, frame_bgr = capture.read()
|
||||
if not ok:
|
||||
break
|
||||
position, within_unit = divmod(index, stride)
|
||||
if within_unit == 0:
|
||||
frames.append(encode(frame_bgr, index))
|
||||
interiors.append([])
|
||||
elif within_unit in sub_strides and position < len(interiors):
|
||||
interiors[position].append(encode(frame_bgr, index))
|
||||
index += 1
|
||||
finally:
|
||||
capture.release()
|
||||
if not frames:
|
||||
raise ValueError(f"video has no decodable frames: {path}")
|
||||
stacked = [concat_frames_b64(interior) if interior else None for interior in interiors]
|
||||
return frames, stacked
|
||||
|
||||
|
||||
def _resolve_video_frames(args: DemoArgs) -> tuple[list[str], list[str | None]]:
|
||||
"""Base64 JPEG frames for the duplex camera track, newest last.
|
||||
|
||||
Returns the base track plus the parallel stacked composites. Only a video
|
||||
source has sub-frames to stack; an explicit frame list or a still image
|
||||
stacks nothing.
|
||||
"""
|
||||
frames = list(getattr(args, "video_frames_b64", None) or [])
|
||||
if frames:
|
||||
stacked = list(getattr(args, "video_stacked_frames_b64", None) or [])
|
||||
stacked += [None] * (len(frames) - len(stacked))
|
||||
return frames, stacked[: len(frames)]
|
||||
input_video = getattr(args, "input_video", None)
|
||||
if input_video:
|
||||
return _video_frames_from_file(
|
||||
Path(input_video).expanduser(),
|
||||
stack_frames=max(1, int(getattr(args, "stack_frames", 1) or 1)),
|
||||
)
|
||||
frame_image = getattr(args, "frame_image", None)
|
||||
if frame_image:
|
||||
return [base64.b64encode(Path(frame_image).expanduser().read_bytes()).decode("ascii")], [None]
|
||||
return [], []
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -373,17 +460,20 @@ class DemoState:
|
||||
commit_indices = [
|
||||
index for index, event in enumerate(self.events) if event.get("type") == "input_audio_buffer.committed"
|
||||
]
|
||||
decision_indices = [
|
||||
index
|
||||
for index, event in enumerate(self.events)
|
||||
if event.get("type") == "response.created"
|
||||
or (
|
||||
event.get("type") == "response.listen"
|
||||
and isinstance(event.get("response"), dict)
|
||||
and isinstance(event["response"].get("metadata"), dict)
|
||||
and event["response"]["metadata"].get("model_listen") is True
|
||||
)
|
||||
]
|
||||
decision_indices: list[int] = []
|
||||
for index, event in enumerate(self.events):
|
||||
event_type = event.get("type")
|
||||
if event_type == "response.created":
|
||||
decision_indices.append(index)
|
||||
continue
|
||||
if event_type != "response.listen":
|
||||
continue
|
||||
response = event.get("response")
|
||||
if not isinstance(response, dict):
|
||||
continue
|
||||
metadata = response.get("metadata")
|
||||
if isinstance(metadata, dict) and metadata.get("model_listen") is True:
|
||||
decision_indices.append(index)
|
||||
first_input_index = self.first_index("input_audio_buffer.speech_started")
|
||||
if first_input_index is None:
|
||||
first_input_index = commit_indices[0] if commit_indices else None
|
||||
@@ -500,7 +590,7 @@ class DemoState:
|
||||
return stale
|
||||
|
||||
|
||||
def _session_update_event(args: argparse.Namespace) -> dict[str, object]:
|
||||
def _session_update_event(args: DemoArgs) -> dict[str, object]:
|
||||
session_payload: dict[str, object] = {
|
||||
"model": args.model,
|
||||
"modalities": ["audio", "text"],
|
||||
@@ -836,7 +926,8 @@ async def _send_pcm16(
|
||||
hints: dict[str, object] | None = None,
|
||||
first_chunk_hints: dict[str, object] | None = None,
|
||||
on_model_unit_ready=None,
|
||||
frame_b64: str | None = None,
|
||||
frames_b64: Sequence[str] | None = None,
|
||||
stacked_frames_b64: Sequence[str | None] | None = None,
|
||||
) -> None:
|
||||
hints = hints or {}
|
||||
first_chunk_hints = first_chunk_hints or {}
|
||||
@@ -844,6 +935,8 @@ async def _send_pcm16(
|
||||
model_unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
|
||||
audio_ms = 0
|
||||
next_model_unit_bytes = model_unit_bytes
|
||||
frames = list(frames_b64 or [])
|
||||
stacked = list(stacked_frames_b64 or [])
|
||||
frames_sent = 0
|
||||
for offset in range(0, len(pcm16), chunk_bytes):
|
||||
chunk = pcm16[offset : offset + chunk_bytes]
|
||||
@@ -852,9 +945,18 @@ async def _send_pcm16(
|
||||
chunk_hints = dict(hints)
|
||||
if offset == 0:
|
||||
chunk_hints.update(first_chunk_hints)
|
||||
if frame_b64 is not None and audio_ms > frames_sent * 1000:
|
||||
# Omni duplex cadence: one camera frame per 1 s of audio.
|
||||
chunk_hints["video_frames"] = [frame_b64]
|
||||
if frames and audio_ms >= duplex_unit_boundary_ms(frames_sent):
|
||||
# Omni duplex cadence: one camera frame per model unit, riding the
|
||||
# append that closes it, so the frame and the second of audio it was
|
||||
# captured during enter the same unit. A video advances one frame per
|
||||
# unit and holds its last frame when the audio outlives the clip; a
|
||||
# still image is a one-element list and therefore repeats.
|
||||
# A stacked composite of that unit's interior sub-frames rides the
|
||||
# same append right after the base frame (wire maximum of 2).
|
||||
index = min(frames_sent, len(frames) - 1)
|
||||
composite = stacked[index] if index < len(stacked) else None
|
||||
base = frames[index]
|
||||
chunk_hints["video_frames"] = [base] if composite is None else [base, composite]
|
||||
frames_sent += 1
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
@@ -935,7 +1037,8 @@ async def _send_clean_turn(
|
||||
realtime_input: bool = False,
|
||||
model_policy_settle_s: float = 2.0,
|
||||
commit_input: bool = True,
|
||||
frame_b64: str | None = None,
|
||||
frames_b64: Sequence[str] | None = None,
|
||||
stacked_frames_b64: Sequence[str | None] | None = None,
|
||||
) -> tuple[str | None, str]:
|
||||
before_created = state.count("response.created")
|
||||
before_model_listen = state.model_listen_count
|
||||
@@ -957,7 +1060,8 @@ async def _send_clean_turn(
|
||||
chunk_ms=chunk_ms,
|
||||
realtime_delay=realtime_input,
|
||||
hints={"transcript": transcript} if send_transcript_hint else {},
|
||||
frame_b64=frame_b64,
|
||||
frames_b64=frames_b64,
|
||||
stacked_frames_b64=stacked_frames_b64,
|
||||
)
|
||||
if commit_input:
|
||||
state.input_commit_sent_at_s.append(time.monotonic())
|
||||
@@ -1163,7 +1267,8 @@ async def _send_listen_only_overlap_pair(
|
||||
model_policy_settle_s: float,
|
||||
validation_mode: str,
|
||||
silence_ms: int,
|
||||
frame_b64: str | None = None,
|
||||
frames_b64: Sequence[str] | None = None,
|
||||
stacked_frames_b64: Sequence[str | None] | None = None,
|
||||
) -> tuple[list[str | None], list[str], bool]:
|
||||
if not realtime_input:
|
||||
raise ValueError("listen-only-overlap requires --realtime-input")
|
||||
@@ -1194,7 +1299,8 @@ async def _send_listen_only_overlap_pair(
|
||||
realtime_delay=True,
|
||||
hints={"transcript": transcripts[0]} if send_transcript_hint else {},
|
||||
on_model_unit_ready=continue_until_first_speak,
|
||||
frame_b64=frame_b64,
|
||||
frames_b64=frames_b64,
|
||||
stacked_frames_b64=stacked_frames_b64,
|
||||
)
|
||||
state.input_commit_sent_at_s.append(time.monotonic())
|
||||
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
|
||||
@@ -1233,7 +1339,8 @@ async def _send_listen_only_overlap_pair(
|
||||
realtime_delay=True,
|
||||
hints={"transcript": transcripts[1]} if send_transcript_hint else {},
|
||||
on_model_unit_ready=record_model_unit_ready,
|
||||
frame_b64=frame_b64,
|
||||
frames_b64=frames_b64,
|
||||
stacked_frames_b64=stacked_frames_b64,
|
||||
)
|
||||
state.input_commit_sent_at_s.append(time.monotonic())
|
||||
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
|
||||
@@ -1297,10 +1404,8 @@ async def _send_listen_only_overlap_pair(
|
||||
return [first_response_id, second_response_id], ["speak", second_outcome], overlap_ok
|
||||
|
||||
|
||||
async def run_demo(args: argparse.Namespace) -> dict[str, object]:
|
||||
demo_frame_b64: str | None = None
|
||||
if getattr(args, "frame_image", None):
|
||||
demo_frame_b64 = base64.b64encode(Path(args.frame_image).read_bytes()).decode("ascii")
|
||||
async def run_demo(args: DemoArgs) -> dict[str, object]:
|
||||
demo_frames_b64, demo_stacked_b64 = _resolve_video_frames(args)
|
||||
turn_input_paths = _turn_input_paths(
|
||||
Path(args.input_wav),
|
||||
list(getattr(args, "turn_input_wav", []) or []),
|
||||
@@ -1371,7 +1476,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
|
||||
model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000),
|
||||
validation_mode=validation_mode,
|
||||
silence_ms=max(0, args.silence_ms),
|
||||
frame_b64=demo_frame_b64,
|
||||
frames_b64=demo_frames_b64,
|
||||
stacked_frames_b64=demo_stacked_b64,
|
||||
)
|
||||
for turn_index in range(2, len(turn_specs)):
|
||||
transcript, duration_ms = turn_specs[turn_index]
|
||||
@@ -1389,7 +1495,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
|
||||
realtime_input=realtime_input,
|
||||
model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000),
|
||||
commit_input=not continuous_input,
|
||||
frame_b64=demo_frame_b64,
|
||||
frames_b64=demo_frames_b64,
|
||||
stacked_frames_b64=demo_stacked_b64,
|
||||
)
|
||||
turn_response_ids.append(response_id)
|
||||
turn_outcomes.append(outcome)
|
||||
@@ -1409,7 +1516,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
|
||||
realtime_input=realtime_input,
|
||||
model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000),
|
||||
commit_input=not continuous_input,
|
||||
frame_b64=demo_frame_b64,
|
||||
frames_b64=demo_frames_b64,
|
||||
stacked_frames_b64=demo_stacked_b64,
|
||||
)
|
||||
turn_response_ids.append(response_id)
|
||||
turn_outcomes.append(outcome)
|
||||
@@ -1584,6 +1692,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
|
||||
"continuous_input": continuous_input,
|
||||
"continuous_input_ok": continuous_input_ok,
|
||||
"input_chunk_ms": args.chunk_ms,
|
||||
"video_frame_count": len(demo_frames_b64),
|
||||
"video_stacked_frame_count": sum(frame is not None for frame in demo_stacked_b64),
|
||||
"model_policy_settle_ms": args.model_policy_settle_ms,
|
||||
"turn_outcomes": turn_outcomes,
|
||||
"turn_outcomes_ok": turn_outcomes_ok,
|
||||
@@ -1641,6 +1751,19 @@ def parse_args() -> argparse.Namespace:
|
||||
default=None,
|
||||
help="Optional image file sent as an omni-duplex camera frame (one per 1 s of audio)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input-video",
|
||||
default=None,
|
||||
help="Optional video file streamed as omni-duplex camera frames (decoded at 1 fps, "
|
||||
"overrides --frame-image; the last frame is held if the audio outlives the clip)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stack-frames",
|
||||
type=int,
|
||||
default=1,
|
||||
choices=[1, 2],
|
||||
help="Frames per append: 2 also resends the previous frame for motion context.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--turn-input-wav",
|
||||
action="append",
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
import base64
|
||||
from pathlib import Path
|
||||
|
||||
from vllm_omni.experimental.fullduplex.client import (
|
||||
PCM16_BYTES_PER_SAMPLE,
|
||||
PCM16_SAMPLE_RATE,
|
||||
RealtimeDuplexClient,
|
||||
RealtimeEventCollector,
|
||||
build_realtime_url,
|
||||
read_pcm16_wav,
|
||||
wait_for,
|
||||
)
|
||||
|
||||
|
||||
def _events(events: list[dict[str, object]], response_id: str, kind: str) -> list[dict[str, object]]:
|
||||
return [e for e in events if e.get("type") == kind and RealtimeEventCollector.response_id(e) == response_id]
|
||||
|
||||
|
||||
async def run_server_vad_interrupt(args) -> dict[str, object]:
|
||||
initial = read_pcm16_wav(Path(args.input_wav))
|
||||
interrupt = read_pcm16_wav(Path(args.interrupt_wav))[: PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * 47 // 10]
|
||||
session_id = f"server-vad-hard-interrupt-{id(args)}"
|
||||
client = RealtimeDuplexClient(build_realtime_url(args.url, args.model, autostart=False, session_id=session_id))
|
||||
|
||||
async def until(predicate, label: str) -> None:
|
||||
await wait_for(predicate, timeout_s=args.timeout_s, label=label)
|
||||
|
||||
try:
|
||||
await client.__aenter__()
|
||||
await client.configure(
|
||||
args.model,
|
||||
ref_audio="data:audio/wav;base64," + base64.b64encode(Path(args.ref_audio).read_bytes()).decode(),
|
||||
session_id=session_id,
|
||||
turn_detection={"type": "server_vad", "interrupt_response": True},
|
||||
timeout_s=args.timeout_s,
|
||||
)
|
||||
await client.stream_pcm16(initial + bytes(16_000 * 2 * 1600 // 1000), chunk_ms=200, realtime=True)
|
||||
await client.commit()
|
||||
await until(lambda: any(e.get("type") == "response.created" for e in client.events.events), "response.created")
|
||||
target_id = RealtimeEventCollector.response_id(
|
||||
next(e for e in client.events.events if e.get("type") == "response.created")
|
||||
)
|
||||
assert target_id
|
||||
await until(lambda: _events(client.events.events, target_id, "response.audio.delta"), "response.audio.delta")
|
||||
cursor = len(client.events.events)
|
||||
await client.stream_pcm16(interrupt + bytes(16_000 * 2 * 800 // 1000), chunk_ms=args.chunk_ms, realtime=True)
|
||||
await until(lambda: _events(client.events.events, target_id, "response.done"), "cancelled response.done")
|
||||
await client.commit()
|
||||
await until(
|
||||
lambda: any(
|
||||
e.get("type") == "response.done" and RealtimeEventCollector.response_id(e) != target_id
|
||||
for e in client.events.events[cursor:]
|
||||
),
|
||||
"follow-up response.done",
|
||||
)
|
||||
events = client.events.events
|
||||
finally:
|
||||
try:
|
||||
await client.close_session(timeout_s=args.timeout_s)
|
||||
finally:
|
||||
await client.__aexit__(None, None, None)
|
||||
|
||||
done = _events(events, target_id, "response.done")
|
||||
terminal = done[0] if len(done) == 1 else None
|
||||
trailing = events[events.index(terminal) + 1 :] if terminal is not None else events
|
||||
stale = any(
|
||||
e.get("type") == "response.audio.delta" and RealtimeEventCollector.response_id(e) == target_id for e in trailing
|
||||
)
|
||||
followup_audio = any(
|
||||
e.get("type") == "response.audio.delta" and RealtimeEventCollector.response_id(e) != target_id
|
||||
for e in events[cursor:]
|
||||
)
|
||||
ok = (
|
||||
terminal is not None
|
||||
and isinstance(terminal.get("response"), dict)
|
||||
and terminal["response"].get("status") == "cancelled"
|
||||
and terminal["response"].get("status_details", {}).get("reason") == "turn_detected"
|
||||
and not stale
|
||||
and sum(e.get("type") == "input_audio_buffer.speech_started" for e in events[cursor:]) == 1
|
||||
and sum(e.get("type") == "input_audio_buffer.speech_stopped" for e in events[cursor:]) == 1
|
||||
and any(e.get("type") == "input_audio_buffer.committed" for e in events[cursor:])
|
||||
and followup_audio
|
||||
and not any(e.get("type") == "error" for e in events)
|
||||
)
|
||||
return {"ok": ok}
|
||||
@@ -5,8 +5,20 @@ Tests of common diffusion feature combinations in online serving mode
|
||||
for HunyuanVideo-1.5-T2V (480p).
|
||||
|
||||
Coverage (H100, since model cannot fit L4):
|
||||
- CacheDiT + Layerwise CPU offloading (1 GPU)
|
||||
- CacheDiT + TP=2 + VAE patch parallel=2 (2 GPUs)
|
||||
- CPU offloading (1 GPU) — ``core_model`` + ``advanced_model``
|
||||
- CacheDiT + Layerwise CPU offloading (1 GPU) — ``full_model``
|
||||
- CacheDiT + TP=2 + VAE patch parallel=2 (2 GPUs) — ``full_model``
|
||||
|
||||
HunyuanVideo-1.5 is a high-priority model, so only the most basic single-card deployment
|
||||
row runs on every PR (L2) and on merge (L3). The heavyweight feature combinations stay
|
||||
nightly-only (L4), together with the video similarity suites in
|
||||
``tests/e2e/accuracy/hunyuanvideo15_{t2v,i2v}/``.
|
||||
|
||||
From ``tests/``::
|
||||
|
||||
pytest -s -v e2e/online_serving/test_hunyuan_video_15_expansion.py -m "core_model and diffusion" --run-level=core_model
|
||||
pytest -s -v e2e/online_serving/test_hunyuan_video_15_expansion.py -m "advanced_model and diffusion" --run-level=advanced_model
|
||||
pytest -s -v e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and diffusion" --run-level=full_model
|
||||
"""
|
||||
|
||||
import pytest
|
||||
@@ -14,7 +26,7 @@ import pytest
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
|
||||
pytestmark = [pytest.mark.diffusion]
|
||||
|
||||
PROMPT = "A cat walking across a sunlit garden, cinematic lighting, slow motion."
|
||||
NEGATIVE_PROMPT = "low quality, blurry, distorted"
|
||||
@@ -24,9 +36,14 @@ MODEL = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
|
||||
SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||
PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
|
||||
|
||||
# Only the most basic single-card deployment row is cheap enough for PR (L2) / merge (L3)
|
||||
# CI; the CacheDiT / parallel feature combinations run nightly (L4).
|
||||
_CORE_MARKS = [pytest.mark.core_model, pytest.mark.advanced_model]
|
||||
_FULL_MARKS = [pytest.mark.full_model]
|
||||
|
||||
|
||||
def _get_diffusion_feature_cases(model: str):
|
||||
"""Return L4 diffusion feature cases for HunyuanVideo-1.5.
|
||||
"""Return diffusion feature cases for HunyuanVideo-1.5.
|
||||
|
||||
Designed for 2x H100 environment per issue #1832.
|
||||
"""
|
||||
@@ -40,7 +57,7 @@ def _get_diffusion_feature_cases(model: str):
|
||||
],
|
||||
),
|
||||
id="single_card_cpu_offload",
|
||||
marks=SINGLE_CARD_MARKS,
|
||||
marks=SINGLE_CARD_MARKS + _CORE_MARKS,
|
||||
),
|
||||
# (1 GPU) CacheDiT + Layerwise CPU offloading
|
||||
pytest.param(
|
||||
@@ -53,7 +70,7 @@ def _get_diffusion_feature_cases(model: str):
|
||||
],
|
||||
),
|
||||
id="single_card_cachedit_layerwise",
|
||||
marks=SINGLE_CARD_MARKS,
|
||||
marks=SINGLE_CARD_MARKS + _FULL_MARKS,
|
||||
),
|
||||
# (2 GPUs) CacheDiT + TP=2 + VAE patch parallel=2
|
||||
pytest.param(
|
||||
@@ -70,7 +87,7 @@ def _get_diffusion_feature_cases(model: str):
|
||||
],
|
||||
),
|
||||
id="parallel_cachedit_tp2_vae2",
|
||||
marks=PARALLEL_MARKS,
|
||||
marks=PARALLEL_MARKS + _FULL_MARKS,
|
||||
),
|
||||
]
|
||||
|
||||
@@ -84,7 +101,7 @@ def test_hunyuan_video_15_t2v(
|
||||
omni_server: OmniServer,
|
||||
online_client: OnlineOmniClient,
|
||||
):
|
||||
"""L4 diffusion feature coverage for HunyuanVideo-1.5-T2V on H100."""
|
||||
"""Diffusion feature coverage for HunyuanVideo-1.5-T2V on H100."""
|
||||
form_data = {
|
||||
"prompt": PROMPT,
|
||||
"negative_prompt": NEGATIVE_PROMPT,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Unit tests for acknowledged abort on AsyncOmniEngine."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -19,6 +22,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
def _make_engine(request_q: queue.Queue, rpc_q: queue.Queue) -> AsyncOmniEngine:
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = False
|
||||
engine.request_queue = SimpleNamespace(sync_q=request_q)
|
||||
engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
|
||||
engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
|
||||
@@ -154,3 +158,110 @@ def test_abort_async_preserves_request_queue_backpressure(mocker: MockerFixture)
|
||||
asyncio.run(_run())
|
||||
finally:
|
||||
engine._correlated_rpc_client.close()
|
||||
|
||||
|
||||
def test_abort_async_is_noop_after_shutdown_starts():
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = True
|
||||
engine.request_queue = None
|
||||
engine._correlated_rpc_client = None
|
||||
|
||||
asyncio.run(engine.abort_async(["req-late"]))
|
||||
|
||||
|
||||
def test_abort_async_is_noop_for_empty_request_ids():
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = False
|
||||
engine.request_queue = None
|
||||
engine._correlated_rpc_client = None
|
||||
|
||||
asyncio.run(engine.abort_async([]))
|
||||
|
||||
|
||||
def test_abort_async_tolerates_request_queue_close_during_shutdown(
|
||||
mocker: MockerFixture,
|
||||
):
|
||||
class SyntheticSyncQueueShutDownError(Exception):
|
||||
pass
|
||||
|
||||
rpc_q: queue.Queue = queue.Queue()
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = False
|
||||
|
||||
class ClosingRequestQueue:
|
||||
def put(self, item) -> None:
|
||||
del item
|
||||
engine._shutdown_called = True
|
||||
raise SyntheticSyncQueueShutDownError
|
||||
|
||||
request_q = ClosingRequestQueue()
|
||||
engine.request_queue = SimpleNamespace(sync_q=request_q)
|
||||
engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
|
||||
mocker.patch(
|
||||
"vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN",
|
||||
SyntheticSyncQueueShutDownError,
|
||||
)
|
||||
|
||||
try:
|
||||
asyncio.run(engine.abort_async(["req-submit-race"], timeout=1))
|
||||
finally:
|
||||
engine._correlated_rpc_client.close()
|
||||
|
||||
|
||||
def test_abort_async_surfaces_request_queue_close_while_engine_is_live(
|
||||
mocker: MockerFixture,
|
||||
):
|
||||
class SyntheticSyncQueueShutDownError(Exception):
|
||||
pass
|
||||
|
||||
class ClosedRequestQueue:
|
||||
def put(self, item) -> None:
|
||||
del item
|
||||
raise SyntheticSyncQueueShutDownError
|
||||
|
||||
request_q = ClosedRequestQueue()
|
||||
rpc_q: queue.Queue = queue.Queue()
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = False
|
||||
engine.request_queue = SimpleNamespace(sync_q=request_q)
|
||||
engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
|
||||
mocker.patch(
|
||||
"vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN",
|
||||
SyntheticSyncQueueShutDownError,
|
||||
)
|
||||
|
||||
try:
|
||||
with pytest.raises(SyntheticSyncQueueShutDownError):
|
||||
asyncio.run(engine.abort_async(["req-live-submit"], timeout=1))
|
||||
finally:
|
||||
engine._correlated_rpc_client.close()
|
||||
|
||||
|
||||
def test_abort_async_tolerates_rpc_router_close_during_shutdown():
|
||||
request_q: queue.Queue = queue.Queue()
|
||||
rpc_q: queue.Queue = queue.Queue()
|
||||
engine = _make_engine(request_q, rpc_q)
|
||||
|
||||
async def _run() -> None:
|
||||
task = asyncio.create_task(engine.abort_async(["req-ack-race"], timeout=1))
|
||||
await asyncio.to_thread(request_q.get, True, 1)
|
||||
engine._shutdown_called = True
|
||||
engine._correlated_rpc_client.close()
|
||||
await task
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_abort_async_surfaces_rpc_router_close_while_engine_is_live():
|
||||
request_q: queue.Queue = queue.Queue()
|
||||
rpc_q: queue.Queue = queue.Queue()
|
||||
engine = _make_engine(request_q, rpc_q)
|
||||
|
||||
async def _run() -> None:
|
||||
task = asyncio.create_task(engine.abort_async(["req-live"], timeout=1))
|
||||
await asyncio.to_thread(request_q.get, True, 1)
|
||||
engine._correlated_rpc_client.close()
|
||||
with pytest.raises(RuntimeError, match="RPC result router closed"):
|
||||
await task
|
||||
|
||||
asyncio.run(_run())
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from pytest_mock import MockerFixture
|
||||
@@ -14,6 +17,104 @@ from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _SyntheticSyncQueueShutDownError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def test_abort_is_noop_for_empty_request_ids(mocker: MockerFixture):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = False
|
||||
engine.request_queue = mocker.Mock()
|
||||
|
||||
engine.abort([])
|
||||
|
||||
engine.request_queue.sync_q.put.assert_not_called()
|
||||
|
||||
|
||||
def test_abort_is_noop_after_shutdown_starts(mocker: MockerFixture):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = True
|
||||
engine.request_queue = mocker.Mock()
|
||||
|
||||
engine.abort(["req-1"])
|
||||
|
||||
engine.request_queue.sync_q.put.assert_not_called()
|
||||
|
||||
|
||||
def test_abort_tolerates_queue_close_race_during_shutdown(mocker: MockerFixture):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = False
|
||||
engine.request_queue = mocker.Mock()
|
||||
mocker.patch(
|
||||
"vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN",
|
||||
_SyntheticSyncQueueShutDownError,
|
||||
)
|
||||
|
||||
def close_queue_during_abort(*args, **kwargs):
|
||||
del args, kwargs
|
||||
engine._shutdown_called = True
|
||||
raise _SyntheticSyncQueueShutDownError
|
||||
|
||||
engine.request_queue.sync_q.put.side_effect = close_queue_during_abort
|
||||
|
||||
engine.abort(["req-1"])
|
||||
|
||||
|
||||
def test_abort_tolerates_legacy_janus_close_error_during_shutdown(mocker: MockerFixture):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = False
|
||||
engine.request_queue = mocker.Mock()
|
||||
mocker.patch(
|
||||
"vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN",
|
||||
None,
|
||||
)
|
||||
|
||||
def close_queue_during_abort(*args, **kwargs):
|
||||
del args, kwargs
|
||||
engine._shutdown_called = True
|
||||
raise RuntimeError("Operation on the closed queue is forbidden")
|
||||
|
||||
engine.request_queue.sync_q.put.side_effect = close_queue_during_abort
|
||||
|
||||
engine.abort(["req-1"])
|
||||
|
||||
|
||||
def test_abort_surfaces_unrelated_legacy_runtime_error_during_shutdown(
|
||||
mocker: MockerFixture,
|
||||
):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = False
|
||||
engine.request_queue = mocker.Mock()
|
||||
mocker.patch(
|
||||
"vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN",
|
||||
None,
|
||||
)
|
||||
|
||||
def fail_during_abort(*args, **kwargs):
|
||||
del args, kwargs
|
||||
engine._shutdown_called = True
|
||||
raise RuntimeError("unrelated queue failure")
|
||||
|
||||
engine.request_queue.sync_q.put.side_effect = fail_during_abort
|
||||
|
||||
with pytest.raises(RuntimeError, match="unrelated queue failure"):
|
||||
engine.abort(["req-1"])
|
||||
|
||||
|
||||
def test_abort_surfaces_unexpected_closed_queue(mocker: MockerFixture):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = False
|
||||
engine.request_queue = mocker.Mock()
|
||||
mocker.patch(
|
||||
"vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN",
|
||||
_SyntheticSyncQueueShutDownError,
|
||||
)
|
||||
engine.request_queue.sync_q.put.side_effect = _SyntheticSyncQueueShutDownError
|
||||
|
||||
with pytest.raises(_SyntheticSyncQueueShutDownError):
|
||||
engine.abort(["req-1"])
|
||||
|
||||
|
||||
def _make_engine_core_request(request_id: str = "req-1") -> EngineCoreRequest:
|
||||
return EngineCoreRequest(
|
||||
request_id=request_id,
|
||||
@@ -67,6 +168,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
|
||||
@@ -2216,7 +2216,12 @@ def test_stage_pool_metrics_use_resumable_segment_token_count() -> None:
|
||||
)
|
||||
output = SimpleNamespace(
|
||||
request_id="req-stream",
|
||||
outputs=[SimpleNamespace(cumulative_token_ids=list(range(11)))],
|
||||
outputs=[
|
||||
SimpleNamespace(
|
||||
cumulative_token_ids=list(range(11)),
|
||||
finish_reason=FinishReason.LENGTH,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
metrics = pool.build_stage_metrics(
|
||||
@@ -2228,6 +2233,7 @@ def test_stage_pool_metrics_use_resumable_segment_token_count() -> None:
|
||||
|
||||
assert metrics.num_tokens_out == 3
|
||||
assert metrics.output_unit_count == 3
|
||||
assert metrics.finish_reason == "length"
|
||||
|
||||
|
||||
def test_image_ttfo_preserves_request_time_and_tracks_stage_time() -> None:
|
||||
|
||||
@@ -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"}
|
||||
@@ -389,6 +398,12 @@ async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_rout
|
||||
async def fake_omni_init_app_state(engine_client, state, args):
|
||||
state.engine_client = engine_client
|
||||
state.initialized_by_omni = True
|
||||
state.openai_serving_video = SimpleNamespace(
|
||||
shutdown=lambda: captured.__setitem__("video_shutdown", True),
|
||||
)
|
||||
state.openai_serving_speech = SimpleNamespace(
|
||||
shutdown=lambda: captured.__setitem__("speech_shutdown", True),
|
||||
)
|
||||
|
||||
async def fake_storage_start():
|
||||
captured["storage_started"] = True
|
||||
@@ -424,6 +439,8 @@ async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_rout
|
||||
assert captured["supported_tasks"] == ("generate",)
|
||||
assert captured["storage_started"] is True
|
||||
assert captured["restrictions"] == {}
|
||||
assert captured["video_shutdown"] is True
|
||||
assert captured["speech_shutdown"] is True
|
||||
assert sock.closed is True
|
||||
assert served_app.state.initialized_by_omni is True
|
||||
|
||||
@@ -437,6 +454,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:
|
||||
@@ -728,7 +747,11 @@ async def test_pure_diffusion_app_state_key_snapshot(monkeypatch) -> None:
|
||||
)
|
||||
|
||||
state = State()
|
||||
await api_server.omni_init_app_state(engine, state, _minimal_args())
|
||||
await api_server.omni_init_app_state(
|
||||
engine,
|
||||
state,
|
||||
_minimal_args(),
|
||||
)
|
||||
|
||||
_assert_app_state_snapshot(
|
||||
state,
|
||||
@@ -739,6 +762,55 @@ async def test_pure_diffusion_app_state_key_snapshot(monkeypatch) -> None:
|
||||
assert state.diffusion_engine is engine
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pure_diffusion_speech_forwards_media_access_args(monkeypatch) -> None:
|
||||
stage = SimpleNamespace(engine_args={})
|
||||
engine = _FakeEngineClient(stage_configs=[stage])
|
||||
speech_kwargs = {}
|
||||
|
||||
def _for_diffusion_factory(label: str):
|
||||
@classmethod
|
||||
def _factory(cls, *args, **kwargs):
|
||||
return _marker(label)
|
||||
|
||||
return _factory
|
||||
|
||||
@classmethod
|
||||
def _speech_factory(cls, *args, **kwargs):
|
||||
speech_kwargs.update(kwargs)
|
||||
return _marker("speech")
|
||||
|
||||
monkeypatch.setattr(api_server, "get_stage_type", lambda _cfg: "diffusion")
|
||||
monkeypatch.setattr(api_server.OmniOpenAIServingChat, "for_diffusion", _for_diffusion_factory("chat"))
|
||||
monkeypatch.setattr(api_server.OmniOpenAIServingChatBatch, "for_diffusion", _for_diffusion_factory("chat_batch"))
|
||||
monkeypatch.setattr(
|
||||
api_server.OmniOpenAIServingAudioGenerate,
|
||||
"for_diffusion",
|
||||
_for_diffusion_factory("audio_generate"),
|
||||
)
|
||||
monkeypatch.setattr(api_server.OmniOpenAIServingVideo, "for_diffusion", _for_diffusion_factory("video"))
|
||||
monkeypatch.setattr(api_server.OmniStreamingVideoOutputHandler, "__init__", lambda self, *a, **k: None)
|
||||
monkeypatch.setattr(api_server.OmniOpenAIServingSpeech, "for_diffusion", _speech_factory)
|
||||
monkeypatch.setattr(
|
||||
api_server.ServingRealtimeRobotOpenPI,
|
||||
"create_policy_server",
|
||||
classmethod(lambda cls, *a, **k: _marker("openpi")),
|
||||
)
|
||||
|
||||
state = State()
|
||||
await api_server.omni_init_app_state(
|
||||
engine,
|
||||
state,
|
||||
_minimal_args(
|
||||
allowed_local_media_path="/allowed/media",
|
||||
allowed_media_domains=["media.example.com"],
|
||||
),
|
||||
)
|
||||
|
||||
assert speech_kwargs["allowed_local_media_path"] == "/allowed/media"
|
||||
assert speech_kwargs["allowed_media_domains"] == ["media.example.com"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multistage_app_state_key_snapshot(monkeypatch) -> None:
|
||||
"""Lock multi-stage ``app.state`` keys after init, including live vs None.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -45,9 +45,15 @@ from vllm_omni.entrypoints.openai.serving_speech import (
|
||||
OmniOpenAIServingSpeech,
|
||||
_create_wav_header,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.base import DEFAULT_TTS_LANGUAGES, PreparedRequest, SpeechServingContext
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.base import (
|
||||
DEFAULT_TTS_LANGUAGES,
|
||||
PreparedRequest,
|
||||
SpeechServingContext,
|
||||
TTSGenerationError,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_supported_speakers
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSCodecLimitError
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter
|
||||
from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
|
||||
FISH_TEXT_ONLY_SYSTEM_PROMPT,
|
||||
@@ -1245,6 +1251,46 @@ class TestTTSMethods:
|
||||
cache_key
|
||||
) == speech_server._make_ref_audio_artifact_cache_key(np.asarray(first[0], dtype=np.float32), 24000)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_diffusion_ref_audio_uses_media_access_config(
|
||||
self,
|
||||
mocker: MockerFixture,
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
):
|
||||
monkeypatch.setenv("SPEAKER_SAMPLES_DIR", str(tmp_path / "speakers"))
|
||||
connector = mocker.MagicMock()
|
||||
connector.fetch_audio_async = mocker.AsyncMock(return_value=(np.zeros(24000, dtype=np.float32), 24000))
|
||||
connector_cls = mocker.patch.object(
|
||||
serving_speech_module,
|
||||
"MediaConnector",
|
||||
return_value=connector,
|
||||
)
|
||||
media_dir = tmp_path / "media"
|
||||
media_dir.mkdir()
|
||||
ref_audio = media_dir / "reference.wav"
|
||||
ref_audio.write_bytes(b"v1")
|
||||
|
||||
server = OmniOpenAIServingSpeech.for_diffusion(
|
||||
diffusion_engine=mocker.MagicMock(),
|
||||
model_name="test-model",
|
||||
allowed_local_media_path=str(media_dir),
|
||||
allowed_media_domains=["media.example.com"],
|
||||
)
|
||||
try:
|
||||
first = await server._resolve_ref_audio(ref_audio.as_uri())
|
||||
ref_audio.write_bytes(b"v2-with-different-size")
|
||||
second = await server._resolve_ref_audio(ref_audio.as_uri())
|
||||
finally:
|
||||
server.shutdown()
|
||||
|
||||
connector_cls.assert_called_once_with(
|
||||
allowed_local_media_path=str(media_dir),
|
||||
allowed_media_domains=["media.example.com"],
|
||||
)
|
||||
assert first[2] != second[2]
|
||||
assert connector.fetch_audio_async.await_count == 2
|
||||
|
||||
# ── ref-audio cache key tests (static helper) ──
|
||||
|
||||
def test_local_file_cache_key_invalidation(self):
|
||||
@@ -2888,6 +2934,40 @@ class TestStreamingResponse:
|
||||
assert "text/event-stream" not in response.headers["content-type"]
|
||||
assert len(response.content) > 0
|
||||
|
||||
def test_raw_stream_passes_tts_params_to_common_guard(self, streaming_app, monkeypatch):
|
||||
"""The raw route must not bypass the shared Qwen3-TTS EOS guard."""
|
||||
speech_server = streaming_app.state.speech_server
|
||||
finalized_tts_params = {"_qwen3_tts_effective_max_tokens": [192]}
|
||||
captured: dict = {}
|
||||
|
||||
async def prepare(_request, request_id=None):
|
||||
return request_id, object(), finalized_tts_params
|
||||
|
||||
async def generate_chunks(
|
||||
_generator,
|
||||
_request_id,
|
||||
_response_format="pcm",
|
||||
raw_request=None,
|
||||
request_start_s=None,
|
||||
include_sample_rate=False,
|
||||
usage_acc=None,
|
||||
tts_params=None,
|
||||
collect=None,
|
||||
):
|
||||
captured["tts_params"] = tts_params
|
||||
yield b"\x00\x00"
|
||||
|
||||
monkeypatch.setattr(speech_server, "_prepare_speech_generation", prepare)
|
||||
monkeypatch.setattr(speech_server, "_generate_audio_chunks", generate_chunks)
|
||||
|
||||
response = TestClient(streaming_app).post(
|
||||
"/v1/audio/speech",
|
||||
json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert captured["tts_params"] is finalized_tts_params
|
||||
|
||||
def test_sse_streaming(self, streaming_app):
|
||||
"""stream_format=sse without stream=True returns audio deltas as SSE."""
|
||||
client = TestClient(streaming_app)
|
||||
@@ -4808,6 +4888,57 @@ class TestTTSAsyncOffloading:
|
||||
assert ("artifact-a", False) not in qwen3_tts_server._ref_audio_model_artifact_ready
|
||||
assert "artifact-b" in {entry[3] for entry in qwen3_tts_server._ref_audio_resolve_cache.values()}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_qwen3_tts_nonstream_retries_codec_limit_once(self, qwen3_tts_server, mocker: MockerFixture) -> None:
|
||||
qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
|
||||
qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(
|
||||
side_effect=[
|
||||
Qwen3TTSCodecLimitError("codec limit"),
|
||||
(b"RIFF" + b"\x00" * 32, "audio/wav"),
|
||||
]
|
||||
)
|
||||
request = OpenAICreateSpeechRequest(input="Hello", task_type="Base")
|
||||
|
||||
response = await qwen3_tts_server.create_speech(request)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert qwen3_tts_server._generate_audio_bytes.await_count == 2
|
||||
first_call, retry_call = qwen3_tts_server._generate_audio_bytes.await_args_list
|
||||
assert first_call.args[0] is request
|
||||
assert first_call.kwargs["request_id"] != retry_call.kwargs["request_id"]
|
||||
assert retry_call.args[0].seed is not None
|
||||
assert request.seed is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nonstream_does_not_retry_nonretryable_generation_error(
|
||||
self, qwen3_tts_server, mocker: MockerFixture
|
||||
) -> None:
|
||||
qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
|
||||
qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(side_effect=TTSGenerationError("invalid generation"))
|
||||
|
||||
response = await qwen3_tts_server.create_speech(OpenAICreateSpeechRequest(input="Hello", task_type="Base"))
|
||||
|
||||
assert response.status_code == 500
|
||||
qwen3_tts_server._generate_audio_bytes.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("request_kwargs", [{"seed": 100423}, {"max_new_tokens": 192}])
|
||||
async def test_qwen3_tts_nonstream_does_not_retry_explicit_sampling_controls(
|
||||
self,
|
||||
qwen3_tts_server,
|
||||
mocker: MockerFixture,
|
||||
request_kwargs: dict[str, int],
|
||||
) -> None:
|
||||
qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
|
||||
qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(side_effect=Qwen3TTSCodecLimitError("codec limit"))
|
||||
|
||||
response = await qwen3_tts_server.create_speech(
|
||||
OpenAICreateSpeechRequest(input="Hello", task_type="Base", **request_kwargs)
|
||||
)
|
||||
|
||||
assert response.status_code == 500
|
||||
qwen3_tts_server._generate_audio_bytes.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_error(self, qwen3_tts_server):
|
||||
async def failing_generator():
|
||||
@@ -4822,6 +4953,80 @@ class TestTTSAsyncOffloading:
|
||||
assert "req-fail" not in qwen3_tts_server._request_ref_audio_artifact_keys
|
||||
assert ("artifact-fail", False) not in qwen3_tts_server._ref_audio_model_artifact_ready
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_audio_chunks_rejects_qwen3_codec_limit_before_success(self, qwen3_tts_server):
|
||||
async def length_limited_generator():
|
||||
yield SimpleNamespace(
|
||||
multimodal_output={"audio": torch.zeros(16, dtype=torch.float32), "sr": 24000},
|
||||
metrics={
|
||||
"stage_metrics": {
|
||||
"0": {
|
||||
"num_tokens_out": 191,
|
||||
"finish_reason": "length",
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
chunks = qwen3_tts_server._generate_audio_chunks(
|
||||
length_limited_generator(),
|
||||
"req-codec-limit",
|
||||
tts_params={"task_type": ["Base"], "_qwen3_tts_effective_max_tokens": [192]},
|
||||
)
|
||||
with pytest.raises(Qwen3TTSCodecLimitError, match="191/192"):
|
||||
async for _ in chunks:
|
||||
pass
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_codec_limit_marks_partial_audio_for_discard(self, qwen3_tts_server):
|
||||
async def length_limited_generator():
|
||||
yield SimpleNamespace(
|
||||
multimodal_output={"audio": torch.zeros(16, dtype=torch.float32), "sr": 24000},
|
||||
metrics={
|
||||
"stage_metrics": {
|
||||
"0": {
|
||||
"num_tokens_out": 191,
|
||||
"finish_reason": "length",
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
events = [
|
||||
event
|
||||
async for event in qwen3_tts_server._generate_audio_sse_events(
|
||||
length_limited_generator(),
|
||||
"req-sse-codec-limit",
|
||||
tts_params={"task_type": ["Base"], "_qwen3_tts_effective_max_tokens": [192]},
|
||||
)
|
||||
]
|
||||
|
||||
assert any("event: speech.audio.delta" in event for event in events)
|
||||
assert not any("event: speech.audio.done" in event for event in events)
|
||||
error_event = next(event for event in events if "event: speech.audio.error" in event)
|
||||
payload = json.loads(next(line for line in error_event.splitlines() if line.startswith("data: "))[6:])
|
||||
assert payload["error"]["partial_audio"] is True
|
||||
assert payload["error"]["action"] == "discard"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_error_before_audio_is_not_marked_partial(self, qwen3_tts_server):
|
||||
async def failing_generator():
|
||||
raise RuntimeError("boom")
|
||||
yield # pragma: no cover
|
||||
|
||||
events = [
|
||||
event
|
||||
async for event in qwen3_tts_server._generate_audio_sse_events(
|
||||
failing_generator(),
|
||||
"req-sse-no-audio",
|
||||
)
|
||||
]
|
||||
|
||||
error_event = next(event for event in events if "event: speech.audio.error" in event)
|
||||
payload = json.loads(next(line for line in error_event.splitlines() if line.startswith("data: "))[6:])
|
||||
assert "partial_audio" not in payload["error"]
|
||||
assert "action" not in payload["error"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server):
|
||||
async def pcm_generator():
|
||||
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
from types import SimpleNamespace
|
||||
@@ -44,7 +47,9 @@ def _build_test_app(
|
||||
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {}))
|
||||
speech_service.forced_aligner_enabled = False
|
||||
|
||||
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
|
||||
async def mock_generate_pcm_chunks(
|
||||
_generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
|
||||
):
|
||||
for chunk in (b"\x01\x02", b"\x03\x04\x05"):
|
||||
yield (chunk, 24000) if include_sample_rate else chunk
|
||||
|
||||
@@ -244,6 +249,7 @@ class TestStreamingSpeechWebSocket:
|
||||
|
||||
def test_streaming_multiple_binary_frames(self, mocker: MockerFixture):
|
||||
captured_requests = []
|
||||
captured_tts_params = []
|
||||
|
||||
speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
|
||||
speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
|
||||
@@ -253,11 +259,12 @@ class TestStreamingSpeechWebSocket:
|
||||
|
||||
async def mock_prepare_speech_generation(request):
|
||||
captured_requests.append(request)
|
||||
return "req-stream", object(), {}
|
||||
return "req-stream", object(), {"_qwen3_tts_effective_max_tokens": [192]}
|
||||
|
||||
speech_service._prepare_speech_generation = mock_prepare_speech_generation
|
||||
|
||||
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
|
||||
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
|
||||
captured_tts_params.append(tts_params)
|
||||
for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"):
|
||||
yield (chunk, 24000) if include_sample_rate else chunk
|
||||
|
||||
@@ -302,6 +309,7 @@ class TestStreamingSpeechWebSocket:
|
||||
assert captured_requests[0].stream is True
|
||||
assert captured_requests[0].response_format == "pcm"
|
||||
assert captured_requests[0].initial_codec_chunk_frames == 12
|
||||
assert captured_tts_params == [{"_qwen3_tts_effective_max_tokens": [192]}]
|
||||
assert speech_service._generate_audio_bytes.await_count == 0
|
||||
|
||||
def test_word_timestamps_requires_configured_aligner(self, mocker: MockerFixture):
|
||||
@@ -344,7 +352,9 @@ class TestStreamingSpeechWebSocket:
|
||||
|
||||
# The forced-aligner stage rides the same generator: its pooling output
|
||||
# is surfaced via the ``collect`` channel once the audio has streamed.
|
||||
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
|
||||
async def mock_generate_pcm_chunks(
|
||||
_generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
|
||||
):
|
||||
for chunk in (first_chunk, second_chunk):
|
||||
yield (chunk, 1000) if include_sample_rate else chunk
|
||||
if collect is not None:
|
||||
@@ -427,7 +437,9 @@ class TestStreamingSpeechWebSocket:
|
||||
speech_service.forced_aligner_enabled = True
|
||||
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {}))
|
||||
|
||||
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
|
||||
async def mock_generate_pcm_chunks(
|
||||
_generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
|
||||
):
|
||||
chunk = b"\x01" * 1000
|
||||
yield (chunk, 1000) if include_sample_rate else chunk
|
||||
if collect is not None:
|
||||
@@ -608,7 +620,7 @@ class TestStreamingSpeechWebSocket:
|
||||
speech_service.engine_client.abort = mocker.AsyncMock()
|
||||
speech_service.forced_aligner_enabled = False
|
||||
|
||||
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
|
||||
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
|
||||
yield b"\x01\x02"
|
||||
raise RuntimeError("stream boom")
|
||||
|
||||
@@ -633,6 +645,8 @@ class TestStreamingSpeechWebSocket:
|
||||
assert ws.receive_json() == {
|
||||
"type": "error",
|
||||
"message": "Generation failed for utterance 0, sentence 0: stream boom",
|
||||
"partial_audio": True,
|
||||
"action": "discard",
|
||||
}
|
||||
assert ws.receive_json() == {
|
||||
"type": "audio.done",
|
||||
@@ -702,7 +716,7 @@ class TestStreamingSpeechWebSocket:
|
||||
speech_service.engine_client.abort = mocker.AsyncMock()
|
||||
speech_service.forced_aligner_enabled = False
|
||||
|
||||
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
|
||||
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
|
||||
yield b"\x01\x02"
|
||||
|
||||
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
|
||||
|
||||
@@ -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()],
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Unit tests for speech token-usage accounting (issue #4646).
|
||||
|
||||
These exercise the pure usage logic without a model/server: input-token
|
||||
@@ -205,6 +206,14 @@ def test_output_tokens_from_stage_metrics():
|
||||
assert acc.total() == 77
|
||||
|
||||
|
||||
def test_output_tokens_preserves_stage0_finish_reason():
|
||||
acc = SpeechOutputTokenCounter()
|
||||
res = _final_res(stage0_out=128)
|
||||
res.metrics["stage_metrics"]["0"]["finish_reason"] = "length"
|
||||
acc.observe(res)
|
||||
assert acc.stage0_finish_reason == "length"
|
||||
|
||||
|
||||
def test_output_tokens_takes_max_across_stages():
|
||||
acc = SpeechOutputTokenCounter()
|
||||
acc.observe(_final_res(stage0_out=42, stage1_out=0))
|
||||
|
||||
@@ -10,11 +10,14 @@ from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
|
||||
from vllm_omni.entrypoints.openai.tts_adapters import (
|
||||
TTS_ADAPTER_REGISTRY,
|
||||
ARTTSAdapter,
|
||||
DiffusionTTSAdapter,
|
||||
SpeechServingContext,
|
||||
all_tts_model_types,
|
||||
detect_tts_model_type,
|
||||
resolve_adapter,
|
||||
@@ -29,7 +32,11 @@ from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import (
|
||||
MossTTSAdapter,
|
||||
MossTTSNanoAdapter,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import (
|
||||
QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY,
|
||||
Qwen3TTSAdapter,
|
||||
Qwen3TTSCodecLimitError,
|
||||
)
|
||||
from vllm_omni.model_executor.models.indextts2 import prompt_utils
|
||||
from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
|
||||
INDEXTTS25_TOKENIZER_FILE,
|
||||
@@ -115,11 +122,100 @@ def test_moss_tts_applies_request_max_new_tokens(adapter_cls):
|
||||
assert stage_defaults[0].max_tokens == 4096
|
||||
|
||||
|
||||
def _build_moss_tts_request(adapter_cls, mocker, *, request_seed):
|
||||
server = mocker.Mock()
|
||||
server._build_moss_tts_params = mocker.AsyncMock(return_value={})
|
||||
adapter = adapter_cls(SpeechServingContext(server=server))
|
||||
request = OpenAICreateSpeechRequest(input="hello", seed=request_seed)
|
||||
|
||||
return asyncio.run(
|
||||
adapter.build(
|
||||
request,
|
||||
[SamplingParams(seed=42)],
|
||||
has_inline_ref_audio=False,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# Full-family coverage pins the adapter contract; only Nano consumes this seed end to end today.
|
||||
@pytest.mark.parametrize("adapter_cls", [MossTTSAdapter, MossTTSNanoAdapter])
|
||||
@pytest.mark.parametrize("request_seed", [0, 1234])
|
||||
def test_moss_tts_request_seed_overrides_stage_default(adapter_cls, request_seed, mocker):
|
||||
prepared = _build_moss_tts_request(
|
||||
adapter_cls,
|
||||
mocker,
|
||||
request_seed=request_seed,
|
||||
)
|
||||
|
||||
assert prepared.tts_params["seed"] == [request_seed]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("adapter_cls", [MossTTSAdapter, MossTTSNanoAdapter])
|
||||
def test_moss_tts_seed_falls_back_to_stage_default(adapter_cls, mocker):
|
||||
prepared = _build_moss_tts_request(
|
||||
adapter_cls,
|
||||
mocker,
|
||||
request_seed=None,
|
||||
)
|
||||
|
||||
assert prepared.tts_params["seed"] == [42]
|
||||
|
||||
|
||||
def test_qwen3_tts_metadata():
|
||||
assert Qwen3TTSAdapter.backend == "ar"
|
||||
assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("task_type", "text_tokens", "request_cap", "expected_cap"),
|
||||
[
|
||||
("Base", 0, None, 4096),
|
||||
("Base", 10, None, 192),
|
||||
("Base", 23, None, 276),
|
||||
("Base", 23, 128, 128),
|
||||
("Base", 23, 512, 512),
|
||||
("Base", 400, None, 4096),
|
||||
("CustomVoice", 10, None, 4096),
|
||||
("CustomVoice", 10, 256, 256),
|
||||
],
|
||||
)
|
||||
def test_qwen3_tts_applies_text_scaled_codec_safety_limit(task_type, text_tokens, request_cap, expected_cap, mocker):
|
||||
server = mocker.Mock()
|
||||
server._count_usage_text_tokens.return_value = text_tokens
|
||||
adapter = Qwen3TTSAdapter(SpeechServingContext(server=server))
|
||||
stage_defaults = [SamplingParams(max_tokens=4096, min_tokens=2)]
|
||||
request = OpenAICreateSpeechRequest(
|
||||
input="test text",
|
||||
task_type=task_type,
|
||||
max_new_tokens=request_cap,
|
||||
)
|
||||
prompt: dict[str, Any] = {"additional_information": {}}
|
||||
|
||||
overridden = adapter.apply_sampling_overrides(stage_defaults, request, prompt)
|
||||
|
||||
assert overridden[0].max_tokens == expected_cap
|
||||
assert prompt["additional_information"][QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY] == [expected_cap]
|
||||
assert stage_defaults[0].max_tokens == 4096
|
||||
|
||||
|
||||
def test_qwen3_tts_rejects_only_length_finished_base_audio(mocker):
|
||||
adapter = Qwen3TTSAdapter(SpeechServingContext(server=mocker.Mock()))
|
||||
params = {
|
||||
"task_type": ["Base"],
|
||||
QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY: [192],
|
||||
}
|
||||
|
||||
# EOS at the exact budget is valid; the token count alone is not a
|
||||
# sufficient failure signal.
|
||||
adapter.validate_generation(params, stage0_finish_reason="stop", output_tokens=192)
|
||||
adapter.validate_generation(params, stage0_finish_reason=None, output_tokens=192)
|
||||
|
||||
# Some frontends expose 191 decoded frames for max_new_tokens=192. The
|
||||
# engine terminal reason still makes this an unambiguous limit failure.
|
||||
with pytest.raises(Qwen3TTSCodecLimitError, match="191/192"):
|
||||
adapter.validate_generation(params, stage0_finish_reason="length", output_tokens=191)
|
||||
|
||||
|
||||
def test_indextts_adapters_are_versioned():
|
||||
assert resolve_adapter("indextts2") is IndexTTS2Adapter
|
||||
assert resolve_adapter("indextts2_5") is IndexTTS25Adapter
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
"""Unit tests for OpenAI-compatible video API encoding helpers."""
|
||||
|
||||
import base64
|
||||
import threading
|
||||
from io import BytesIO
|
||||
|
||||
import av
|
||||
@@ -249,6 +250,95 @@ def test_channel_first_video_tensor_uses_direct_planar_mux(monkeypatch):
|
||||
assert options["fps"] == 12.0
|
||||
|
||||
|
||||
def test_planar_converter_serial_and_parallel_outputs_match():
|
||||
frames = [np.full((2, 3, 3), fill_value=value / 10, dtype=np.float32) for value in (1, 4, 7)]
|
||||
|
||||
serial_converter = video_api_utils._PlanarFrameConverter(max_workers=1)
|
||||
parallel_converter = video_api_utils._PlanarFrameConverter(max_workers=2)
|
||||
try:
|
||||
serial = [
|
||||
frame.to_ndarray(format="rgb24") for frame in serial_converter.iter_frames(frames, np.dtype(np.float32))
|
||||
]
|
||||
parallel = [
|
||||
frame.to_ndarray(format="rgb24") for frame in parallel_converter.iter_frames(frames, np.dtype(np.float32))
|
||||
]
|
||||
finally:
|
||||
serial_converter.shutdown()
|
||||
parallel_converter.shutdown()
|
||||
|
||||
np.testing.assert_array_equal(np.stack(serial), np.stack(parallel))
|
||||
|
||||
|
||||
def test_planar_converter_preserves_fifo_order_and_bounds_pending_futures(monkeypatch):
|
||||
converter = video_api_utils._PlanarFrameConverter(max_workers=2)
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
four_submitted = threading.Event()
|
||||
outstanding = 0
|
||||
peak_outstanding = 0
|
||||
lock = threading.Lock()
|
||||
real_submit = converter._executor.submit
|
||||
|
||||
def fake_build(frame, _common_dtype):
|
||||
started.set()
|
||||
release.wait(timeout=2)
|
||||
return int(frame[0, 0, 0])
|
||||
|
||||
def submit(*args, **kwargs):
|
||||
nonlocal outstanding, peak_outstanding
|
||||
with lock:
|
||||
outstanding += 1
|
||||
peak_outstanding = max(peak_outstanding, outstanding)
|
||||
if peak_outstanding >= 4:
|
||||
four_submitted.set()
|
||||
future = real_submit(*args, **kwargs)
|
||||
|
||||
def done(_future):
|
||||
nonlocal outstanding
|
||||
with lock:
|
||||
outstanding -= 1
|
||||
|
||||
future.add_done_callback(done)
|
||||
return future
|
||||
|
||||
monkeypatch.setattr(converter, "_build_frame", fake_build)
|
||||
monkeypatch.setattr(converter._executor, "submit", submit)
|
||||
frames = [np.full((1, 1, 3), value, dtype=np.uint8) for value in range(8)]
|
||||
result: list[int] = []
|
||||
iterator = converter.iter_frames(frames, np.dtype(np.uint8))
|
||||
consumer = threading.Thread(target=lambda: result.extend(iterator))
|
||||
consumer.start()
|
||||
try:
|
||||
assert started.wait(timeout=2)
|
||||
assert four_submitted.wait(timeout=2)
|
||||
release.set()
|
||||
consumer.join(timeout=2)
|
||||
assert not consumer.is_alive()
|
||||
finally:
|
||||
release.set()
|
||||
iterator.close()
|
||||
converter.shutdown()
|
||||
|
||||
assert result == list(range(8))
|
||||
assert peak_outstanding <= 4
|
||||
|
||||
|
||||
def test_planar_converter_reuses_executor_and_shutdowns_after_cancel():
|
||||
converter = video_api_utils._PlanarFrameConverter(max_workers=2)
|
||||
executor = converter._executor
|
||||
frames = [np.zeros((2, 2, 3), dtype=np.uint8) for _ in range(4)]
|
||||
try:
|
||||
list(converter.iter_frames(frames, np.dtype(np.uint8)))
|
||||
iterator = converter.iter_frames(frames, np.dtype(np.uint8))
|
||||
next(iterator)
|
||||
iterator.close()
|
||||
assert converter._executor is executor
|
||||
finally:
|
||||
converter.shutdown()
|
||||
|
||||
assert executor._shutdown
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("video", "expected_path", "expected_reason", "expected_mux"),
|
||||
[
|
||||
@@ -347,6 +437,35 @@ def test_video_only_does_not_log_default_audio_sample_rate(monkeypatch):
|
||||
assert log_messages == []
|
||||
|
||||
|
||||
def test_planar_mux_accepts_none_rate_and_defaults_audio_to_44100():
|
||||
def make_frame():
|
||||
return next(
|
||||
video_api_utils._iter_planar_video_frames(
|
||||
[np.zeros((2, 2, 3), dtype=np.uint8)],
|
||||
np.dtype(np.uint8),
|
||||
)
|
||||
)
|
||||
|
||||
video_only = media_utils.mux_av_video_audio_bytes(
|
||||
[make_frame()],
|
||||
width=2,
|
||||
height=2,
|
||||
audio_sample_rate=None,
|
||||
)
|
||||
with av.open(BytesIO(video_only), mode="r", format="mp4") as container:
|
||||
assert not container.streams.audio
|
||||
|
||||
with_audio = media_utils.mux_av_video_audio_bytes(
|
||||
[make_frame()],
|
||||
width=2,
|
||||
height=2,
|
||||
audio_waveform=np.zeros((1, 1024), dtype=np.float32),
|
||||
audio_sample_rate=None,
|
||||
)
|
||||
with av.open(BytesIO(with_audio), mode="r", format="mp4") as container:
|
||||
assert container.streams.audio[0].rate == 44100
|
||||
|
||||
|
||||
def test_interleaved_video_uses_legacy_fallback_automatically(monkeypatch):
|
||||
calls = []
|
||||
path_logs = []
|
||||
@@ -383,6 +502,7 @@ def test_interleaved_video_uses_legacy_fallback_automatically(monkeypatch):
|
||||
assert "reason=non_contiguous_rgb_planes" in path_logs[0]
|
||||
assert "audio_present=False" in path_logs[0]
|
||||
assert "effective_audio_sample_rate=None" in path_logs[0]
|
||||
assert "effective_frame_conversion_workers=0" in path_logs[0]
|
||||
assert len(coerce_calls) == 1
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -89,12 +95,13 @@ class FakeAsyncOmni:
|
||||
yield MockVideoResult(videos)
|
||||
|
||||
|
||||
def test_raw_and_base64_encoders_receive_no_policy_config(mocker: MockerFixture):
|
||||
def test_raw_and_base64_encoders_receive_persistent_converter(mocker: MockerFixture):
|
||||
engine = FakeAsyncOmni()
|
||||
handler = OmniOpenAIServingVideo.for_diffusion(
|
||||
engine,
|
||||
model_name="test-model",
|
||||
)
|
||||
assert handler._video_frame_converter.max_workers == 8
|
||||
raw_encoder = mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
|
||||
return_value=b"encoded-video",
|
||||
@@ -111,8 +118,9 @@ def test_raw_and_base64_encoders_receive_no_policy_config(mocker: MockerFixture)
|
||||
|
||||
asyncio.run(_generate_both_response_types())
|
||||
|
||||
assert "encoding_config" not in raw_encoder.call_args.kwargs
|
||||
assert "encoding_config" not in base64_encoder.call_args.kwargs
|
||||
assert raw_encoder.call_args.kwargs["frame_converter"] is handler._video_frame_converter
|
||||
assert base64_encoder.call_args.kwargs["frame_converter"] is handler._video_frame_converter
|
||||
handler.shutdown()
|
||||
|
||||
|
||||
def test_resolve_diffusion_od_config_falls_back_to_attribute():
|
||||
@@ -828,6 +836,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
|
||||
|
||||
@@ -1213,8 +1233,9 @@ def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFi
|
||||
audio=None,
|
||||
audio_sample_rate=None,
|
||||
video_codec_options=None,
|
||||
frame_converter=None,
|
||||
):
|
||||
del video, fps, audio, video_codec_options
|
||||
del video, fps, audio, video_codec_options, frame_converter
|
||||
audio_sample_rates.append(audio_sample_rate)
|
||||
return b"fake-video"
|
||||
|
||||
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
from types import SimpleNamespace
|
||||
@@ -8,6 +11,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
@@ -77,7 +81,7 @@ def test_generate_submits_randomized_id_to_engine():
|
||||
"""Ensure the engine receives a UUID-suffixed ID, not the raw request ID"""
|
||||
|
||||
async def run():
|
||||
submitted_ids = []
|
||||
submitted_ids: list[str] = []
|
||||
omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids))
|
||||
|
||||
req_id = "my-req-1"
|
||||
@@ -106,8 +110,8 @@ def test_generate_forwards_lora_request_to_engine():
|
||||
"""
|
||||
|
||||
async def run():
|
||||
submitted_ids = []
|
||||
submitted_loras = []
|
||||
submitted_ids: list[str] = []
|
||||
submitted_loras: list[LoRARequest | None] = []
|
||||
omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids, submitted_loras))
|
||||
|
||||
lora = LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/fake")
|
||||
@@ -143,7 +147,7 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi
|
||||
user provided request ID will be aborted."""
|
||||
|
||||
async def run():
|
||||
aborted_batches = []
|
||||
aborted_batches: list[list[str]] = []
|
||||
omni = get_async_omni_instance(
|
||||
fake_abort_request=get_fake_abort(aborted_batches),
|
||||
)
|
||||
@@ -240,8 +244,8 @@ def test_abort_propagates_engine_errors_without_popping_state():
|
||||
@pytest.mark.cpu
|
||||
def test_generate_accepts_request_after_repeated_cancellations():
|
||||
async def run_test():
|
||||
submitted_request_ids = []
|
||||
aborted_request_batches = []
|
||||
submitted_request_ids: list[str] = []
|
||||
aborted_request_batches: list[list[str]] = []
|
||||
|
||||
async def collect_outputs(request_id):
|
||||
outputs = []
|
||||
@@ -286,6 +290,43 @@ def test_generate_accepts_request_after_repeated_cancellations():
|
||||
asyncio.run(run_test())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_generate_cancellation_converges_after_engine_shutdown_starts(mocker):
|
||||
async def run_test():
|
||||
submitted_request_ids: list[str] = []
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = True
|
||||
engine.request_queue = mocker.Mock()
|
||||
engine.num_stages = 1
|
||||
engine.add_request_async = get_fake_add_request(submitted_request_ids)
|
||||
|
||||
omni = get_async_omni_instance()
|
||||
omni.engine = engine
|
||||
|
||||
async def collect_outputs():
|
||||
async for _ in omni.generate(
|
||||
prompt={"prompt": "prompt"},
|
||||
request_id="cancel-shutdown",
|
||||
sampling_params_list=[SimpleNamespace()],
|
||||
output_modalities=["image"],
|
||||
):
|
||||
pass
|
||||
|
||||
task = asyncio.create_task(collect_outputs())
|
||||
await asyncio.sleep(0)
|
||||
task.cancel()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
assert len(submitted_request_ids) == 1
|
||||
assert submitted_request_ids[0].startswith("cancel-shutdown-")
|
||||
assert omni.request_states == {}
|
||||
engine.request_queue.sync_q.put.assert_not_called()
|
||||
|
||||
asyncio.run(run_test())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_generate_yields_streaming_diffusion_chunks_before_final():
|
||||
"""AsyncOmni.generate yields every intermediate diffusion chunk before the final one."""
|
||||
|
||||
Loaded 100 of 261 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user