Compare commits
31
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
26cef6e1ff | ||
|
|
596c16a550 | ||
|
|
90a08c0eda | ||
|
|
dbc0dd6d85 | ||
|
|
83d1f11449 | ||
|
|
4b66f62e10 | ||
|
|
f41451f64c | ||
|
|
bbe6ccc512 | ||
|
|
bbb709cfb5 | ||
|
|
c98517ce7f | ||
|
|
d77a35a32b | ||
|
|
877beeaa49 | ||
|
|
8b8d3f56db | ||
|
|
c57178c0e1 | ||
|
|
9a59475015 | ||
|
|
12f8168179 | ||
|
|
72a02b492a | ||
|
|
5358f5db01 | ||
|
|
e677360168 | ||
|
|
208c6776a8 | ||
|
|
fcf1085395 | ||
|
|
38a2500bdf | ||
|
|
7cd826d28c | ||
|
|
31384cb39c | ||
|
|
5be7541d01 | ||
|
|
9b0df3909a | ||
|
|
a15b5b4bcd | ||
|
|
3c3989ed04 | ||
|
|
854bff7fe5 | ||
|
|
b792578d89 | ||
|
|
bdb1d996be |
No files matched your search
@@ -345,7 +345,7 @@ steps:
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
|
||||
@@ -361,7 +361,7 @@ steps:
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
|
||||
@@ -377,7 +377,7 @@ steps:
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
|
||||
@@ -393,7 +393,7 @@ steps:
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
|
||||
@@ -409,7 +409,7 @@ steps:
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
|
||||
@@ -456,7 +456,7 @@ steps:
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx2_expansion.py -m "full_model and cuda" --run-level "full_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx.py -m "full_model and cuda" --run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test"
|
||||
@@ -477,10 +477,16 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_one_stage_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · LTX Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
|
||||
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx23_two_stage_layer_fused_i2v]' -m full_model --run-level full_model
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Perf Test"
|
||||
key: nightly-diffusion-x2v-performance
|
||||
timeout_in_minutes: 180
|
||||
|
||||
@@ -227,6 +227,18 @@ steps:
|
||||
"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Omni · MiniCPM-o 4.5 Perf Test"
|
||||
timeout_in_minutes: 60
|
||||
source_file_dependencies:
|
||||
- tests/dfx/perf/tests/test_minicpmo_4_5_ready.json
|
||||
- vllm_omni/deploy/minicpmo_4_5.yaml
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and cuda'
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "TTS · VoxCPM2 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_voxcpm2_tts.py
|
||||
|
||||
@@ -115,7 +115,7 @@ steps:
|
||||
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
|
||||
|
||||
@@ -17,6 +17,15 @@ steps:
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'core_model and npu and A3 and omni' --run-level 'core_model'
|
||||
- label: "Omni · MiniCPM-o 4.5 Perf Test"
|
||||
timeout_in_minutes: 60
|
||||
mirror_hardwares: a3_npu_2
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and npu and A3 and omni'
|
||||
|
||||
- group: ":card_index_dividers: TTS Model Test"
|
||||
key: ready-tts-group
|
||||
|
||||
@@ -40,9 +40,11 @@ steps:
|
||||
commands:
|
||||
- "bash .buildkite/release/scripts/publish-release-wheels-pypi.sh"
|
||||
|
||||
# Auto-skip when NIGHTLY=1 (scheduled); otherwise require manual unblock.
|
||||
- block: "Unblock to build release Docker images"
|
||||
depends_on: ~
|
||||
key: block-build-release-images
|
||||
if: build.env("NIGHTLY") != "1"
|
||||
|
||||
- group: "Build release Docker images"
|
||||
key: "build-release-images"
|
||||
@@ -63,7 +65,7 @@ steps:
|
||||
--progress plain \
|
||||
-f docker/Dockerfile.cuda .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
|
||||
# re-tag to default image tag and push, just in case arm64 build fails
|
||||
# Fallback default tag (no arch suffix) if arm64 build fails.
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT"
|
||||
env:
|
||||
@@ -86,6 +88,26 @@ steps:
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
# NIGHTLY=1 only: promote ECR images to DockerHub nightly tags, then prune.
|
||||
- label: "Publish nightly multi-arch image to DockerHub"
|
||||
depends_on:
|
||||
- build-release-image-x86
|
||||
- build-release-image-arm64
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
key: publish-nightly-images-dockerhub
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/release/scripts/publish-release-images.sh nightly"
|
||||
- "bash .buildkite/release/scripts/cleanup-nightly-builds.sh"
|
||||
plugins:
|
||||
- docker-login#v3.0.0:
|
||||
username: vllmbot
|
||||
password-env: DOCKERHUB_TOKEN
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
DOCKERHUB_USERNAME: "vllmbot"
|
||||
|
||||
- block: "Publish release images to DockerHub"
|
||||
key: block-publish-release-images
|
||||
depends_on:
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Clean up old nightly builds from DockerHub (vllm/vllm-omni), keeping the
|
||||
# newest 14 tags with the nightly- prefix.
|
||||
|
||||
set -ex
|
||||
|
||||
TAG_PREFIX="nightly-"
|
||||
REPO="vllm/vllm-omni"
|
||||
REPO_API_URL="https://hub.docker.com/v2/repositories/${REPO}/tags"
|
||||
|
||||
echo "Cleaning up tags with prefix: $TAG_PREFIX in repository: $REPO"
|
||||
|
||||
if [ -z "$DOCKERHUB_TOKEN" ]; then
|
||||
echo "Error: DOCKERHUB_TOKEN environment variable is not set"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [ -z "$DOCKERHUB_USERNAME" ]; then
|
||||
echo "Error: DOCKERHUB_USERNAME environment variable is not set"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Getting DockerHub bearer token..."
|
||||
set +x
|
||||
BEARER_TOKEN=$(curl -s -X POST \
|
||||
-H "Content-Type: application/json" \
|
||||
-d "{\"username\": \"$DOCKERHUB_USERNAME\", \"password\": \"$DOCKERHUB_TOKEN\"}" \
|
||||
"https://hub.docker.com/v2/users/login" | jq -r '.token')
|
||||
set -x
|
||||
|
||||
if [ -z "$BEARER_TOKEN" ] || [ "$BEARER_TOKEN" = "null" ]; then
|
||||
echo "Error: Failed to get DockerHub bearer token"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
get_all_tags() {
|
||||
local page=1
|
||||
local all_tags=""
|
||||
# Hub `name` is a partial match (not a strict prefix). Use it to narrow the
|
||||
# list, then enforce startswith() client-side. Exhaust API pages; do not
|
||||
# stop merely because a page has no matching tags after filtering.
|
||||
local name_query
|
||||
name_query=$(printf '%s' "$TAG_PREFIX" | jq -sRr @uri)
|
||||
|
||||
while true; do
|
||||
set +x
|
||||
local response
|
||||
response=$(curl -s -H "Authorization: Bearer $BEARER_TOKEN" \
|
||||
"$REPO_API_URL?page=$page&page_size=100&name=${name_query}")
|
||||
set -x
|
||||
|
||||
local page_count
|
||||
page_count=$(echo "$response" | jq -r '.results | length')
|
||||
if [ -z "$page_count" ] || [ "$page_count" = "0" ] || [ "$page_count" = "null" ]; then
|
||||
break
|
||||
fi
|
||||
|
||||
local tags
|
||||
tags=$(echo "$response" | jq -r --arg prefix "$TAG_PREFIX" \
|
||||
'.results[] | select(.name | startswith($prefix)) | "\(.last_updated)|\(.name)"')
|
||||
|
||||
if [ -n "$tags" ]; then
|
||||
all_tags="$all_tags$tags"$'\n'
|
||||
fi
|
||||
page=$((page + 1))
|
||||
done
|
||||
|
||||
echo "$all_tags" | sort -r | cut -d'|' -f2
|
||||
}
|
||||
|
||||
delete_tag() {
|
||||
local tag_name="$1"
|
||||
echo "Deleting tag: $tag_name"
|
||||
|
||||
local delete_url="https://hub.docker.com/v2/repositories/${REPO}/tags/$tag_name"
|
||||
set +x
|
||||
local response http_code body
|
||||
response=$(curl -s -w "\n%{http_code}" -X DELETE \
|
||||
-H "Authorization: Bearer $BEARER_TOKEN" "$delete_url")
|
||||
set -x
|
||||
http_code=$(printf '%s\n' "$response" | tail -n1)
|
||||
body=$(printf '%s\n' "$response" | sed '$d')
|
||||
|
||||
if [ "$http_code" -ge 200 ] && [ "$http_code" -lt 300 ]; then
|
||||
echo "Successfully deleted tag: $tag_name (HTTP ${http_code})"
|
||||
return 0
|
||||
fi
|
||||
|
||||
local detail=""
|
||||
if echo "$body" | jq -e '.detail' > /dev/null 2>&1; then
|
||||
detail=$(echo "$body" | jq -r '.detail')
|
||||
elif [ -n "$body" ]; then
|
||||
detail="$body"
|
||||
else
|
||||
detail="(empty response body)"
|
||||
fi
|
||||
echo "Error: Failed to delete tag $tag_name (HTTP ${http_code}): ${detail}"
|
||||
return 1
|
||||
}
|
||||
|
||||
echo "Fetching all tags from DockerHub..."
|
||||
all_tags=$(get_all_tags)
|
||||
|
||||
if [ -z "$all_tags" ]; then
|
||||
echo "No tags found to clean up"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
total_tags=$(echo "$all_tags" | wc -l)
|
||||
echo "Found $total_tags tags"
|
||||
|
||||
tags_to_keep=14
|
||||
tags_to_delete=$((total_tags - tags_to_keep))
|
||||
|
||||
if [ $tags_to_delete -le 0 ]; then
|
||||
echo "No tags need to be deleted (only $total_tags tags found, keeping $tags_to_keep)"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
echo "Will delete $tags_to_delete old tags, keeping the newest $tags_to_keep"
|
||||
|
||||
tags_to_delete_list=$(echo "$all_tags" | tail -n +$((tags_to_keep + 1)))
|
||||
|
||||
if [ -z "$tags_to_delete_list" ]; then
|
||||
echo "No tags to delete"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
echo "Deleting old tags..."
|
||||
delete_failures=0
|
||||
while IFS= read -r tag; do
|
||||
if [ -n "$tag" ]; then
|
||||
if ! delete_tag "$tag"; then
|
||||
delete_failures=$((delete_failures + 1))
|
||||
fi
|
||||
sleep 1
|
||||
fi
|
||||
done <<< "$tags_to_delete_list"
|
||||
|
||||
if [ "$delete_failures" -gt 0 ]; then
|
||||
echo "Cleanup finished with ${delete_failures} deletion failure(s)"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Cleanup completed successfully"
|
||||
@@ -1,50 +1,87 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Publish release Docker images from ECR to DockerHub.
|
||||
# Pulls per-arch images, tags with `latest` and versioned tags, pushes them,
|
||||
# then creates and pushes a multi-arch manifest.
|
||||
# Publish Docker images from ECR staging to DockerHub.
|
||||
#
|
||||
# Usage:
|
||||
# publish-release-images.sh # release mode (default)
|
||||
# publish-release-images.sh release
|
||||
# publish-release-images.sh nightly
|
||||
#
|
||||
# Release tags: latest, v<version>
|
||||
# Nightly tags: nightly, nightly-<commit>
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
|
||||
if [ -z "${RELEASE_VERSION}" ]; then
|
||||
echo "ERROR: release-version metadata not set"
|
||||
exit 1
|
||||
fi
|
||||
MODE="${1:-release}"
|
||||
|
||||
COMMIT="$BUILDKITE_COMMIT"
|
||||
DOCKERHUB_REPO="vllm/vllm-omni"
|
||||
ECR_REPO="public.ecr.aws/q9t5s3a7/vllm-omni-release-repo"
|
||||
COMMIT="$BUILDKITE_COMMIT"
|
||||
PRIMARY_TAGS=()
|
||||
EXTRA_MANIFEST_TAGS=()
|
||||
|
||||
case "${MODE}" in
|
||||
release)
|
||||
RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
|
||||
if [ -z "${RELEASE_VERSION}" ]; then
|
||||
echo "ERROR: release-version metadata not set"
|
||||
exit 1
|
||||
fi
|
||||
PRIMARY_TAGS=("latest" "v${RELEASE_VERSION}")
|
||||
;;
|
||||
nightly)
|
||||
# Commit-pinned manifest reuses nightly-{x86_64,aarch64} arch images.
|
||||
PRIMARY_TAGS=("nightly")
|
||||
EXTRA_MANIFEST_TAGS=("nightly-${COMMIT}")
|
||||
;;
|
||||
*)
|
||||
echo "Usage: $0 {release|nightly}"
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
|
||||
echo "========================================"
|
||||
echo "Publishing release images v${RELEASE_VERSION}"
|
||||
echo "Publishing ${MODE} images to ${DOCKERHUB_REPO}"
|
||||
echo " Commit: ${COMMIT}"
|
||||
echo " Primary tags: ${PRIMARY_TAGS[*]}"
|
||||
if [ ${#EXTRA_MANIFEST_TAGS[@]} -gt 0 ]; then
|
||||
echo " Extra manifests: ${EXTRA_MANIFEST_TAGS[*]}"
|
||||
fi
|
||||
echo "========================================"
|
||||
|
||||
# Login to ECR to pull staging images
|
||||
aws ecr-public get-login-password --region us-east-1 | \
|
||||
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
|
||||
|
||||
docker pull "${ECR_REPO}:${COMMIT}-x86_64"
|
||||
docker pull "${ECR_REPO}:${COMMIT}-aarch64"
|
||||
SRC_X86="${ECR_REPO}:${COMMIT}-x86_64"
|
||||
SRC_ARM="${ECR_REPO}:${COMMIT}-aarch64"
|
||||
|
||||
docker tag "${ECR_REPO}:${COMMIT}-x86_64" "${DOCKERHUB_REPO}:latest-x86_64"
|
||||
docker tag "${ECR_REPO}:${COMMIT}-x86_64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64"
|
||||
docker push "${DOCKERHUB_REPO}:latest-x86_64"
|
||||
docker push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64"
|
||||
docker pull "${SRC_X86}"
|
||||
docker pull "${SRC_ARM}"
|
||||
|
||||
docker tag "${ECR_REPO}:${COMMIT}-aarch64" "${DOCKERHUB_REPO}:latest-aarch64"
|
||||
docker tag "${ECR_REPO}:${COMMIT}-aarch64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
|
||||
docker push "${DOCKERHUB_REPO}:latest-aarch64"
|
||||
docker push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
|
||||
publish_manifest() {
|
||||
local tag="$1"
|
||||
local arch_tag_base="$2"
|
||||
docker manifest rm "${DOCKERHUB_REPO}:${tag}" || true
|
||||
docker manifest create \
|
||||
"${DOCKERHUB_REPO}:${tag}" \
|
||||
"${DOCKERHUB_REPO}:${arch_tag_base}-x86_64" \
|
||||
"${DOCKERHUB_REPO}:${arch_tag_base}-aarch64"
|
||||
docker manifest push "${DOCKERHUB_REPO}:${tag}"
|
||||
}
|
||||
|
||||
docker manifest rm "${DOCKERHUB_REPO}:latest" || true
|
||||
docker manifest rm "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" || true
|
||||
docker manifest create "${DOCKERHUB_REPO}:latest" "${DOCKERHUB_REPO}:latest-x86_64" "${DOCKERHUB_REPO}:latest-aarch64"
|
||||
docker manifest create "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
|
||||
docker manifest push "${DOCKERHUB_REPO}:latest"
|
||||
docker manifest push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}"
|
||||
for tag in "${PRIMARY_TAGS[@]}"; do
|
||||
docker tag "${SRC_X86}" "${DOCKERHUB_REPO}:${tag}-x86_64"
|
||||
docker tag "${SRC_ARM}" "${DOCKERHUB_REPO}:${tag}-aarch64"
|
||||
docker push "${DOCKERHUB_REPO}:${tag}-x86_64"
|
||||
docker push "${DOCKERHUB_REPO}:${tag}-aarch64"
|
||||
publish_manifest "${tag}" "${tag}"
|
||||
done
|
||||
|
||||
ARCH_BASE="${PRIMARY_TAGS[0]}"
|
||||
for tag in "${EXTRA_MANIFEST_TAGS[@]}"; do
|
||||
publish_manifest "${tag}" "${ARCH_BASE}"
|
||||
done
|
||||
|
||||
echo ""
|
||||
echo "Successfully published release images for v${RELEASE_VERSION}"
|
||||
echo "Successfully published ${MODE} images"
|
||||
@@ -1,6 +1,6 @@
|
||||
---
|
||||
name: precheck-pr
|
||||
description: Self-check your branch before creating a PR — catch dead code, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub.
|
||||
description: Self-check your branch before creating a PR — catch dead code, prevent new model-specific Python examples, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub.
|
||||
---
|
||||
|
||||
# PR Pre-Check
|
||||
@@ -65,6 +65,8 @@ Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type
|
||||
|
||||
**Also run the Code-Quality sweep on every PR**, regardless of type or mode: the five diff-scoped checks in [references/code-quality.md](references/code-quality.md) — kwargs fragility, broad-except swallow, `Any`/wrong type hints, hot-path `.clone()`/`deepcopy`, and event-loop blocking — plus the advisory conventions (log level, structured logging, synchronization, cleanup, dependencies, naming) in [checklists.md](references/checklists.md). These count only lines the PR *adds* — the pre-existing backlog across the repo is out of scope.
|
||||
|
||||
**Also run the Examples-Policy check on every PR** using [references/examples-policy.md](references/examples-policy.md). Inspect only Python paths introduced by the diff. Treat a new model-, checkpoint-, vendor-, or family-specific Python example as blocking; do not report pre-existing example debt.
|
||||
|
||||
### Step 5: Print Report
|
||||
|
||||
```
|
||||
@@ -77,6 +79,7 @@ Pre-check report for <branch>
|
||||
───────────────── ──────
|
||||
PR title format ✓
|
||||
Code quality ⚠ 1 broad except, 2 Any hints
|
||||
Examples policy ✗ new model-specific Python example
|
||||
PR desc integrity ✓
|
||||
Registry/config ✓
|
||||
Dead code ⚠ 2 warnings
|
||||
|
||||
@@ -47,6 +47,19 @@ Diff-scoped sweep for five fragility patterns that are pervasive in the existing
|
||||
|
||||
Roll the ⚠/✗ counts into the report as a single **Code quality** dimension row.
|
||||
|
||||
---
|
||||
|
||||
## All PRs: Examples Policy (run on every PR, quick & full)
|
||||
|
||||
Inspect Python files introduced under `examples/` using
|
||||
[examples-policy.md](examples-policy.md). This is a diff-scoped ratchet: existing
|
||||
model-specific examples may be edited or removed without triggering the check.
|
||||
|
||||
- [ ] **No new model-specific Python examples:** the diff adds, copies, or renames no Python file whose path or behavior is specific to one model, checkpoint, vendor, or model family. ✗ for any violation. Route inference through a shared task/protocol entrypoint, model request behavior through `vllm_omni/model_extras`, and commands through task docs or `recipes/`.
|
||||
- [ ] **New generic example code is justified:** any new task-neutral Python file under `examples/` is an actual user-facing entrypoint rather than reusable production logic, a benchmark, an app, a setup tool, or a reproducer. ⚠ when the placement is questionable.
|
||||
|
||||
Roll the result into the report as a separate **Examples policy** dimension.
|
||||
|
||||
### Conventions (eyeball — no grep)
|
||||
|
||||
Judgment calls that don't grep cleanly — apply them while reading the diff, not as a mechanical check:
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
# Examples Python Policy
|
||||
|
||||
Prevent new model-specific Python examples without forcing cleanup of the
|
||||
existing backlog. Run this check on every PR in quick and full modes.
|
||||
|
||||
## Scope
|
||||
|
||||
Set the merge base as in the main workflow, then list Python paths introduced
|
||||
by the PR, including additions, copies, and renames:
|
||||
|
||||
```bash
|
||||
BASE="$(git merge-base HEAD origin/main)"
|
||||
git diff --name-status --diff-filter=ACR --find-renames --find-copies \
|
||||
"${BASE}"...HEAD -- 'examples/**/*.py'
|
||||
```
|
||||
|
||||
Inspect destination paths for copies and renames. Existing model-specific files
|
||||
that are only modified or deleted are grandfathered and must not be reported.
|
||||
|
||||
## Blocking classification
|
||||
|
||||
Mark ✗ when a newly introduced Python file does any of the following:
|
||||
|
||||
- Names or scopes a path to one model, checkpoint, vendor, or model family.
|
||||
- Hard-codes a model identifier, model-specific prompt, request shape, output
|
||||
conversion, or launch configuration behind an otherwise generic filename.
|
||||
- Duplicates an existing task/protocol runner to demonstrate one model.
|
||||
- Adds a per-model helper below a generic task hub such as `text_to_speech/`.
|
||||
|
||||
Do not infer safety from a generic filename alone. For example, a file named
|
||||
`text_to_audio.py` is still model-specific if it only implements one model's
|
||||
contract.
|
||||
|
||||
## Acceptable additions
|
||||
|
||||
A new Python example may pass when it is a genuinely model-neutral,
|
||||
user-facing task or protocol entrypoint. It must accept model selection through
|
||||
configuration and keep model-specific behavior out of the script.
|
||||
|
||||
Use ⚠ when the file is model-neutral but reusable logic appears misplaced under
|
||||
`examples/`. Route code by purpose:
|
||||
|
||||
| Code purpose | Preferred location |
|
||||
|---|---|
|
||||
| Model prompt, defaults, request/output adaptation | `vllm_omni/model_extras/` or production model modules |
|
||||
| Runnable model command and validation evidence | Task documentation or `recipes/` |
|
||||
| UI or interactive application | `apps/` |
|
||||
| Benchmark or evaluator | `benchmarks/` |
|
||||
| Regression reproducer | `tests/` |
|
||||
| Download, conversion, or setup utility | `tools/` |
|
||||
|
||||
If a model cannot use an existing shared runner, report the missing generic
|
||||
capability. Do not solve the gap by adding a model-specific Python example.
|
||||
|
||||
## Report format
|
||||
|
||||
Add one row to the pre-check report:
|
||||
|
||||
```text
|
||||
Examples policy ✓ no new Python example paths
|
||||
Examples policy ⚠ generic helper belongs in tools/
|
||||
Examples policy ✗ examples/offline_inference/new_model/end2end.py is model-specific
|
||||
Examples policy ✗ examples/offline_inference/x_to_y_model_name.py is model-specific
|
||||
```
|
||||
|
||||
For a blocker, name the path, the model-specific behavior, and the appropriate
|
||||
shared runner or destination. Do not propose deleting unrelated existing
|
||||
examples as part of the PR.
|
||||
@@ -97,6 +97,7 @@ If docs and live code disagree, verify the code/tests and report the drift.
|
||||
| [test-quality-evaluation.md](references/checks/test-quality-evaluation.md) | Tests change, are absent for risky code, or may not exercise production behavior. |
|
||||
| [tests-docs-checklist.md](references/checks/tests-docs-checklist.md) | Coverage, CI markers, examples, user docs, or PR evidence need review. |
|
||||
| [verification.md](references/checks/verification.md) | Hardware, a server, or a runnable affected path is available for active verification. |
|
||||
| [examples-policy.md](../precheck-pr/references/examples-policy.md) | The PR adds, copies, or renames Python under `examples/`; apply the canonical policy shared with `precheck-pr`. |
|
||||
|
||||
### Delivery and reviewer coordination
|
||||
|
||||
@@ -152,6 +153,13 @@ contract they protect or use only the applicable evidence checks.
|
||||
Apply every category in [general-checks.md](references/process/general-checks.md) before
|
||||
lower-priority comments.
|
||||
|
||||
If the diff census contains an added, copied, or renamed Python path under
|
||||
`examples/`, read and apply the canonical
|
||||
[examples policy](../precheck-pr/references/examples-policy.md). Treat a new
|
||||
model-specific Python example as blocking. Do not flag model-specific example
|
||||
debt that the PR only modifies or removes, and do not run the rest of the
|
||||
author-oriented `precheck-pr` workflow.
|
||||
|
||||
For each changed value or behavior, trace:
|
||||
|
||||
```text
|
||||
|
||||
+7
-2
@@ -11,10 +11,15 @@ build:
|
||||
post_checkout:
|
||||
- git fetch --unshallow || true
|
||||
pre_install:
|
||||
# `pre_install` runs before RTD upgrades pip, and the seeded pip 23.1
|
||||
# rejects the `typing_extensions` wheel torch needs (fixed in pip 23.2).
|
||||
- python -m pip install --upgrade pip
|
||||
# Pre-install CPU torch so the subsequent `pip install .[docs]` does not
|
||||
# pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
|
||||
# 1 GB memory cap and time out the mkdocs build.
|
||||
- python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
|
||||
# 1 GB memory cap and time out the mkdocs build. Wheel-only: `--index-url`
|
||||
# replaces PyPI, so an sdist fallback cannot reach its build deps. The
|
||||
# quotes keep YAML from reading `:all:` as a mapping key.
|
||||
- "python -m pip install --only-binary=:all: --index-url https://download.pytorch.org/whl/cpu torch"
|
||||
|
||||
mkdocs:
|
||||
configuration: mkdocs.yml
|
||||
|
||||
@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
||||
---
|
||||
|
||||
*Latest News* 🔥
|
||||
- [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0).
|
||||
- [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support.
|
||||
- [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
|
||||
- [2026/06] Starting with [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0), vLLM-Omni publishes a stable release aligned with every even-numbered upstream vLLM minor version. [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0), [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0), [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0), and [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) continued this cadence, expanding omni and world-model support with [NVIDIA Cosmos3](recipes/cosmos3/Cosmos3-Nano.md) and DreamZero, adding models such as MiniCPM-o 4.5, MOSS-TTS, and Lance, and advancing TTS, diffusion, distributed execution, quantization, RL integration through [VeRL-Omni](https://github.com/verl-project/verl-omni), and CUDA/ROCm/MUSA/NPU/XPU coverage.
|
||||
|
||||
@@ -158,7 +158,6 @@ def _build_omni_kwargs(args, quantization=None):
|
||||
def _generate_image(omni, args, prompt, seed):
|
||||
"""Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
|
||||
@@ -176,7 +175,7 @@ def _generate_image(omni, args, prompt, seed):
|
||||
elapsed = time.perf_counter() - start
|
||||
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
|
||||
|
||||
req_out = OmniRequestOutput.unwrap_result(outputs)
|
||||
req_out = outputs[0] if isinstance(outputs, list) else outputs
|
||||
if not req_out.images:
|
||||
raise ValueError("Could not extract image output from result.")
|
||||
img = req_out.images[0]
|
||||
@@ -206,15 +205,15 @@ def _generate_video(omni, args, prompt, seed):
|
||||
elapsed = time.perf_counter() - start
|
||||
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
|
||||
|
||||
first = outputs[0]
|
||||
if hasattr(first, "request_output") and isinstance(first.request_output, list):
|
||||
inner = first.request_output[0]
|
||||
if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
|
||||
frames = inner.images[0] if inner.images else None
|
||||
else:
|
||||
frames = inner
|
||||
elif hasattr(first, "images") and first.images:
|
||||
frames = first.images[0]
|
||||
if isinstance(outputs, list) and isinstance(outputs[0], OmniRequestOutput):
|
||||
first = outputs[0]
|
||||
elif isinstance(outputs, OmniRequestOutput):
|
||||
first = outputs
|
||||
else:
|
||||
raise ValueError("Could not extract video frames from output.")
|
||||
|
||||
if hasattr(first, "images"):
|
||||
frames = first.images[0] if first.images else None
|
||||
else:
|
||||
raise ValueError("Could not extract video frames from output.")
|
||||
|
||||
|
||||
@@ -300,7 +300,7 @@ def benchmark(args: argparse.Namespace) -> None:
|
||||
output_idx = 0
|
||||
for stage_outputs in omni.generate(all_prompts, sampling_params_list, py_generator=True):
|
||||
if stage_outputs.final_output_type == "image":
|
||||
request_output = stage_outputs.request_output
|
||||
request_output = stage_outputs
|
||||
request_id = getattr(request_output, "request_id", "")
|
||||
|
||||
images = getattr(request_output, "images", [])
|
||||
|
||||
@@ -167,9 +167,11 @@ def main() -> None:
|
||||
]
|
||||
native_cmd = [
|
||||
sys.executable,
|
||||
"examples/offline_inference/text_to_video/text_to_video_lingbot.py",
|
||||
"examples/offline_inference/text_to_video/text_to_video.py",
|
||||
"--model",
|
||||
args.model,
|
||||
"--model-class-name",
|
||||
"LingBotVideoPipeline",
|
||||
"--prompt",
|
||||
args.prompt,
|
||||
"--output",
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.26.0
|
||||
ARG VLLM_ASCEND_TAG=v0.27.0
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.26.0-a3
|
||||
ARG VLLM_ASCEND_TAG=v0.27.0-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.26.0
|
||||
ARG VLLM_ASCEND_TAG=v0.27.0
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.26.0-a3
|
||||
ARG VLLM_ASCEND_TAG=v0.27.0-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
|
||||
+2
-1
@@ -163,7 +163,8 @@ nav:
|
||||
- CLI Reference: cli
|
||||
- Community:
|
||||
- Governance: community/governance.md
|
||||
- community/*
|
||||
- Contact Us: community/contact_us.md
|
||||
- Community Meetings: community/meetings.md
|
||||
- Slack: https://slack.vllm.ai
|
||||
- Blog: https://blog.vllm.ai
|
||||
- Forum: https://discuss.vllm.ai
|
||||
@@ -1,5 +1,9 @@
|
||||
# Contact Us
|
||||
|
||||
- For technical questions and feature requests, please use GitHub [Issues](https://github.com/vllm-project/vllm-omni/issues)
|
||||
- For coordinating contributions and development and discussing with other users and developers, please join `sig-omni` channel in our [Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/)
|
||||
- For security disclosures, please use GitHub's [Security Advisories](https://github.com/vllm-project/vllm-omni/security/advisories) feature
|
||||
- For bugs and feature requests, search or open a GitHub [issue](https://github.com/vllm-project/vllm-omni/issues).
|
||||
- For technical and usage questions, join the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/).
|
||||
- For coordinating contributions and development, use `sig-omni`, the forum, or join a [community meeting](meetings.md).
|
||||
|
||||
## Security Disclosures
|
||||
|
||||
Do not share vulnerability details in public GitHub issues, pull requests, discussions, or Slack channels. To arrange a private report, send a direct message to one of the [lead maintainers](governance.md#lead-maintainers) in [vLLM Slack](https://slack.vllm.ai/). If you cannot send a direct message, you may post a detail-free request in `sig-omni` asking a lead maintainer to contact you privately; do not include any vulnerability details in that request.
|
||||
@@ -52,31 +52,11 @@ Committers have write access and merge rights. They typically have deep expertis
|
||||
- [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
|
||||
- [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
|
||||
|
||||
## Diffusion Workload Division
|
||||
## Reviewer Routing
|
||||
|
||||
This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review.
|
||||
For path-based reviewer routing, consult [`.github/CODEOWNERS`](https://github.com/vllm-project/vllm-omni/blob/main/.github/CODEOWNERS). Detailed subsystem boundaries and safe-change guidance live in the [module design documents](../design/index.md#module-design-documents).
|
||||
|
||||
### Runtime-related
|
||||
|
||||
- Diffusion models: @RuixiangMa, @wtomin, @xuechendi
|
||||
|
||||
### Optimization
|
||||
|
||||
- Parallel: @RuixiangMa, @wtomin
|
||||
- Attention backend: @ZJY0516
|
||||
- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
|
||||
- Offload: @yuanheng-zhao
|
||||
- Intel GPU support: @xuechendi
|
||||
- Quantization: @lishunyang12, @david6666666
|
||||
- Kernel fusion / communication-computation: @ZJY0516
|
||||
|
||||
### Docs & Test
|
||||
|
||||
- Diffusion docs: @RuixiangMa, @wtomin
|
||||
|
||||
## Meetings
|
||||
|
||||
Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
|
||||
`CODEOWNERS` is operational routing metadata, not a governance or approval policy. Being listed does not grant committer status or merge rights, and it does not mean every listed reviewer must approve a change. The maintainer roster above is authoritative for project roles and merge rights.
|
||||
|
||||
## Committer Nomination Process
|
||||
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
# Community Meetings
|
||||
|
||||
vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. Schedules and access details can change, so use these maintained documents as the source of truth:
|
||||
|
||||
- [Chinese-language meeting document](https://docs.google.com/document/d/1pdUBiS_7mdOUNDtdwy-9OUf7jsMWN324BIbps5olDME/edit?tab=t.0)
|
||||
- [English-language meeting document](https://docs.google.com/document/d/19AlD08T8uJ7wqfs8VtYCi0iGMTwlTYVVlAxFZ3S7OIg/edit?tab=t.0)
|
||||
|
||||
Each document contains the current schedule, access details, agendas, and past meeting notes. For broader vLLM community events, see the [vLLM events calendar](https://vllm.ai/events/).
|
||||
@@ -1 +0,0 @@
|
||||
# Meetups
|
||||
@@ -1,12 +0,0 @@
|
||||
# Volunteers for Bugfix and CI
|
||||
|
||||
We encourage you to check current docs and [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions for your questions. If non of these can solve it, please propose an issue to describe your questions about bug or CI problems for developing.
|
||||
|
||||
If you have urgent need for locating and solving bugfix or CI problems, please find community volunteers below.
|
||||
|
||||
| Dec 4-Dec 12 | Dec 15-Dec 19 | Dec 22-Dec 26 | Dec 29- Jan 2, 2026| Jan 5-Jan 9 | Jan 12-Jan 16 |
|
||||
|----------|----------|----------|----------|----------|----------|
|
||||
| <a href="https://github.com/congw729">Conw729</a> | <a href="https://github.com/yinpeiqi">yinpeiqi</a> | <a href="https://github.com/tzhouam">tzhouam</a> | <a href="https://github.com/SamitHuang">SamitHuang</a> | <a href="https://github.com/gcanlin">gcanlin</a> | <a href="https://github.com/natureofnature">natureofnature</a> |
|
||||
| <a href="https://github.com/david6666666">david6666666</a> | <a href="https://github.com/R2-Y">R2-Y</a> | <a href="https://github.com/hsliuustc0106">hsliuustc0106</a> | <a href="https://github.com/Gaohan123">Gaohan123</a> | <a href="https://github.com/ZJY0516">ZJY0516</a> | <a href="https://github.com/qibaoyuan">qibaoyuan</a> |
|
||||
|
||||
We kindly welcome more contributors to fix bugs and contribute new features!
|
||||
@@ -3,7 +3,7 @@
|
||||
Thank you for your interest in contributing to vLLM-Omni! This document provides guidelines and instructions for contributing.
|
||||
|
||||
!!! note
|
||||
We host weekly developer-facing online meetings to discuss milestones and updates **every Tuesday at 19:30 PDT**. Meeting link as well as the past meeting notes can be found [here](https://tinyurl.com/vllm-omni-meeting).
|
||||
vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. See [Community Meetings](../community/meetings.md) for current schedules, access details, agendas, and past notes.
|
||||
|
||||
## Getting Started
|
||||
|
||||
@@ -122,7 +122,7 @@ Repository skills guide the agent, but they do not replace contributor judgment,
|
||||
If you encounter a bug or have a feature request, please search existing issues first to see if it has already been reported. If not, please file a new issue, providing as much relevant information as possible.
|
||||
|
||||
!!! important
|
||||
If you discover a security vulnerability, please report it by creating a GitHub issue with the `security` label.
|
||||
Do not report suspected security vulnerabilities through a public issue, pull request, discussion, or Slack channel. Follow the [security disclosure instructions](../community/contact_us.md#security-disclosures) to arrange a private report.
|
||||
|
||||
## Pull Requests & Code Reviews
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# Prometheus Metrics Design
|
||||
|
||||
This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. This iteration covers the text and audio surface; diffusion / image / video families land in a follow-up.
|
||||
This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics.
|
||||
|
||||
## Objectives
|
||||
|
||||
|
||||
@@ -42,7 +42,8 @@ th {
|
||||
| `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` | ✅︎ | | | | — |
|
||||
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
|
||||
| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | — |
|
||||
| `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 | ✅︎ | ✅︎ | | | — |
|
||||
| `LTX2DistilledPipeline` | LTX-2 / LTX-2.3 full-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + `Lightricks/LTX-2.3` upsampler | ✅︎ | ✅︎ | | | — |
|
||||
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
|
||||
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
|
||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
@@ -58,6 +59,7 @@ th {
|
||||
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
|
||||
+3
-3
@@ -10,12 +10,12 @@ A: Refer to [GPU memory calculation and configuration](../configuration/gpu_memo
|
||||
|
||||
> Q: I encounter some bugs or CI problems, which is urgent. How can I solve it?
|
||||
|
||||
A: At first, you can check current [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions. If none of these satisfy your demand and it is urgent, please find these [volunteers](https://docs.vllm.ai/projects/vllm-omni/en/latest/community/volunteers/) for help.
|
||||
A: Search the current [issues](https://github.com/vllm-project/vllm-omni/issues) first. If no existing issue matches a bug or CI failure, open a [bug report](https://github.com/vllm-project/vllm-omni/issues/new?template=400-bug-report.yml) with a minimal reproducer, environment details, and relevant logs. For technical or usage questions, use the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or the [vLLM Forum](https://discuss.vllm.ai/).
|
||||
|
||||
> Q: Does vLLM-Omni support AWQ or any other quantization?
|
||||
|
||||
A: We plan to introduce GGUF FP8 prequantized models and online FP8 quantization in version 0.16.0. Support for other quantization types will follow in future releases. For details, please see our [Q1 quantization roadmap](https://github.com/vllm-project/vllm-omni/issues/1057).
|
||||
A: AWQ is available through the inherited vLLM quantization registry, and vLLM-Omni supports several additional methods. Actual support and validation vary by model, component, hardware, and checkpoint. See the [quantization overview](../user_guide/quantization/overview.md) for validated combinations and method-specific guides.
|
||||
|
||||
> Q: Does vLLM-Omni support multimodal streaming input and output?
|
||||
|
||||
A: Not yet. We already put it on the [Roadmap](https://github.com/vllm-project/vllm-omni/issues/165). Please stay tuned!
|
||||
A: Yes, for specific models and serving paths. Current capabilities include [streaming video input](../serving/video_stream_api.md), [streaming diffusion output](../user_guide/diffusion/execution_modes.md#streaming-output), and experimental full-duplex real-time audio described in the [project overview](../README.md). Check the model and API documentation because input and output support is model-specific.
|
||||
+50
-26
@@ -1,23 +1,23 @@
|
||||
# Production Metrics
|
||||
|
||||
vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. This page covers the text and audio surface; diffusion / image / video metrics are tracked in a follow-up PR.
|
||||
vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server.
|
||||
|
||||
```bash
|
||||
vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats
|
||||
curl http://localhost:8000/metrics
|
||||
```
|
||||
|
||||
**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm:omni_*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
|
||||
**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm_omni:*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
|
||||
|
||||
## Metric Namespaces
|
||||
|
||||
| Prefix | Source | Present when |
|
||||
|--------|--------|--------------|
|
||||
| `vllm:omni_` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
|
||||
| `vllm_omni:` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
|
||||
| `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage |
|
||||
| `http_` / `process_` | Uvicorn / Python runtime | Always |
|
||||
|
||||
## Pipeline-Level Metrics (`vllm:omni_`)
|
||||
## Pipeline-Level Metrics (`vllm_omni:`)
|
||||
|
||||
Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline.
|
||||
|
||||
@@ -25,42 +25,66 @@ Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the
|
||||
|
||||
| Metric | Type | Labels | Description |
|
||||
|--------|------|--------|-------------|
|
||||
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
|
||||
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
|
||||
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
|
||||
| `vllm_omni:num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
|
||||
| `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
|
||||
| `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
|
||||
|
||||
### Latency
|
||||
|
||||
| Metric | Type | Labels | Description |
|
||||
|--------|------|--------|-------------|
|
||||
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
|
||||
| `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
|
||||
|
||||
## Audio Modality Metrics (`vllm:omni_`)
|
||||
## Audio Modality Metrics (`vllm_omni:`)
|
||||
|
||||
Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label.
|
||||
|
||||
| Metric | Type | Extra label | Description |
|
||||
|--------|------|-------------|-------------|
|
||||
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
|
||||
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
|
||||
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
|
||||
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
|
||||
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
|
||||
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
|
||||
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
|
||||
| `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
|
||||
| `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
|
||||
| `vllm_omni:audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
|
||||
| `vllm_omni:audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
|
||||
| `vllm_omni:audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
|
||||
| `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
|
||||
| `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
|
||||
|
||||
The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition.
|
||||
|
||||
## Cross-Stage Transfer Metrics (`vllm:omni_`)
|
||||
## Diffusion Metrics (`vllm_omni:`)
|
||||
|
||||
Per-request timing breakdowns for diffusion (image/video) stages. Emitted at request finalize when `stage_metrics.diffusion_metrics` is present. All carry `{model_name, stage, replica}`.
|
||||
|
||||
| Metric | Type | Description |
|
||||
|--------|------|-------------|
|
||||
| `vllm_omni:diffusion_exec_s` | Histogram | DiT forward pass execution time per request in seconds |
|
||||
| `vllm_omni:diffusion_exec_per_step_s` | Histogram | DiT forward pass execution time per denoising step in seconds (`exec_s / num_inference_steps`) |
|
||||
| `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time per request in seconds |
|
||||
| `vllm_omni:diffusion_postprocess_s` | Histogram | Diffusion output postprocessing (VAE decode) time per request in seconds |
|
||||
|
||||
`diffusion_exec_s` uses `SECONDS_BUCKETS` (50 ms → 300 s); the other three use `SECONDS_FAST_BUCKETS` (1 ms → 60 s). `diffusion_exec_per_step_s` is skipped when `num_inference_steps` is not set.
|
||||
|
||||
### Example PromQL
|
||||
|
||||
```promql
|
||||
# P99 DiT exec time by stage
|
||||
histogram_quantile(0.99, rate(vllm_omni:diffusion_exec_s_bucket[5m]))
|
||||
|
||||
# Ratio of preprocess to total exec wall time
|
||||
rate(vllm_omni:diffusion_preprocess_s_sum[5m])
|
||||
/ rate(vllm_omni:diffusion_exec_s_sum[5m])
|
||||
```
|
||||
|
||||
## Cross-Stage Transfer Metrics (`vllm_omni:`)
|
||||
|
||||
Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed.
|
||||
|
||||
| Metric | Type | Description |
|
||||
|--------|------|-------------|
|
||||
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
|
||||
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
|
||||
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
|
||||
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
|
||||
| `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
|
||||
| `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
|
||||
| `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
|
||||
| `vllm_omni:transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
|
||||
|
||||
## vLLM Engine Metrics (`vllm:`)
|
||||
|
||||
@@ -83,9 +107,9 @@ For the full list of upstream metrics, see [the vLLM docs](https://github.com/vl
|
||||
|
||||
| Metric group | Multi-stage LLM (Qwen3-Omni) |
|
||||
|---|---|
|
||||
| `vllm:omni_` request tracking + latency | With `--log-stats` |
|
||||
| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage |
|
||||
| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
|
||||
| `vllm_omni:` request tracking + latency | With `--log-stats` |
|
||||
| `vllm_omni:` audio modality | With `--log-stats`, if pipeline has a talker stage |
|
||||
| `vllm_omni:` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
|
||||
| `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` |
|
||||
| `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` |
|
||||
|
||||
@@ -94,5 +118,5 @@ For the full list of upstream metrics, see [the vLLM docs](https://github.com/vl
|
||||
- All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values.
|
||||
- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
|
||||
- Sizes use the `_bytes` suffix.
|
||||
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
|
||||
- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label.
|
||||
- All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm_omni:`.
|
||||
- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm_omni:audio_ttfp_s` for audio) rather than a single metric with a `modality` label.
|
||||
@@ -12,6 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back
|
||||
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|
||||
|---|---|---|---|---|---|
|
||||
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
|
||||
| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) |
|
||||
| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
|
||||
@@ -35,6 +36,25 @@ python text_to_video.py \
|
||||
--output t2v_out.mp4
|
||||
```
|
||||
|
||||
### LingBot-Video
|
||||
|
||||
The shared runner recognizes both official LingBot checkpoint IDs and selects
|
||||
the native `LingBotVideoPipeline` automatically:
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model robbyant/lingbot-video-dense-1.3b \
|
||||
--prompt "a robotic arm picks up a red block" \
|
||||
--height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
|
||||
--guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
|
||||
--output lingbot_t2v.mp4
|
||||
```
|
||||
|
||||
Use `robbyant/lingbot-video-moe-30b-a3b` for the MoE checkpoint. For local
|
||||
paths without `lingbot` in the directory name, add
|
||||
`--model-class-name LingBotVideoPipeline`. Model-only options use
|
||||
`--extra-body`, for example `'{"batch_cfg": true, "output_type": "np"}'`.
|
||||
|
||||
### LTX-2
|
||||
|
||||
```bash
|
||||
|
||||
@@ -80,12 +80,14 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
### Diffusion Models
|
||||
|
||||
| Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization |
|
||||
|-------|-----------|:-------:|:------:|----------------|------------------|------------------|
|
||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | |
|
||||
| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | |
|
||||
| LTX-2 | `Lightricks/LTX-2`, `rootonchair/LTX-2-19b-distilled` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
|
||||
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers`, `diffusers/LTX-2.3-Distilled-Diffusers` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
|
||||
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None | ✅︎ |
|
||||
| FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None | |
|
||||
| FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | Yes | Not validated | All layers | None | ✅︎ |
|
||||
|
||||
@@ -12,20 +12,20 @@ For the internal architecture and backend extension points, see the
|
||||
|
||||
| Mode | Guide | Description | Methods |
|
||||
|------|-------|-------------|---------|
|
||||
| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8, MXFP4 W4A4 |
|
||||
| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, BitsAndBytes W4, MXFP8 W8A8, MXFP4 W4A4 |
|
||||
| Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA |
|
||||
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale |
|
||||
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale |
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | GGUF | AutoRound | msModelSlim |
|
||||
|--------|----------|-----------|----------|------------|------------|------|-----------|-------------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
|
||||
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ⭕ | ⭕ | ❌ |
|
||||
| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
|
||||
| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ |
|
||||
| Device | FP8 W8A8 | Int8 W8A8 | BitsAndBytes W4 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | AutoRound | msModelSlim |
|
||||
|--------|----------|-----------|------------------|----------|------------|------------|-----------|-------------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
|
||||
| AMD ROCm | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ✅ | ⭕ | ❌ |
|
||||
| Intel XPU | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
|
||||
| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide. FP8 on Ampere may use a weight-only path where available.
|
||||
@@ -43,10 +43,10 @@ otherwise.
|
||||
|--------|-------|------|----------------|--------|
|
||||
| FP8 W8A8 | [FP8](fp8.md) | Online W8A8 or checkpoint FP8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
|
||||
| Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
|
||||
| BitsAndBytes W4 | [BitsAndBytes](bitsandbytes.md) | Online W4 weight-only | Z-Image-Turbo; Qwen-Image/Wan2.2 not validated | Validated for Z-Image-Turbo |
|
||||
| ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints |
|
||||
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Online W8A8 or offline pre-quantized | Wan2.2-T2V-A14B, I2V-A14B, TI2V-5B | Ascend NPU only; validated for Wan2.2 |
|
||||
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | `mxfp4`: online single-scale only; `mxfp4_dualscale`: online or offline dual-scale (offline recommended) | Wan2.2-T2V-A14B, I2V-A14B | Ascend NPU only; validated for Wan2.2 A14B cascade models; TI2V-5B not supported; offline `mxfp4_dualscale` uses calibrated `mul_scale` for best accuracy |
|
||||
| GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists |
|
||||
| AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven |
|
||||
| msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path |
|
||||
|
||||
@@ -61,9 +61,9 @@ in BF16 unless the model guide explicitly adds support.
|
||||
|--------|-------|-------|----------------|--------|
|
||||
| ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
|
||||
| Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints |
|
||||
| msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
|
||||
@@ -76,10 +76,10 @@ attached to the intended stage rather than applied globally.
|
||||
|--------|-------|-------|----------------|--------|
|
||||
| FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
|
||||
| Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
|
||||
| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Stage-specific transformer or DiT module | BAGEL, GLM-Image | Not validated |
|
||||
| ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation |
|
||||
| MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
|
||||
| MXFP4 | [MXFP4](mxfp4.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
|
||||
| GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed |
|
||||
| AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed |
|
||||
| msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
|
||||
|
||||
@@ -106,7 +106,7 @@ config = build_quant_config({
|
||||
|
||||
| Component | Default quantized? | Notes |
|
||||
|-----------|--------------------|-------|
|
||||
| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, MXFP8, MXFP4, GGUF, AutoRound, and msModelSlim |
|
||||
| Diffusion transformer | Yes | Primary target for FP8, Int8, BitsAndBytes, ModelOpt, MXFP8, MXFP4, AutoRound, and msModelSlim |
|
||||
| Text encoder | No | Keep BF16 unless a method-specific guide documents support |
|
||||
| VAE | No | Keep BF16; storage-only paths are method-specific |
|
||||
| Scheduler/tokenizer | No | Loaded from the base model repository |
|
||||
@@ -139,8 +139,8 @@ from vllm_omni.quantization import build_quant_config
|
||||
|
||||
build_quant_config("fp8")
|
||||
build_quant_config({"method": "fp8", "activation_scheme": "static"})
|
||||
build_quant_config("bitsandbytes")
|
||||
build_quant_config("auto-round", bits=4, group_size=128)
|
||||
build_quant_config({"method": "gguf", "gguf_model": "/path/to/model.gguf"})
|
||||
build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
|
||||
build_quant_config(None)
|
||||
```
|
||||
|
||||
@@ -131,7 +131,7 @@ def build_sampling_params() -> list[SamplingParams]:
|
||||
|
||||
|
||||
def save_stage_output(stage_outputs, output_dir: str) -> None:
|
||||
output = stage_outputs.request_output
|
||||
output = stage_outputs
|
||||
request_id = output.request_id
|
||||
if stage_outputs.final_output_type == "text":
|
||||
text_output = output.outputs[0].text
|
||||
|
||||
@@ -59,7 +59,7 @@ def main(args):
|
||||
|
||||
omni = Omni(
|
||||
model=args.model_name,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
deploy_config=args.deploy_config,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
batch_timeout=args.batch_timeout,
|
||||
@@ -104,7 +104,7 @@ def main(args):
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
for stage_outputs in omni_outputs:
|
||||
output = stage_outputs.request_output
|
||||
output = stage_outputs
|
||||
if stage_outputs.final_output_type == "text":
|
||||
request_id = output.request_id
|
||||
text_output = output.outputs[0].text
|
||||
@@ -167,10 +167,10 @@ def parse_args():
|
||||
help="Sampling rate for audio loading (default: 16000).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to stage configs YAML file.",
|
||||
help="Path to a deploy config YAML file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
|
||||
@@ -228,7 +228,7 @@ async def main():
|
||||
if not outputs:
|
||||
raise ValueError("No output produced from omni.generate()")
|
||||
|
||||
first_out = outputs[0].request_output
|
||||
first_out = outputs[0]
|
||||
req_out: OmniRequestOutput = first_out
|
||||
|
||||
# Verify trajectory data (from custom pipeline)
|
||||
|
||||
@@ -741,7 +741,7 @@ def make_mmu_prompt(
|
||||
|
||||
def iter_mm_outputs(outputs: list[Any]):
|
||||
for omni_out in outputs:
|
||||
req_out = getattr(omni_out, "request_output", None)
|
||||
req_out = omni_out
|
||||
req_list = req_out if isinstance(req_out, list) else [req_out]
|
||||
for item in req_list:
|
||||
if item is None:
|
||||
@@ -876,7 +876,7 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
|
||||
parser.add_argument("--task", type=str, required=True, choices=TASK_CHOICES)
|
||||
parser.add_argument("--model", type=str, required=True, help="HF repo id or local model directory.")
|
||||
parser.add_argument(
|
||||
"--stage-config-path",
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
|
||||
help="Path to deploy config yaml.",
|
||||
@@ -1394,8 +1394,8 @@ def main() -> None:
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
stage_config_path = str(Path(args.stage_config_path).expanduser())
|
||||
omni = Omni(model=model_source, stage_configs_path=stage_config_path, dtype=args.dtype)
|
||||
deploy_config_path = str(Path(args.deploy_config).expanduser())
|
||||
omni = Omni(model=model_source, deploy_config=deploy_config_path, dtype=args.dtype)
|
||||
sampling_params_list = [
|
||||
SamplingParams(max_tokens=int(args.max_tokens_per_stage), temperature=0.0, top_p=1.0, detokenize=False)
|
||||
for _ in range(omni.num_stages)
|
||||
|
||||
@@ -73,7 +73,6 @@ def parse_args():
|
||||
)
|
||||
parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
|
||||
parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
|
||||
parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
|
||||
parser.add_argument(
|
||||
"--stream",
|
||||
action="store_true",
|
||||
@@ -146,13 +145,7 @@ def main():
|
||||
else:
|
||||
bot_task = args.bot_task
|
||||
|
||||
if args.deploy_config is not None and args.stage_configs_path is not None:
|
||||
raise ValueError("--deploy-config and --stage-configs-path are mutually exclusive.")
|
||||
|
||||
deploy_config = args.deploy_config
|
||||
stage_configs_path = args.stage_configs_path
|
||||
if deploy_config is None and stage_configs_path is None:
|
||||
deploy_config = _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
|
||||
deploy_config = args.deploy_config or _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
|
||||
|
||||
omni_kwargs = {
|
||||
"model": args.model,
|
||||
@@ -168,10 +161,7 @@ def main():
|
||||
|
||||
if additional_config is not None:
|
||||
omni_kwargs["additional_config"] = additional_config
|
||||
if deploy_config is not None:
|
||||
omni_kwargs["deploy_config"] = deploy_config
|
||||
else:
|
||||
omni_kwargs["stage_configs_path"] = stage_configs_path
|
||||
omni_kwargs["deploy_config"] = deploy_config
|
||||
|
||||
omni = Omni(**omni_kwargs)
|
||||
|
||||
@@ -272,10 +262,7 @@ def main():
|
||||
print(f" Modality: {args.modality}")
|
||||
print(f" Prompt task: {task}")
|
||||
print(f" Bot task: {bot_task}")
|
||||
if deploy_config is not None:
|
||||
print(f" Deploy config: {deploy_config}")
|
||||
else:
|
||||
print(f" Stage config: {stage_configs_path}")
|
||||
print(f" Deploy config: {deploy_config}")
|
||||
print(f" Num stages: {omni.num_stages}")
|
||||
if args.modality in ("text2img", "img2img"):
|
||||
print(f" Inference steps: {args.steps}")
|
||||
@@ -303,7 +290,7 @@ def main():
|
||||
)
|
||||
img_idx = 0
|
||||
for req_output in omni_outputs:
|
||||
ro = getattr(req_output, "request_output", None)
|
||||
ro = req_output
|
||||
stage_id = getattr(req_output, "stage_id", None)
|
||||
|
||||
# AR stage text — each CompletionOutput.text is already a delta when
|
||||
|
||||
@@ -621,7 +621,7 @@ def main():
|
||||
images = getattr(output, "images", None)
|
||||
if images:
|
||||
break
|
||||
req_out = getattr(output, "request_output", None)
|
||||
req_out = output
|
||||
images = getattr(req_out, "images", None) if req_out is not None else None
|
||||
if images:
|
||||
break
|
||||
|
||||
@@ -646,8 +646,8 @@ def main():
|
||||
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)
|
||||
if frames.is_pipeline_output and frames.request_output is not None:
|
||||
inner_output = frames.request_output
|
||||
if frames.is_pipeline_output and frames is not None:
|
||||
inner_output = frames
|
||||
if isinstance(inner_output, OmniRequestOutput):
|
||||
if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
|
||||
audio = inner_output.multimodal_output["audio"]
|
||||
|
||||
@@ -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()
|
||||
@@ -297,7 +297,7 @@ def main(args):
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
for stage_outputs in omni_outputs:
|
||||
output = stage_outputs.request_output
|
||||
output = stage_outputs
|
||||
if stage_outputs.final_output_type == "text":
|
||||
request_id = output.request_id
|
||||
text_output = output.outputs[0].text
|
||||
|
||||
@@ -347,7 +347,7 @@ def main(args):
|
||||
omni.start_profile(stages=args.profiler_stages)
|
||||
|
||||
for stage_outputs in omni.generate(prompts, sampling_params_list):
|
||||
output = stage_outputs.request_output
|
||||
output = stage_outputs
|
||||
if stage_outputs.final_output_type == "text":
|
||||
request_id = output.request_id
|
||||
text_output = output.outputs[0].text
|
||||
|
||||
@@ -452,12 +452,6 @@ def parse_args():
|
||||
default=None,
|
||||
help="Path to a .txt file with one prompt per line (preferred).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a stage configs file (deprecated; prefer --deploy-config).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
|
||||
@@ -0,0 +1,229 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline speech-to-speech with NVIDIA-NemotronLabs-VoiceChat-11B.
|
||||
|
||||
Feeds a 16 kHz user utterance (plus a spoken-style system prompt) through the
|
||||
3-stage nemotron_voicechat pipeline and writes the agent's reply as text and a
|
||||
22.05 kHz WAV.
|
||||
|
||||
Usage:
|
||||
python end2end.py \
|
||||
--checkpoint /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
|
||||
--wav /path/to/user_question_16k.wav \
|
||||
--output-dir results/nemotron_voicechat
|
||||
|
||||
The system prompt defaults to the NeMo reference default. The text tokenizer
|
||||
resolves from the ``nvidia/NVIDIA-Nemotron-Nano-9B-v2`` HF id (or the
|
||||
``NEMOTRON_VOICECHAT_LLM_PATH`` env var for air-gapped runs).
|
||||
|
||||
Sizing the input WAV: the timeline is frame-locked, so the reply budget IS the
|
||||
input duration — the acoustic channel trails the text channel, and a reply
|
||||
that has not finished when the frames run out is truncated SILENTLY (no error;
|
||||
ASR of the WAV will just disagree with the emitted text). Leave generous
|
||||
trailing silence after the question (a question ending at ~4.5 s truncated in
|
||||
an 8 s WAV but completed cleanly in a 16 s one). The script warns when the
|
||||
text channel is still emitting non-PAD tokens near the last frame.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
# Warn if the text channel is still speaking within this many final frames
|
||||
# (12 frames = ~1 s) — the frame-locked reply likely got cut off.
|
||||
TRUNCATION_GUARD_FRAMES = 12
|
||||
|
||||
# The NeMo reference offline script's default system prompt.
|
||||
DEFAULT_SYSTEM_PROMPT = (
|
||||
"You are an AI voice assistant developed by NVIDIA. "
|
||||
"Your name is NVIDIA Voice Chat. "
|
||||
"Answer in a spoken, conversational style rather than a written one. "
|
||||
"Do not repeat the same sentence over and over again. "
|
||||
"Start the conversation by greeting the user."
|
||||
)
|
||||
|
||||
|
||||
def load_wav_16k_mono(path: str) -> np.ndarray:
|
||||
wav, sr = sf.read(path, dtype="float32")
|
||||
if wav.ndim > 1:
|
||||
wav = wav.mean(axis=1)
|
||||
if sr != 16000:
|
||||
from vllm.multimodal.audio import resample_audio_scipy
|
||||
|
||||
wav = resample_audio_scipy(wav, orig_sr=sr, target_sr=16000)
|
||||
return wav.astype(np.float32)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--checkpoint", required=True, help="NVIDIA-NemotronLabs-VoiceChat-11B directory")
|
||||
parser.add_argument("--wav", required=True, help="user speech (any sr; resampled to 16 kHz mono)")
|
||||
parser.add_argument("--output-dir", default="results/nemotron_voicechat")
|
||||
parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
default=None,
|
||||
help=(
|
||||
"Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml "
|
||||
"for the async-chunk streaming pipeline (stages overlap; audio starts decoding "
|
||||
"before the thinker finishes). Default: the bundled non-streaming yaml (bit-exact "
|
||||
"parity path)."
|
||||
),
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if not Path(args.wav).is_file():
|
||||
raise FileNotFoundError(f"Input WAV not found: {args.wav}")
|
||||
ckpt = Path(args.checkpoint)
|
||||
raw_cfg = json.loads((ckpt / "config.json").read_text())
|
||||
stt_cfg = raw_cfg["model"]["stt"]["model"]
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
|
||||
compute_acoustic_frame_count,
|
||||
)
|
||||
|
||||
tok_ref = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or stt_cfg.get(
|
||||
"pretrained_llm", "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained(tok_ref, trust_remote_code=False)
|
||||
bos_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("bos_token", "<s>"))
|
||||
eos_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("eos_token", "</s>"))
|
||||
pad_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("pad_token", "<SPECIAL_12>"))
|
||||
# NeMo encode_system_prompt: [bos] + tokens + [eos].
|
||||
prompt_ids = [bos_id] + tokenizer.encode(args.system_prompt, add_special_tokens=False) + [eos_id]
|
||||
|
||||
wav = load_wav_16k_mono(args.wav)
|
||||
n_frames = compute_acoustic_frame_count(stt_cfg, int(wav.shape[0]))
|
||||
print(f"input: {wav.shape[0]} samples @16kHz -> {n_frames} acoustic frames; prompt {len(prompt_ids)} tokens")
|
||||
|
||||
inputs = {
|
||||
# vLLM prompt = system prompt + one placeholder position (acoustic
|
||||
# frame 0); the thinker replaces every position with fused embeddings.
|
||||
"prompt_token_ids": prompt_ids + [pad_id],
|
||||
"additional_information": {
|
||||
# Plain float list: raw ndarrays do not survive the engine-core
|
||||
# message serialization.
|
||||
"nvc_audio": wav.astype(np.float32).tolist(),
|
||||
"nvc_sr": 16000,
|
||||
"nvc_prompt_token_ids": prompt_ids,
|
||||
"nvc_expected_frames": n_frames,
|
||||
},
|
||||
}
|
||||
sampling_params_list = [
|
||||
# thinker: greedy, frame-locked generation length. NO min_tokens: the
|
||||
# tokenizer's eos_token is <SPECIAL_12> — the frame-locked PAD/silence
|
||||
# token — and vLLM's min_tokens masks "EOS" to -inf, which would forbid
|
||||
# silence and force the agent to babble through the whole utterance.
|
||||
SamplingParams(temperature=0.0, max_tokens=n_frames, ignore_eos=True, seed=0),
|
||||
# talker: placeholder loop, stops on the stage's stop token. 16383:
|
||||
# the 1-token placeholder prompt + max_tokens must fit the stage's
|
||||
# max_model_len (16384).
|
||||
SamplingParams(temperature=0.0, max_tokens=16383, stop_token_ids=[1], detokenize=False, seed=0),
|
||||
# code2wav: single decode step.
|
||||
SamplingParams(temperature=0.0, max_tokens=1, detokenize=False, seed=0),
|
||||
]
|
||||
|
||||
omni_kwargs: dict = {}
|
||||
if args.stage_configs_path:
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path
|
||||
omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs)
|
||||
outputs = omni.generate(inputs, sampling_params_list)
|
||||
|
||||
out_dir = Path(args.output_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
stem = Path(args.wav).stem
|
||||
|
||||
agent_text = None
|
||||
audio = None
|
||||
sr = 22050
|
||||
for stage_outputs in outputs:
|
||||
# One entry per final stage (text from the thinker, audio from code2wav);
|
||||
# mirrors the qwen2_5_omni offline example's collection loop.
|
||||
output = getattr(stage_outputs, "request_output", stage_outputs)
|
||||
final_type = getattr(stage_outputs, "final_output_type", None)
|
||||
completions = getattr(output, "outputs", None) or []
|
||||
if not completions:
|
||||
continue
|
||||
completion = completions[0]
|
||||
if final_type in (None, "text"):
|
||||
text = getattr(completion, "text", None)
|
||||
if text:
|
||||
agent_text = text
|
||||
token_ids = getattr(completion, "token_ids", None)
|
||||
if token_ids:
|
||||
(out_dir / f"{stem}_text_tokens.json").write_text(json.dumps(list(map(int, token_ids))))
|
||||
# Frame lock: the reply budget IS the input length. If the text
|
||||
# channel is still mid-sentence when the frames run out, the
|
||||
# spoken answer is silently truncated — warn instead.
|
||||
tail = list(map(int, token_ids))[-TRUNCATION_GUARD_FRAMES:]
|
||||
if any(t != pad_id for t in tail):
|
||||
print(
|
||||
f"WARNING: the text channel emits non-PAD tokens within the last "
|
||||
f"{TRUNCATION_GUARD_FRAMES} frames ({TRUNCATION_GUARD_FRAMES * 80} ms) — the reply "
|
||||
"likely ran out of acoustic frames and is cut off mid-sentence. "
|
||||
"Pad the input WAV with more trailing silence and rerun."
|
||||
)
|
||||
# Best-effort: dump the function-channel debug timeline if the
|
||||
# engine surfaced the request info on this output.
|
||||
info = getattr(output, "additional_information", None)
|
||||
func_timeline = info.get("nvc_function_tokens") if isinstance(info, dict) else None
|
||||
if func_timeline is not None:
|
||||
values = func_timeline.tolist() if hasattr(func_timeline, "tolist") else list(func_timeline)
|
||||
(out_dir / f"{stem}_function_tokens.json").write_text(json.dumps([int(v) for v in values]))
|
||||
print(f"function-channel timeline: {len(values)} entries")
|
||||
if final_type in (None, "audio"):
|
||||
# multimodal_output is a MultimodalPayload (Mapping) — duck-type it.
|
||||
mm = getattr(completion, "multimodal_output", None)
|
||||
get = getattr(mm, "get", None)
|
||||
if not callable(get):
|
||||
continue
|
||||
candidate = None
|
||||
for key in ("audio", "model_outputs", "wav", "waveform"):
|
||||
value = get(key)
|
||||
if value is not None:
|
||||
candidate = value
|
||||
break
|
||||
if isinstance(candidate, list):
|
||||
candidate = candidate[0] if candidate else None
|
||||
if candidate is not None:
|
||||
arr = np.asarray(
|
||||
candidate.float().cpu() if hasattr(candidate, "cpu") else candidate, dtype=np.float32
|
||||
).reshape(-1)
|
||||
if arr.size > 0:
|
||||
audio = arr
|
||||
sr_val = get("sr")
|
||||
if sr_val is not None:
|
||||
sr_item = sr_val[0] if isinstance(sr_val, list) else sr_val
|
||||
sr = int(sr_item.item() if hasattr(sr_item, "item") else sr_item)
|
||||
|
||||
if agent_text is not None:
|
||||
# The frame-locked text channel emits PAD tokens on silent frames;
|
||||
# strip them (NeMo's post-inference detokenizer does the same).
|
||||
pad_token = stt_cfg.get("pad_token", "<SPECIAL_12>")
|
||||
agent_text = " ".join(agent_text.replace(pad_token, " ").split())
|
||||
text_path = out_dir / f"{stem}_output.txt"
|
||||
text_path.write_text(agent_text + "\n")
|
||||
print(f"agent text -> {text_path}\n{agent_text}")
|
||||
else:
|
||||
print("WARNING: no agent text in outputs")
|
||||
|
||||
if audio is not None and audio.size > 0:
|
||||
wav_path = out_dir / f"{stem}_output.wav"
|
||||
sf.write(wav_path, audio, sr)
|
||||
print(f"agent audio -> {wav_path} ({audio.size / sr:.2f}s @ {sr} Hz)")
|
||||
else:
|
||||
raise RuntimeError("No audio produced by the code2wav stage.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -392,7 +392,7 @@ def main(args):
|
||||
total_requests = len(prompts)
|
||||
processed_count = 0
|
||||
for stage_outputs in omni_generator:
|
||||
output = stage_outputs.request_output
|
||||
output = stage_outputs
|
||||
if stage_outputs.final_output_type == "text":
|
||||
request_id = output.request_id
|
||||
text_output = output.outputs[0].text
|
||||
|
||||
@@ -367,7 +367,7 @@ def main(args):
|
||||
print(f"query type: {args.query_type}")
|
||||
|
||||
for stage_outputs in omni_generator:
|
||||
output = stage_outputs.request_output
|
||||
output = stage_outputs
|
||||
if stage_outputs.final_output_type == "text":
|
||||
request_id = output.request_id
|
||||
text_output = output.outputs[0].text
|
||||
@@ -484,10 +484,10 @@ def parse_args():
|
||||
help="Path to a .txt file with one prompt per line (preferred).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a stage configs file.",
|
||||
help="Path to a deploy config YAML.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video-path",
|
||||
|
||||
@@ -237,7 +237,7 @@ async def run_single_request(
|
||||
sampling_params_list=sampling_params_list,
|
||||
output_modalities=output_modalities,
|
||||
):
|
||||
output = omni_output.request_output
|
||||
output = omni_output
|
||||
if omni_output.final_output_type == "text":
|
||||
if stage_0_first_output_ts is None:
|
||||
stage_0_first_output_ts = time.perf_counter()
|
||||
|
||||
@@ -26,7 +26,6 @@ import torch
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
@@ -352,8 +351,8 @@ def main():
|
||||
else:
|
||||
print("[Profiler] No valid profiling data returned.")
|
||||
|
||||
# Extract output from result
|
||||
output = OmniRequestOutput.unwrap_result(result)
|
||||
# omni.generate() returns a list for sync calls; unwrap single-result list.
|
||||
output = result[0] if isinstance(result, list) else result
|
||||
|
||||
if not output.images:
|
||||
raise ValueError("No video frames found in OmniRequestOutput.")
|
||||
|
||||
@@ -232,15 +232,8 @@ def main(args):
|
||||
raise ValueError(f"Unknown query type: {args.query_type}")
|
||||
|
||||
# Initialize vLLM-Omni with Step-Audio2
|
||||
# Resolve deploy config path. Keep explicit stage_configs_path for legacy
|
||||
# custom configs, but use bundled deploy configs by default.
|
||||
deploy_config_path = None
|
||||
stage_config_path = None
|
||||
if args.stage_configs_path and args.deploy_config:
|
||||
raise ValueError("--stage-configs-path and --deploy-config are mutually exclusive")
|
||||
if args.stage_configs_path:
|
||||
stage_config_path = args.stage_configs_path
|
||||
elif args.deploy_config:
|
||||
# Resolve an explicit deploy config or use the bundled task default.
|
||||
if args.deploy_config:
|
||||
deploy_config_path = args.deploy_config
|
||||
else:
|
||||
configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/deploy"
|
||||
@@ -262,10 +255,7 @@ def main(args):
|
||||
"ray_address": args.ray_address,
|
||||
"trust_remote_code": True,
|
||||
}
|
||||
if stage_config_path is not None:
|
||||
omni_kwargs["stage_configs_path"] = stage_config_path
|
||||
else:
|
||||
omni_kwargs["deploy_config"] = deploy_config_path
|
||||
omni_kwargs["deploy_config"] = deploy_config_path
|
||||
|
||||
omni_llm = Omni(**omni_kwargs)
|
||||
|
||||
@@ -345,7 +335,7 @@ def main(args):
|
||||
# Process outputs from each stage
|
||||
for stage_idx, stage_outputs in enumerate(omni_outputs):
|
||||
# request_output may be a single object or a list
|
||||
raw_output = stage_outputs.request_output
|
||||
raw_output = stage_outputs
|
||||
outputs_list = raw_output if isinstance(raw_output, list) else [raw_output]
|
||||
|
||||
if stage_outputs.final_output_type == "text":
|
||||
@@ -418,12 +408,6 @@ def parse_args():
|
||||
" - HuggingFace ID: stepfun-ai/Step-Audio-2-mini"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a legacy stage config YAML file (default: use bundled deploy config)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
|
||||
@@ -43,6 +43,7 @@ from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
AUDIOX_TASKS = ("t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m")
|
||||
@@ -388,9 +389,9 @@ def main():
|
||||
raise ValueError("No output generated from omni.generate()")
|
||||
|
||||
output = outputs[0]
|
||||
if not hasattr(output, "request_output") or not output.request_output:
|
||||
if not isinstance(output, OmniRequestOutput) or not output:
|
||||
raise ValueError("No request_output found in OmniRequestOutput")
|
||||
request_output = output.request_output
|
||||
request_output = output
|
||||
if not hasattr(request_output, "multimodal_output"):
|
||||
raise ValueError("No multimodal_output found in request_output")
|
||||
|
||||
|
||||
@@ -115,7 +115,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
|
||||
)
|
||||
images_outputs = []
|
||||
for output in outputs:
|
||||
req_out = output.request_output
|
||||
req_out = output
|
||||
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
|
||||
raise ValueError("Invalid request_output structure or missing 'images' key")
|
||||
images = req_out.images
|
||||
|
||||
@@ -92,12 +92,6 @@ def parse_args() -> argparse.Namespace:
|
||||
"meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
|
||||
"stabilityai/stable-diffusion-3.5-medium, Tongyi-MAI/Z-Image-Turbo and etc.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="[Deprecated] Path to a legacy stage_args-format YAML. Prefer --deploy-config.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
@@ -475,8 +469,6 @@ def main():
|
||||
omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
|
||||
if args.enforce_eager is not None:
|
||||
omni_kwargs["enforce_eager"] = args.enforce_eager
|
||||
if args.stage_configs_path:
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path
|
||||
if args.deploy_config:
|
||||
omni_kwargs["deploy_config"] = args.deploy_config
|
||||
if use_nextstep:
|
||||
@@ -515,8 +507,8 @@ def main():
|
||||
print(f" Image size: {args.width}x{args.height}")
|
||||
if args.lora_path:
|
||||
print(f" LoRA: scale={args.lora_scale}")
|
||||
if args.stage_configs_path:
|
||||
print(f" stage-configs-path: {args.stage_configs_path}")
|
||||
if args.deploy_config:
|
||||
print(f" deploy-config: {args.deploy_config}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
# Build LoRA request when --lora-path is set
|
||||
@@ -656,7 +648,7 @@ def main():
|
||||
images = getattr(output, "images", None)
|
||||
if images:
|
||||
break
|
||||
req_out = getattr(output, "request_output", None)
|
||||
req_out = output
|
||||
images = getattr(req_out, "images", None) if req_out is not None else None
|
||||
if images:
|
||||
break
|
||||
|
||||
@@ -23,6 +23,7 @@ list of supported architectures across all modalities, see
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
|
||||
| dots.tts | `rednote-hilab/dots.tts-soar` | single (native AR) | — (not wired yet) | — | — | 48 kHz |
|
||||
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz |
|
||||
| IndexTTS-2.5 | native `checkpoints/` bundle | 2 (AR talker + EnhancedCodec + S2Mel DiT + BigVGAN) | ✓ (required) | — | multilingual (`--lang`) + emotion control | 22.05 kHz |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
|
||||
@@ -372,18 +373,10 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
Streaming requires `async_chunk: true` in the stage config.
|
||||
|
||||
### Word Timestamps
|
||||
Generate a WAV offline and a JSON sidecar with word-level timestamps from
|
||||
`Qwen/Qwen3-ForcedAligner-0.6B`:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--forced-aligner Qwen/Qwen3-ForcedAligner-0.6B \
|
||||
--text "Hello world." \
|
||||
--output-dir /tmp/qwen3_tts_timestamps
|
||||
```
|
||||
The script writes `qwen3_tts_word_timestamps.wav` and
|
||||
`qwen3_tts_word_timestamps.json`. On machines without a local CUDA toolkit,
|
||||
set `VLLM_USE_FLASHINFER_SAMPLER=0` to avoid FlashInfer sampler JIT.
|
||||
Word-level timestamps are currently a serving-path feature: launch
|
||||
`vllm-omni serve` with `--forced-aligner` and request `word_timestamps`
|
||||
(see `examples/online_serving/text_to_speech/README.md`). An offline
|
||||
example is not provided in this release.
|
||||
|
||||
### Batched decoding
|
||||
The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, point `--stage-configs-path` at a stage configs YAML with the desired values (see `vllm_omni/model_executor/stage_configs/` for templates):
|
||||
@@ -438,6 +431,31 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [
|
||||
|
||||
---
|
||||
|
||||
## dots.tts
|
||||
|
||||
Single-stage native AR TTS at 48 kHz (rednote-hilab). Pipeline: `Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR loopback → AudioVAE (streaming decode)`. Same "vLLM-native base LM + side-path computation" pattern as VoxCPM2, with a plain Qwen2 backbone instead of MiniCPM4 and no FSQ / residual LM stage.
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
|
||||
--model rednote-hilab/dots.tts-soar \
|
||||
--text "Hello, this is a test of dots TTS running on vLLM Omni."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
Not wired in this release — generation is zero-shot only. The CAM++ x-vector speaker encoder weights load, but `end2end.py` has no `--ref-audio`/`--ref-text` flags yet.
|
||||
|
||||
### Streaming
|
||||
The AudioVAE decoder has an internal streaming path (`init_stream_state` / `stream_step` / `stream_flush`) used to avoid boundary artifacts between 160 ms patches, but it is not yet exposed through an online serving endpoint or example.
|
||||
|
||||
### Notes
|
||||
- Output: 48 kHz mono WAV.
|
||||
- Deploy config: `vllm_omni/deploy/dots_tts.yaml` (auto-loaded by HF `model_type`).
|
||||
- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default. `dots.tts-base` shares the same architecture but is unvalidated in this repo. `dots.tts-mf` (MeanFlow, 2-4 step) is not supported yet.
|
||||
- Known limitation: no CUDA graph capture and no batched side-path computation yet, so concurrent requests do not currently scale (each request's DiT Euler steps run serially). See `recipes/rednote-hilab/dots.tts.md` for details and the roadmap.
|
||||
|
||||
---
|
||||
|
||||
## IndexTTS-2 and IndexTTS-2.5
|
||||
|
||||
Both versions use a 2-stage TTS pipeline and produce 22.05 kHz mono audio. Stage 0 is a GPT AR talker. IndexTTS-2 Stage 1 uses RepCodec semantic embeddings, GPT latent, S2Mel CFM/DiT, and BigVGAN. IndexTTS-2.5 replaces RepCodec with EnhancedCodec and uses the official code-only path (`use_gpt_latent=false`). Every request requires reference audio for zero-shot voice cloning.
|
||||
|
||||
@@ -168,7 +168,7 @@ def run_e2e():
|
||||
print(f"Received {len(outputs)} outputs.")
|
||||
for i, output in enumerate(outputs):
|
||||
try:
|
||||
ro = output.request_output
|
||||
ro = output
|
||||
if ro is None:
|
||||
print("No request_output found.")
|
||||
continue
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
"""Offline dots.tts inference example (single-stage native AR pipeline).
|
||||
|
||||
Uses the single-stage native AR config (vllm_omni/deploy/dots_tts.yaml).
|
||||
Reference checkpoint: rednote-hilab/dots.tts-soar.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
SAMPLE_RATE = 48_000
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="Offline dots.tts native AR inference")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="rednote-hilab/dots.tts-soar",
|
||||
help="dots.tts model path or HuggingFace repo ID.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
type=str,
|
||||
default="Hello, this is a test of dots TTS running on vLLM Omni.",
|
||||
help="Text to synthesize.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default="output_audio",
|
||||
help="Directory for output WAV files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override the deploy config path. If unset, auto-loads "
|
||||
"vllm_omni/deploy/dots_tts.yaml based on the HF model_type.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def extract_audio(multimodal_output: dict) -> torch.Tensor:
|
||||
"""Extract the final complete audio tensor from multimodal output.
|
||||
|
||||
The output processor concatenates per-step delta tensors under
|
||||
``model_outputs``. Falls back to ``audio`` for backwards compat.
|
||||
"""
|
||||
audio = multimodal_output.get("model_outputs")
|
||||
if audio is None:
|
||||
audio = multimodal_output.get("audio")
|
||||
if audio is None:
|
||||
raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}")
|
||||
|
||||
if isinstance(audio, list):
|
||||
valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
|
||||
if not valid:
|
||||
raise ValueError("Audio list is empty or all elements are None.")
|
||||
return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
|
||||
|
||||
return torch.as_tensor(audio).float().cpu().reshape(-1)
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
engine = Omni(
|
||||
model=args.model,
|
||||
deploy_config=args.deploy_config,
|
||||
)
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import (
|
||||
build_dots_tts_prompt,
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
|
||||
|
||||
prompt = build_dots_tts_prompt(
|
||||
tokenizer=tokenizer,
|
||||
text=args.text,
|
||||
)
|
||||
|
||||
print(f"Model : {args.model}")
|
||||
print(f"Text : {args.text}")
|
||||
print(f"Output dir : {output_dir}")
|
||||
print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens")
|
||||
|
||||
t_start = time.perf_counter()
|
||||
outputs = engine.generate([prompt])
|
||||
elapsed = time.perf_counter() - t_start
|
||||
|
||||
request_output = outputs[0]
|
||||
mm = request_output.outputs[0].multimodal_output
|
||||
audio = extract_audio(mm)
|
||||
|
||||
duration = audio.numel() / SAMPLE_RATE
|
||||
rtf = elapsed / duration if duration > 0 else float("inf")
|
||||
|
||||
output_path = output_dir / "output.wav"
|
||||
sf.write(str(output_path), audio.numpy(), SAMPLE_RATE, format="WAV")
|
||||
|
||||
print(f"Saved : {output_path}")
|
||||
print(f"Duration : {duration:.2f}s")
|
||||
print(f"Inference : {elapsed:.2f}s")
|
||||
print(f"RTF : {rtf:.3f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
main()
|
||||
@@ -38,7 +38,7 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_MODEL = "fishaudio/s2-pro"
|
||||
DEFAULT_STAGE_CONFIG = os.path.join(
|
||||
DEFAULT_DEPLOY_CONFIG = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
"..",
|
||||
"..",
|
||||
@@ -152,10 +152,7 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
if args.deploy_config:
|
||||
omni_kwargs["deploy_config"] = args.deploy_config
|
||||
else:
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path or DEFAULT_STAGE_CONFIG
|
||||
omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
|
||||
return omni_kwargs
|
||||
|
||||
|
||||
@@ -179,7 +176,7 @@ def main(args):
|
||||
|
||||
t_start = time.perf_counter()
|
||||
for stage_outputs in omni.generate(inputs):
|
||||
request_output = stage_outputs.request_output
|
||||
request_output = stage_outputs
|
||||
if request_output is None or not request_output.outputs:
|
||||
continue
|
||||
_save_wav(
|
||||
@@ -215,7 +212,7 @@ async def main_streaming(args):
|
||||
sample_rate = None
|
||||
|
||||
async for stage_output in omni.generate(prompt, request_id=request_id):
|
||||
mm = stage_output.request_output.outputs[0].multimodal_output
|
||||
mm = stage_output.outputs[0].multimodal_output
|
||||
if not stage_output.finished:
|
||||
t_now = time.perf_counter()
|
||||
sample_rate, n = _collect_audio_chunk(mm, all_audio_chunks, sample_rate)
|
||||
@@ -278,12 +275,6 @@ def parse_args():
|
||||
default=None,
|
||||
help="Deploy config YAML path (under vllm_omni/deploy/).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to stage configs YAML (legacy, prefer --deploy-config).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default="output_audio",
|
||||
|
||||
@@ -145,7 +145,7 @@ def main(args):
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
stage_configs_path=deploy_config_path,
|
||||
deploy_config=deploy_config_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
@@ -137,7 +137,7 @@ def main(args) -> None:
|
||||
params_list = sampling_params
|
||||
|
||||
for stage_outputs in omni.generate(inputs, params_list):
|
||||
for i, req_output in enumerate(stage_outputs.request_output):
|
||||
for i, req_output in enumerate(stage_outputs):
|
||||
for j, out in enumerate(req_output.outputs):
|
||||
mm = out.multimodal_output
|
||||
if mm is None:
|
||||
|
||||
@@ -135,7 +135,7 @@ def run_e2e():
|
||||
print(f"Received {len(outputs)} outputs.")
|
||||
for i, output in enumerate(outputs):
|
||||
try:
|
||||
ro = output.request_output
|
||||
ro = output
|
||||
if ro is None:
|
||||
print("No request_output found.")
|
||||
continue
|
||||
|
||||
@@ -391,7 +391,7 @@ def main(args):
|
||||
for batch_start in range(0, len(inputs), batch_size):
|
||||
batch = inputs[batch_start : batch_start + batch_size]
|
||||
for stage_outputs in omni.generate(batch):
|
||||
output = stage_outputs.request_output
|
||||
output = stage_outputs
|
||||
_save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
|
||||
|
||||
|
||||
@@ -412,7 +412,7 @@ async def main_streaming(args):
|
||||
t_prev = t_start
|
||||
chunk_idx = 0
|
||||
async for stage_output in omni.generate(prompt, request_id=request_id):
|
||||
mm = stage_output.request_output.outputs[0].multimodal_output
|
||||
mm = stage_output.outputs[0].multimodal_output
|
||||
if not stage_output.finished:
|
||||
t_now = time.perf_counter()
|
||||
audio = mm.get("audio")
|
||||
@@ -490,10 +490,10 @@ def parse_args():
|
||||
help="Path to a .txt file with one prompt per line (preferred).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a stage configs file.",
|
||||
help="Path to a deploy config YAML.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-path",
|
||||
|
||||
@@ -116,7 +116,7 @@ def main(args: Any) -> None:
|
||||
prompt = _build_custom_voice_input(args)
|
||||
final_output = None
|
||||
for stage_outputs in omni.generate([prompt]):
|
||||
final_output = stage_outputs.request_output
|
||||
final_output = stage_outputs
|
||||
if final_output is None:
|
||||
raise RuntimeError("Qwen3-TTS did not produce an output.")
|
||||
|
||||
|
||||
@@ -38,7 +38,7 @@ def _stage_durations_from_output(outputs) -> dict[str, float]:
|
||||
durations = getattr(omni_out, "stage_durations", None) or {}
|
||||
if durations:
|
||||
return dict(durations)
|
||||
ro = getattr(omni_out, "request_output", None)
|
||||
ro = omni_out
|
||||
if ro is not None:
|
||||
return dict(getattr(ro, "stage_durations", None) or {})
|
||||
return {}
|
||||
|
||||
@@ -99,7 +99,7 @@ def extract_audio(outputs) -> tuple[np.ndarray, int]:
|
||||
out = outputs[0]
|
||||
mm = getattr(out, "multimodal_output", None)
|
||||
if mm is None:
|
||||
ro = getattr(out, "request_output", None)
|
||||
ro = out
|
||||
if ro is not None:
|
||||
mm = getattr(ro, "multimodal_output", None)
|
||||
if mm is None and ro.outputs:
|
||||
|
||||
@@ -8,6 +8,7 @@ A unified script for text-to-video generation. Supports multiple models with mod
|
||||
|---|---|---|---|---|---|
|
||||
| `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) |
|
||||
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
|
||||
| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) |
|
||||
| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
|
||||
@@ -57,18 +58,31 @@ python text_to_video.py \
|
||||
|
||||
### LingBot-Video
|
||||
|
||||
The dedicated LingBot example explicitly selects the video output modality:
|
||||
The shared runner recognizes the official dense and MoE checkpoint IDs, selects
|
||||
`LingBotVideoPipeline`, and builds the canonical video request envelope:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
|
||||
python text_to_video.py \
|
||||
--model robbyant/lingbot-video-dense-1.3b \
|
||||
--prompt "a robotic arm picks up a red block" \
|
||||
--height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
|
||||
--guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
|
||||
--output lingbot_t2v.mp4
|
||||
```
|
||||
|
||||
Use the same command with `robbyant/lingbot-video-moe-30b-a3b` for the MoE
|
||||
checkpoint. Requested frame counts are rounded upward to the causal
|
||||
VAE's `4n+1` grid.
|
||||
VAE's `4n+1` grid. For a local checkpoint path whose name does not contain
|
||||
`lingbot`, pass `--model-class-name LingBotVideoPipeline`.
|
||||
|
||||
LingBot-only options such as `batch_cfg` and `output_type` use the shared
|
||||
model-extra channel, for example:
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model robbyant/lingbot-video-dense-1.3b \
|
||||
--extra-body '{"batch_cfg": true, "output_type": "np"}'
|
||||
```
|
||||
|
||||
### LTX-2
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@ from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.lora.request import LoRARequest
|
||||
from vllm_omni.lora.utils import stable_lora_int_id
|
||||
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
|
||||
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name, get_output_tensor_range
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
@@ -80,6 +80,17 @@ _MODEL_PRESETS = {
|
||||
"fps": 16,
|
||||
"output": "helios_output.mp4",
|
||||
},
|
||||
"lingbot": {
|
||||
"model_class_name": "LingBotVideoPipeline",
|
||||
"height": 192,
|
||||
"width": 320,
|
||||
"num_frames": 9,
|
||||
"num_inference_steps": 2,
|
||||
"guidance_scale": 3.0,
|
||||
"fps": 24,
|
||||
"flow_shift": 3.0,
|
||||
"output": "lingbot_video_output.mp4",
|
||||
},
|
||||
"ltx2": {
|
||||
"height": 512,
|
||||
"width": 768,
|
||||
@@ -110,27 +121,51 @@ _MODEL_PRESETS = {
|
||||
def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
|
||||
model_lower = model.lower()
|
||||
class_lower = (model_class_name or "").lower()
|
||||
if "distilled" in class_lower or "distilled" in model_lower:
|
||||
return _MODEL_PRESETS["ltx2_distilled"]
|
||||
if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
|
||||
return _MODEL_PRESETS["ltx23"]
|
||||
if "ltx2" in class_lower or "ltx-2" in model_lower or "ltx_2" in model_lower:
|
||||
if "lingbot" in model_lower or "lingbotvideo" in class_lower:
|
||||
return _MODEL_PRESETS["lingbot"]
|
||||
if "ltx" in class_lower or "ltx" in model_lower:
|
||||
if "distilled" in class_lower or "distilled" in model_lower:
|
||||
return _MODEL_PRESETS["ltx2_distilled"]
|
||||
if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
|
||||
return _MODEL_PRESETS["ltx23"]
|
||||
return _MODEL_PRESETS["ltx2"]
|
||||
if "vace" in model_lower:
|
||||
if "vace" in model_lower or "vace" in class_lower:
|
||||
return _MODEL_PRESETS["vace"]
|
||||
# Edge must be matched before the generic cosmos branch (its "cosmos" substring would
|
||||
# otherwise pick up the Nano/Super 720p / gs 6.0 / flow_shift 10.0 preset).
|
||||
if ("cosmos" in model_lower or "cosmos" in class_lower) and ("edge" in model_lower or "edge" in class_lower):
|
||||
return _MODEL_PRESETS["cosmos3_edge"]
|
||||
if "cosmos" in model_lower:
|
||||
if "cosmos" in model_lower or "cosmos" in class_lower:
|
||||
return _MODEL_PRESETS["cosmos"]
|
||||
if "hunyuan" in model_lower:
|
||||
if "hunyuan" in model_lower or "hunyuan" in class_lower:
|
||||
return _MODEL_PRESETS["hunyuan"]
|
||||
if "helios" in model_lower:
|
||||
if "helios" in model_lower or "helios" in class_lower:
|
||||
return _MODEL_PRESETS["helios"]
|
||||
return _MODEL_PRESETS["wan"]
|
||||
|
||||
|
||||
def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]:
|
||||
"""Build the canonical request envelope for the shared T2V example."""
|
||||
result: dict[str, Any] = {
|
||||
"prompt": prompt,
|
||||
"modalities": ["video"],
|
||||
}
|
||||
if negative_prompt is not None:
|
||||
result["negative_prompt"] = negative_prompt
|
||||
return result
|
||||
|
||||
|
||||
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():
|
||||
return tensor
|
||||
if source_range == "negative_one_to_one":
|
||||
return tensor.clamp(-1, 1) * 0.5 + 0.5
|
||||
if source_range == "zero_to_one":
|
||||
return tensor.clamp(0, 1)
|
||||
raise ValueError(f"Unsupported floating-point tensor range: {source_range!r}")
|
||||
|
||||
|
||||
def parse_profiler_config(value: str) -> dict[str, Any]:
|
||||
try:
|
||||
config = json.loads(value)
|
||||
@@ -160,7 +195,7 @@ def parse_extra_body(value: str) -> dict[str, Any]:
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate a video from a text prompt. "
|
||||
"Supports Wan2.2, HunyuanVideo-1.5, Helios, and other text-to-video models."
|
||||
"Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, and other text-to-video models."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
@@ -390,7 +425,7 @@ def _extract_peak_memory_mb(result: Any) -> float:
|
||||
return 0.0
|
||||
val = getattr(result, "peak_memory_mb", 0.0)
|
||||
if not val:
|
||||
inner = getattr(result, "request_output", None)
|
||||
inner = result
|
||||
if isinstance(inner, list):
|
||||
inner = inner[0] if inner else None
|
||||
val = getattr(inner, "peak_memory_mb", 0.0)
|
||||
@@ -408,6 +443,7 @@ def main():
|
||||
for key, default_val in preset.items():
|
||||
if getattr(args, key.replace("-", "_"), None) is None:
|
||||
setattr(args, key.replace("-", "_"), default_val)
|
||||
model_class_name = args.model_class_name
|
||||
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
# Cache-dit config (Wan2.2 only)
|
||||
@@ -480,6 +516,7 @@ def main():
|
||||
omni = Omni(**omni_kwargs)
|
||||
model_class_name = get_model_class_name(omni) or model_class_name
|
||||
declared_extra_body_params = get_extra_body_params(model_class_name)
|
||||
output_tensor_range = get_output_tensor_range(model_class_name)
|
||||
|
||||
if profiler_enabled:
|
||||
print("[Profiler] Starting profiling...")
|
||||
@@ -513,12 +550,11 @@ def main():
|
||||
lora_path=lora_path,
|
||||
)
|
||||
|
||||
prompt_dict = {"prompt": args.prompt}
|
||||
if args.negative_prompt is not None:
|
||||
prompt_dict["negative_prompt"] = args.negative_prompt
|
||||
elif preset not in (_MODEL_PRESETS["ltx2"], _MODEL_PRESETS["ltx23"]):
|
||||
negative_prompt = args.negative_prompt
|
||||
if negative_prompt is None and all(preset is not _MODEL_PRESETS[name] for name in ("lingbot", "ltx2", "ltx23")):
|
||||
# Preserve the historical empty-prompt behavior for non-LTX examples.
|
||||
prompt_dict["negative_prompt"] = ""
|
||||
negative_prompt = ""
|
||||
prompt_dict = build_text_to_video_prompt(args.prompt, negative_prompt)
|
||||
|
||||
extra_args = {}
|
||||
if lora_request:
|
||||
@@ -575,8 +611,8 @@ def main():
|
||||
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)
|
||||
if frames.is_pipeline_output and frames.request_output is not None:
|
||||
inner_output = frames.request_output
|
||||
if frames.is_pipeline_output and frames is not None:
|
||||
inner_output = frames
|
||||
if isinstance(inner_output, OmniRequestOutput):
|
||||
if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
|
||||
audio = inner_output.multimodal_output["audio"]
|
||||
@@ -630,8 +666,7 @@ def main():
|
||||
frame_tensor = frame_tensor[0]
|
||||
if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4):
|
||||
frame_tensor = frame_tensor.permute(1, 2, 0)
|
||||
if frame_tensor.is_floating_point():
|
||||
frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5
|
||||
frame_tensor = _normalize_float_tensor(frame_tensor, output_tensor_range)
|
||||
return frame_tensor.float().numpy()
|
||||
if isinstance(frame, np.ndarray):
|
||||
frame_array = frame
|
||||
@@ -683,8 +718,7 @@ def main():
|
||||
video_tensor = video_tensor[0]
|
||||
elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
|
||||
video_tensor = video_tensor.permute(1, 2, 3, 0)
|
||||
if video_tensor.is_floating_point():
|
||||
video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
|
||||
video_tensor = _normalize_float_tensor(video_tensor, output_tensor_range)
|
||||
video_array = video_tensor.float().numpy()
|
||||
elif isinstance(frames, np.ndarray):
|
||||
video_array = frames
|
||||
|
||||
@@ -1,131 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate text-to-video with a LingBot-Video checkpoint.")
|
||||
parser.add_argument("--model", default="robbyant/lingbot-video-dense-1.3b")
|
||||
parser.add_argument("--prompt", default="a robotic arm picks up a red block")
|
||||
parser.add_argument("--negative-prompt", default=None)
|
||||
parser.add_argument("--output", default="lingbot_video_output.mp4")
|
||||
parser.add_argument("--height", type=int, default=192)
|
||||
parser.add_argument("--width", type=int, default=320)
|
||||
parser.add_argument("--num-frames", type=int, default=9)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=2)
|
||||
parser.add_argument("--guidance-scale", type=float, default=3.0)
|
||||
parser.add_argument("--flow-shift", type=float, default=3.0)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--fps", type=int, default=24)
|
||||
parser.add_argument("--output-type", choices=["pt", "np"], default="pt")
|
||||
parser.add_argument("--batch-cfg", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _peak_memory_mb(output: Any) -> float:
|
||||
result = output[0] if isinstance(output, list) and output else output
|
||||
return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
|
||||
|
||||
|
||||
def _extract_video(output: Any) -> Any:
|
||||
if isinstance(output, list):
|
||||
output = output[0] if output else None
|
||||
if isinstance(output, OmniRequestOutput):
|
||||
if output.final_output_type != "image":
|
||||
raise ValueError(f"Unexpected output type {output.final_output_type!r}.")
|
||||
if not output.images:
|
||||
raise ValueError("No video frames found in OmniRequestOutput.")
|
||||
return output.images[0] if len(output.images) == 1 else output.images
|
||||
return output
|
||||
|
||||
|
||||
def _to_frame_array(video: Any) -> np.ndarray:
|
||||
if isinstance(video, torch.Tensor):
|
||||
tensor = video.detach().cpu()
|
||||
if tensor.dim() == 5:
|
||||
tensor = tensor[0]
|
||||
if tensor.dim() == 4 and tensor.shape[-1] not in (3, 4):
|
||||
if tensor.shape[0] in (3, 4):
|
||||
tensor = tensor.permute(1, 2, 3, 0)
|
||||
elif tensor.shape[1] in (3, 4):
|
||||
tensor = tensor.permute(0, 2, 3, 1)
|
||||
if tensor.is_floating_point():
|
||||
if float(tensor.min()) < 0.0:
|
||||
tensor = tensor.clamp(-1, 1) * 0.5 + 0.5
|
||||
else:
|
||||
tensor = tensor.clamp(0, 1)
|
||||
return tensor.float().numpy()
|
||||
if isinstance(video, np.ndarray):
|
||||
array = video[0] if video.ndim == 5 else video
|
||||
if np.issubdtype(array.dtype, np.integer):
|
||||
return array.astype(np.float32) / 255.0
|
||||
return array.astype(np.float32, copy=False)
|
||||
if isinstance(video, list):
|
||||
return np.stack([_to_frame_array(frame) for frame in video], axis=0)
|
||||
return np.asarray(video, dtype=np.float32)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
model_class_name="LingBotVideoPipeline",
|
||||
flow_shift=args.flow_shift,
|
||||
parallel_config=DiffusionParallelConfig(),
|
||||
)
|
||||
|
||||
prompt: dict[str, Any] = {
|
||||
"prompt": args.prompt,
|
||||
"modalities": ["video"],
|
||||
"num_frames": args.num_frames,
|
||||
}
|
||||
if args.negative_prompt is not None:
|
||||
prompt["negative_prompt"] = args.negative_prompt
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
generator=generator,
|
||||
output_type=args.output_type,
|
||||
extra_args={"batch_cfg": args.batch_cfg, "shift": args.flow_shift},
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
output = omni.generate(prompt, sampling_params)
|
||||
elapsed = time.perf_counter() - start
|
||||
print(f"Total generation time: {elapsed:.4f} seconds ({elapsed * 1000:.2f} ms)")
|
||||
peak_mb = _peak_memory_mb(output)
|
||||
if peak_mb:
|
||||
print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
|
||||
|
||||
video = _to_frame_array(_extract_video(output))
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
export_to_video(list(video), str(output_path), fps=args.fps)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
omni.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -138,7 +138,7 @@ def _extract_peak_memory_mb(result: Any) -> float:
|
||||
return 0.0
|
||||
val = getattr(result, "peak_memory_mb", 0.0)
|
||||
if not val:
|
||||
inner = getattr(result, "request_output", None)
|
||||
inner = result
|
||||
if isinstance(inner, list):
|
||||
inner = inner[0] if inner else None
|
||||
val = getattr(inner, "peak_memory_mb", 0.0)
|
||||
|
||||
@@ -605,6 +605,13 @@ python qwen3_tts/streaming_speech_client.py \
|
||||
The client writes one PCM file per sentence and a matching
|
||||
`sentence_XXX_timestamps.json` sidecar.
|
||||
|
||||
Non-streaming requests can also ask for timestamps: pass
|
||||
`"word_timestamps": true` to `POST /v1/audio/speech` and read the
|
||||
`X-Word-Timestamps` response header (JSON list of `{word, start_ms, end_ms}`,
|
||||
ASCII-escaped). Past 4 KB the header is replaced by
|
||||
`X-Word-Timestamps-Omitted: oversize; bytes=<n>; limit=4096` and the audio
|
||||
still returns — use the WebSocket path for long transcripts.
|
||||
|
||||
To *see* the alignment instead of reading a JSON sidecar, run the
|
||||
word-timestamp Gradio demo (server must be launched with `--forced-aligner`):
|
||||
```bash
|
||||
|
||||
@@ -73,6 +73,8 @@ plugins:
|
||||
- awesome-nav
|
||||
- redirects:
|
||||
redirect_maps:
|
||||
community/meetups.md: community/meetings.md
|
||||
community/volunteers.md: community/contact_us.md
|
||||
user_guide/diffusion/request_batching.md: user_guide/diffusion/execution_modes.md
|
||||
user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
|
||||
user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
|
||||
@@ -111,6 +113,7 @@ plugins:
|
||||
- "vllm_omni.entrypoints.openai" # avoid importing vllm in mkdocs building
|
||||
- "vllm_omni.entrypoints.openai.protocol" # avoid importing vllm in mkdocs building
|
||||
- "vllm_omni.experimental" # avoid importing optional serving deps (openai, fastapi, uvicorn, httpx) in mkdocs building
|
||||
- "vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored" # vendored NeMo code; upstream docstrings trip griffe strict warnings
|
||||
- mkdocstrings:
|
||||
handlers:
|
||||
python:
|
||||
|
||||
@@ -206,6 +206,12 @@ ignore = [
|
||||
# These files keep upstream conventions (star imports, long lines, math-symbol
|
||||
# argument names) so they stay easy to diff against upstream releases.
|
||||
"vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py" = ["E402", "E501", "F403", "F405", "N803"]
|
||||
# NemotronVoiceChat inference modules vendored from NVIDIA-NeMo/Speech (branch
|
||||
# nemotron-labs-voicechat). Numerics-bearing code is kept byte-close to
|
||||
# upstream to stay diffable and parity-safe, so upstream style (long lines,
|
||||
# legacy typing, whitespace, unused-import shims) is exempted. TID251 and the
|
||||
# F-correctness rules still apply.
|
||||
"vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/**" = ["E501", "E721", "E741", "F401", "F841", "I001", "N803", "N814", "UP006", "UP007", "UP008", "UP011", "UP032", "UP035", "W291", "W293"]
|
||||
"vllm_omni/experimental/fullduplex/joyvl/decision/prompts.py" = ["E501"] # System/summarizer prompts are long prose lines
|
||||
|
||||
[tool.ruff.lint.flake8-tidy-imports.banned-api]
|
||||
@@ -314,6 +320,10 @@ extend-ignore-identifiers-re = [
|
||||
".*[Nn][Dd]3.*",
|
||||
".*updator.*",
|
||||
"byt5.*",
|
||||
# NemotronVoiceChat vendored triton kernel: dout/DOUT abbreviate d_out
|
||||
# (output feature dim), not "doubt".
|
||||
".*dout.*",
|
||||
".*DOUT.*",
|
||||
]
|
||||
extend-ignore-re = [
|
||||
"\\.\\.\\.n,d->\\.\\.\\.nd",
|
||||
@@ -334,6 +344,9 @@ byt5 = "byt5"
|
||||
# "seeked" is the standard HTML media element DOM event name (fires when a
|
||||
# seek completes), used by the word-timestamps Gradio demo. Not a typo.
|
||||
seeked = "seeked"
|
||||
# Ryan Leary: author credited in the NemotronVoiceChat vendored NeMo
|
||||
# features.py copyright header. A surname, not a typo of "leery".
|
||||
Leary = "Leary"
|
||||
# IndexTTS 2.5 language codes and speech-emotion-recognition event token.
|
||||
fo = "fo"
|
||||
ba = "ba"
|
||||
|
||||
+72
-46
@@ -4,11 +4,20 @@
|
||||
|
||||
## Pipelines
|
||||
|
||||
| Model | Checkpoint | Task | `--model-class-name` | Batching |
|
||||
|---|---|---|---|---|
|
||||
| LTX-2 | `Lightricks/LTX-2` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
|
||||
| LTX-2 distilled | `rootonchair/LTX-2-19b-distilled` | Two-stage T2V/I2V | `LTX2DistilledPipeline` | No |
|
||||
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
|
||||
| `--model-class-name` | Task | Required checkpoint repositories |
|
||||
|---|---|---|
|
||||
| `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` |
|
||||
| `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` |
|
||||
| `LTX2DistilledPipeline` | LTX-2 full-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
|
||||
| `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` |
|
||||
| `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers`<br>`Lightricks/LTX-2.3` |
|
||||
| `LTX2DistilledPipeline` | LTX-2.3 full-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |
|
||||
|
||||
Repositories in the table are download units. A full pipeline repository
|
||||
contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler,
|
||||
and tokenizer; an additional repository supplies LoRA or upsampler sidecars.
|
||||
The runtime first searches the model root for sidecars under their official
|
||||
filenames, then downloads them from the matching Lightricks Hub repository.
|
||||
|
||||
`LTX2Pipeline` is the unified one-stage entry. Checkpoint metadata selects the
|
||||
LTX-2 or LTX-2.3 profile; omitting an image selects T2V, while one initial
|
||||
@@ -16,10 +25,10 @@ image selects I2V. Both one-stage repositories declare this class, so
|
||||
`--model-class-name` is optional. LTX-2.3 requires the Diffusers checkpoint;
|
||||
the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable.
|
||||
|
||||
`LTX2DistilledPipeline` is the unified distilled entry: omit `image` for T2V
|
||||
or provide one initial image for I2V. As in the official LTX pipelines,
|
||||
distilled mode is selected explicitly rather than inferred from the checkpoint
|
||||
name or metadata; always pass `--model-class-name LTX2DistilledPipeline`.
|
||||
`LTX2TwoStagePipeline` samples the regular model at half resolution, upsamples,
|
||||
then refines with the distilled LoRA. `LTX2DistilledPipeline` uses a fully
|
||||
merged distilled Transformer in both stages. Both entries support T2V and I2V;
|
||||
select their class explicitly.
|
||||
|
||||
## API Migration
|
||||
|
||||
@@ -68,24 +77,28 @@ The model recipe also supplies the default negative prompt. Top-level
|
||||
API defaults `num_frames` to `1`, so set it explicitly; the offline LTX scripts
|
||||
default to `121`.
|
||||
|
||||
## Distilled Two-Stage Defaults
|
||||
## Two-Stage Defaults
|
||||
|
||||
| Parameter | Value |
|
||||
|---|---:|
|
||||
| Final width × height | 1536 × 1024 |
|
||||
| Stage 1 width × height | 768 × 512 |
|
||||
| Frames / frame rate | 121 / 24 |
|
||||
| Stage 1 / Stage 2 denoise steps | 8 / 3 |
|
||||
| Guidance | Fixed positive-only |
|
||||
| Parameter | Ordinary | Full-distilled |
|
||||
|---|---:|---:|
|
||||
| Final width × height | 1536 × 1024 | 1536 × 1024 |
|
||||
| Stage 1 width × height | 768 × 512 | 768 × 512 |
|
||||
| Frames / frame rate | 121 / 24 | 121 / 24 |
|
||||
| Stage 1 / Stage 2 steps | 40 (LTX-2) or 30 (LTX-2.3) / 3 | 8 / 3 |
|
||||
| Guidance | Stage 1 guided; Stage 2 positive-only | Fixed positive-only |
|
||||
|
||||
The dimensions passed to the API are the final output dimensions and must be
|
||||
divisible by 64. Both one- and two-stage requests require `num_frames = 8k+1`.
|
||||
The distilled recipe fixes both sigma schedules; `num_inference_steps` must be
|
||||
`8` when supplied.
|
||||
API dimensions are final dimensions and must be divisible by 64. All LTX
|
||||
requests require `num_frames = 8k+1`. Ordinary Stage 1 uses the LTX-2 or
|
||||
LTX-2.3 one-stage defaults shown above. Distilled schedules are fixed, so
|
||||
`num_inference_steps`, when supplied, must be `8`. Both entries reject custom
|
||||
sigmas and input latents.
|
||||
|
||||
Ordinary two-stage uses layer-fused LoRA for an unquantized BF16 Transformer
|
||||
and automatically switches to dynamic LoRA when quantization is enabled.
|
||||
|
||||
## Serving
|
||||
|
||||
Start either checkpoint:
|
||||
Start either one-stage checkpoint:
|
||||
|
||||
```bash
|
||||
vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
|
||||
@@ -95,11 +108,19 @@ vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
|
||||
vllm serve diffusers/LTX-2.3-Diffusers --omni --stage-init-timeout 600
|
||||
```
|
||||
|
||||
Start the distilled checkpoint with its two-stage entry:
|
||||
Start a two-stage checkpoint with its explicit class:
|
||||
|
||||
```bash
|
||||
vllm serve rootonchair/LTX-2-19b-distilled --omni \
|
||||
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
|
||||
# LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent
|
||||
vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \
|
||||
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
|
||||
# Ordinary LTX-2.3 two-stage
|
||||
vllm serve diffusers/LTX-2.3-Diffusers --omni \
|
||||
--model-class-name LTX2TwoStagePipeline \
|
||||
--enable-layerwise-offload \
|
||||
--stage-init-timeout 600
|
||||
```
|
||||
|
||||
The same server handles T2V and I2V. A T2V request using the selected recipe's
|
||||
@@ -135,8 +156,9 @@ together with `input_reference`.
|
||||
|
||||
## Guidance
|
||||
|
||||
LTX supports independent video/audio CFG, spatio-temporal guidance (STG),
|
||||
cross-modality guidance, and rescaling:
|
||||
One-stage and ordinary Stage 1 support independent video/audio CFG,
|
||||
spatio-temporal guidance (STG), cross-modality guidance, and rescaling.
|
||||
Distilled stages and ordinary Stage 2 are fixed positive-only.
|
||||
|
||||
| Parameter | Default | Effect | Alias |
|
||||
|---|---:|---|---|
|
||||
@@ -172,9 +194,9 @@ complete guidance plan in parallel. Despite the legacy `cfg` name, it covers
|
||||
text CFG, STG, and cross-modality guidance passes; guidance rescaling is
|
||||
applied after all pass predictions are gathered.
|
||||
|
||||
The default one-stage recipe has four Transformer passes per denoise step:
|
||||
`cond`, `uncond`, `ptb` (STG), and `mod` (cross-modality). The useful balanced
|
||||
configurations are therefore:
|
||||
The default one-stage recipe and ordinary Stage 1 have four Transformer passes
|
||||
per denoise step: `cond`, `uncond`, `ptb` (STG), and `mod`
|
||||
(cross-modality). The useful balanced configurations are therefore:
|
||||
|
||||
| `--cfg-parallel-size` | Passes per rank | Guidance-slot utilization | Notes |
|
||||
|---:|---:|---:|---|
|
||||
@@ -195,9 +217,8 @@ vllm serve diffusers/LTX-2.3-Diffusers --omni \
|
||||
```
|
||||
|
||||
The total device count is the product of `cfg_parallel_size` and the other
|
||||
configured parallel dimensions. Distilled pipelines use positive-only
|
||||
guidance, so `cfg_parallel_size > 1` does not provide guidance speedup for
|
||||
them.
|
||||
configured parallel dimensions. Positive-only distilled phases do not benefit
|
||||
from `cfg_parallel_size > 1`.
|
||||
|
||||
### Python API
|
||||
|
||||
@@ -242,22 +263,22 @@ noted below.
|
||||
| `image` | image or batch, `None` | Direct value wins over request images; no image selects T2V. I2V accepts one image per prompt, and a batch cannot mix T2V/I2V. |
|
||||
| `prompt` | string or list, `None` | Positive-text fallback; request prompts win. Mutually exclusive with `prompt_embeds`. |
|
||||
| `negative_prompt` | string or list, `None` | Fallback after request values and before the recipe default; mutually exclusive with negative embeddings. |
|
||||
| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 distilled. |
|
||||
| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 two-stage. |
|
||||
| `width` | `int`, `None` | Same precedence and alignment as `height`. |
|
||||
| `num_frames` | `int`, `None` | Request → direct value → recipe default; must be `8k+1` and also determines audio duration with `frame_rate`. |
|
||||
| `frame_rate` | `float`, `None` | Request `frame_rate` → request `fps` → direct value → recipe default. |
|
||||
| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
|
||||
| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage and ordinary Stage 1, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
|
||||
| `sigmas` | list of float, `None` | One-stage only. Request values win; every request in a fused batch must use the same schedule. |
|
||||
| `timesteps` | list of int, `None` | Compatibility slot; LTX accepts only `None`. Use `sigmas`. |
|
||||
| `guidance_scale` | `float`, `None` | One-stage common video/audio CFG fallback; an explicit request value wins. Distilled two-stage uses fixed positive-only guidance. |
|
||||
| `guidance_scale` | `float`, `None` | One-stage and ordinary Stage 1 common video/audio CFG fallback; an explicit request value wins. Distilled phases use fixed positive-only guidance. |
|
||||
| `guidance_rescale` | `float`, `None` | Accepts only `None` or `0.0`; use the modality rescale fields. |
|
||||
| `noise_scale` | `float`, `0.0` | Compatibility slot; LTX accepts only `0.0`. |
|
||||
| `num_videos_per_prompt` | `int`, `1` | Output-count fallback; positive request `num_outputs_per_prompt` wins. |
|
||||
| `generator` | generator or list, `None` | Explicit RNG; otherwise request generators/seeds are collated. Lists must match the effective output batch. |
|
||||
| `latents` | tensor, `None` | Request tensors win. One-stage accepts packed `[B, S, C]` and validated unpacked video layouts. |
|
||||
| `audio_latents` | tensor, `None` | One-stage initial audio latents; request tensors win and are collated. |
|
||||
| `latents` | tensor, `None` | One-stage only. Request tensors win; packed `[B, S, C]` and validated unpacked video layouts are accepted. |
|
||||
| `audio_latents` | tensor, `None` | One-stage only. Request tensors win and are collated. |
|
||||
| `prompt_embeds` | tensor, `None` | Precomputed positive conditioning; requires `prompt_attention_mask` and cannot accompany `prompt`. |
|
||||
| `negative_prompt_embeds` | tensor, `None` | One-stage precomputed negative conditioning; requires its mask and cannot accompany `negative_prompt`. |
|
||||
| `negative_prompt_embeds` | tensor, `None` | One-stage and ordinary Stage 1 conditioning; requires its mask and cannot accompany `negative_prompt`. |
|
||||
| `prompt_attention_mask` | tensor, `None` | Mask for positive embeddings; request `prompt_attention_mask` or `attention_mask` wins. |
|
||||
| `negative_prompt_attention_mask` | tensor, `None` | Mask for negative embeddings; request negative mask fields win. |
|
||||
| `decode_timestep` | float or list, `0.0` | Video-VAE decode timestep; request value wins. Lists may match 1, prompt batch, or output batch. |
|
||||
@@ -273,13 +294,13 @@ request prompt payload; LTX guidance fields live in sampling `extra_args`.
|
||||
|
||||
### Recipe-Specific Request Capabilities
|
||||
|
||||
| Override | One-stage | Distilled two-stage |
|
||||
|---|---|---|
|
||||
| Guidance | Supported | Fixed positive-only |
|
||||
| Negative prompt/embeddings | Supported | Rejected |
|
||||
| `num_inference_steps` | Supported | Fixed at 8 for Stage 1; Stage 2 uses 3 |
|
||||
| Custom `sigmas` | Supported | Rejected; both phases use fixed schedules |
|
||||
| Video/audio latents | Supported | Rejected |
|
||||
| Override | One-stage | Ordinary two-stage | Distilled two-stage |
|
||||
|---|---|---|---|
|
||||
| Guidance | Supported | Stage 1 only; Stage 2 is positive-only | Fixed positive-only |
|
||||
| Negative prompt/embeddings | Supported | Supported by Stage 1 | Rejected |
|
||||
| `num_inference_steps` | Supported | Controls Stage 1; Stage 2 uses 3 | Fixed at 8 for Stage 1; Stage 2 uses 3 |
|
||||
| Custom `sigmas` | Supported | Rejected | Rejected; both phases use fixed schedules |
|
||||
| Video/audio latents | Supported | Rejected | Rejected |
|
||||
|
||||
These capability checks apply equally to direct `forward` keywords and
|
||||
values in `OmniDiffusionSamplingParams`; unsupported values fail instead of
|
||||
@@ -309,6 +330,10 @@ bundled offline CLI do not currently expose `sigmas`.
|
||||
`--max-num-seqs 1`.
|
||||
- `--cfg-parallel-size` shards the complete LTX guidance plan, including STG,
|
||||
modality guidance, and rescale-compatible prediction gathering.
|
||||
- Sequence parallelism may pad audio latents. Pure Ulysses masks the padding;
|
||||
Ring cannot, so audio length must be SP-divisible. Use `ring_degree=1` or a
|
||||
divisible request shape.
|
||||
- Cache-DiT is one-stage only; multi-stage configurations are rejected.
|
||||
|
||||
## Operational Notes
|
||||
|
||||
@@ -323,14 +348,15 @@ bundled offline CLI do not currently expose `sigmas`.
|
||||
request parameter.
|
||||
- For benchmarks, use `tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with
|
||||
`tests/dfx/perf/scripts/run_diffusion_benchmark.py`.
|
||||
- Current two-stage support is limited to the distilled checkpoint; official
|
||||
LTX two-stage/HQ execution remains out of scope.
|
||||
- Ordinary and full-distilled two-stage T2V/I2V are supported; HQ execution
|
||||
remains out of scope.
|
||||
|
||||
## References
|
||||
|
||||
- <https://huggingface.co/Lightricks/LTX-2>
|
||||
- <https://huggingface.co/Lightricks/LTX-2.3>
|
||||
- <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
|
||||
- <https://huggingface.co/diffusers/LTX-2.3-Distilled-Diffusers>
|
||||
- [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md)
|
||||
- [Diffusion execution modes](../../docs/user_guide/diffusion/execution_modes.md)
|
||||
- [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md)
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
# NemotronLabs VoiceChat 11B
|
||||
|
||||
> Offline speech-to-speech: a 16 kHz user utterance in, the agent's spoken reply
|
||||
> (text + 22.05 kHz WAV) out, on a 3-stage vLLM-Omni pipeline.
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: NVIDIA
|
||||
- Model: [`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
|
||||
- Task: Speech-to-speech voice chat. The model runs a frame-locked 12.5 Hz
|
||||
timeline: a Conformer + NemotronH hybrid-Mamba thinker emits one text token
|
||||
per acoustic frame, a Gemma3-1B EAR-TTS talker turns the text timeline into
|
||||
31-quantizer RVQ code stacks, and an RVQ-VAE codec decodes them to audio.
|
||||
- Mode: Offline single-turn inference (batch=1). Online/duplex serving,
|
||||
function calling, and batch>1 are documented follow-ups.
|
||||
- Maintainer: [`@yuekaizhang`](https://github.com/yuekaizhang)
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe to run a single-turn voice-chat exchange with
|
||||
`NVIDIA-NemotronLabs-VoiceChat-11B` on one GPU: you provide a user utterance as
|
||||
a WAV file (any sample rate; it is resampled to 16 kHz mono) plus an optional
|
||||
spoken-style system prompt, and get back the agent's reply as text and a
|
||||
22.05 kHz WAV. The integration is NeMo-free at runtime — the perception
|
||||
Conformer, EAR-TTS talker, and RVQ-VAE codec are vendored, dependency-stripped
|
||||
NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed.
|
||||
|
||||
## References
|
||||
|
||||
- Offline example:
|
||||
[`examples/offline_inference/nemotron_voicechat/end2end.py`](../../examples/offline_inference/nemotron_voicechat/end2end.py)
|
||||
- Model modules (thinker / talker / code2wav / vendored NeMo):
|
||||
[`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/)
|
||||
- Staged pipeline config:
|
||||
[`vllm_omni/deploy/nemotron_labs_voicechat.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat.yaml)
|
||||
- Upstream model card:
|
||||
[`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
|
||||
- Reference implementation: NVIDIA-NeMo/Speech, branch `nemotron-labs-voicechat`
|
||||
|
||||
## Pipeline
|
||||
|
||||
| stage | arch | dtype | role |
|
||||
|---|---|---|---|
|
||||
| 0 thinker | `NemotronVoiceChatThinkerForConditionalGeneration` (LLM_AR) | fp32 | WAV + system prompt -> frame-locked text-token timeline (+ function channel) |
|
||||
| 1 talker | `NemotronVoiceChatTalker` (LLM_AR) | fp32 | text timeline -> 31-quantizer RVQ code stacks (one per 80 ms frame) |
|
||||
| 2 code2wav | `NemotronVoiceChatCode2Wav` (LLM_GENERATION) | fp32 | RVQ-VAE decode -> 22.05 kHz PCM |
|
||||
|
||||
Stages 0/1 default to fp32 for exact parity with the NeMo reference
|
||||
implementation (greedy decoding matches it token for token on the acceptance
|
||||
fixture). The deploy yaml documents a ~2x-faster bf16 thinker option whose
|
||||
output stayed within one word of the reference in testing.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
## GPU
|
||||
|
||||
### 1x H100 80GB
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.12
|
||||
- vLLM version: 0.26.0
|
||||
- vLLM-Omni version or commit: this PR / current `main`
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
# Tokenizer: the checkpoint ships no HF tokenizer; it resolves from the
|
||||
# nvidia/NVIDIA-Nemotron-Nano-9B-v2 HF id automatically. For air-gapped runs,
|
||||
# point NEMOTRON_VOICECHAT_LLM_PATH at a local snapshot of that repo instead.
|
||||
python examples/offline_inference/nemotron_voicechat/end2end.py \
|
||||
--checkpoint /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
|
||||
--wav /path/to/user_question.wav \
|
||||
--output-dir results/nemotron_voicechat
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
```bash
|
||||
ls results/nemotron_voicechat
|
||||
# <stem>_output.txt the agent reply as text
|
||||
# <stem>_output.wav the agent reply as 22.05 kHz audio
|
||||
# <stem>_text_tokens.json the frame-locked text-token timeline
|
||||
```
|
||||
|
||||
The reply text should read as a coherent spoken-style answer to the question in
|
||||
the input WAV, and the WAV should transcribe to (approximately) the same text
|
||||
with any ASR model.
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory usage: the shipped yaml runs all three stages on one GPU
|
||||
(`gpu_memory_utilization` 0.62 / 0.12 / 0.06); peak usage is dominated by the
|
||||
fp32 thinker. The fp32 default has a hard floor of roughly 43 GB of thinker
|
||||
weights alone (9B backbone + 587M `embed_tokens` + 587M `function_head` +
|
||||
0.6B Conformer), so 48 GB cards cannot run it — use the bf16 thinker option
|
||||
documented in the deploy yaml on anything smaller than an 80 GB part.
|
||||
- Input sizing: the timeline is frame-locked, so the reply budget IS the input
|
||||
duration. The acoustic channel trails the text channel; if the WAV does not
|
||||
carry enough trailing silence for the reply to finish, the spoken answer is
|
||||
truncated silently. Leave generous trailing silence (a question ending at
|
||||
~4.5 s truncated in an 8 s WAV but completed cleanly in 16 s); the offline
|
||||
example warns when the text channel is still speaking near the last frame.
|
||||
- Key flags: sampling is greedy end to end. The thinker is frame-locked —
|
||||
`max_tokens` equals the acoustic frame count with `ignore_eos=True`. Do NOT
|
||||
set `min_tokens` on the thinker: the tokenizer's EOS token is also the
|
||||
frame-locked PAD/silence token, so masking it forces the model to babble
|
||||
instead of pausing.
|
||||
- The talker's `max_tokens` is 16383 (its stage prompt is one placeholder
|
||||
token, and the stage context is 16384).
|
||||
- Known limitations: batch=1 only; offline single turn (no online/duplex
|
||||
serving yet); the function-call channel is decoded but not yet acted on.
|
||||
@@ -66,6 +66,7 @@ recipes/
|
||||
| [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
|
||||
| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
|
||||
| [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
|
||||
| [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
|
||||
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
|
||||
| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
|
||||
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
|
||||
|
||||
@@ -27,7 +27,7 @@ parallel, cache, quantized, or expert-kernel backends.
|
||||
- Dense checkpoint: <https://huggingface.co/robbyant/lingbot-video-dense-1.3b>
|
||||
- MoE checkpoint: <https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b>
|
||||
- Upstream project: <https://github.com/Robbyant/lingbot-video>
|
||||
- Related offline example: [`examples/offline_inference/text_to_video/text_to_video_lingbot.py`](../../examples/offline_inference/text_to_video/text_to_video_lingbot.py)
|
||||
- Related offline example: [`examples/offline_inference/text_to_video/text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py)
|
||||
- Related online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
|
||||
|
||||
## Hardware Support
|
||||
@@ -52,7 +52,7 @@ The MoE path is validated with BF16 expert weights.
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 \
|
||||
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model robbyant/lingbot-video-dense-1.3b \
|
||||
--prompt "a robotic arm picks up a red block" \
|
||||
--output lingbot_t2v.mp4 \
|
||||
@@ -70,7 +70,7 @@ python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 \
|
||||
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model robbyant/lingbot-video-moe-30b-a3b \
|
||||
--prompt "a robotic arm picks up a red block" \
|
||||
--output lingbot_moe_t2v.mp4 \
|
||||
|
||||
@@ -82,7 +82,7 @@ example and the DiT sections below.
|
||||
- Optional environment variables:
|
||||
|
||||
```bash
|
||||
export CACHE_DIT_VERSION=1.3.0
|
||||
export CACHE_DIT_VERSION=1.5.0
|
||||
```
|
||||
|
||||
HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA`
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
# dots.tts
|
||||
|
||||
> Offline continuous-AR TTS at 48 kHz (rednote-hilab)
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: rednote-hilab
|
||||
- Model: `rednote-hilab/dots.tts-soar`
|
||||
- Task: Text-to-speech, zero-shot synthesis only
|
||||
- Mode: Offline end-to-end example only (not yet wired for online serving —
|
||||
see [Known limitations](#known-limitations))
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe as a known-good starting point for running
|
||||
`rednote-hilab/dots.tts-soar` offline on vLLM-Omni on consumer-class GPUs.
|
||||
dots.tts is a ~1.7B-parameter continuous-AR TTS model (Qwen2.5-1.5B base LM
|
||||
+ 344M DiT flow-matching head + 180M AudioVAE) that emits 48 kHz mono audio.
|
||||
It follows the same "vLLM-native base LM + side-path computation" pattern as
|
||||
VoxCPM2 — single-stage pipeline
|
||||
`Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR
|
||||
loopback → AudioVAE (streaming decode)` — with a plain Qwen2 backbone
|
||||
instead of MiniCPM4, and no FSQ / residual-LM stage.
|
||||
|
||||
This is an early integration (see [Known limitations](#known-limitations)
|
||||
below before depending on it for anything beyond offline experimentation).
|
||||
|
||||
## References
|
||||
|
||||
- Offline end-to-end script:
|
||||
[`examples/offline_inference/text_to_speech/dots_tts/end2end.py`](../../examples/offline_inference/text_to_speech/dots_tts/end2end.py)
|
||||
- Example guide:
|
||||
[`examples/offline_inference/text_to_speech/README.md`](../../examples/offline_inference/text_to_speech/README.md#dotstts)
|
||||
- Default deploy config:
|
||||
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
|
||||
- Talker / pipeline source:
|
||||
[`vllm_omni/model_executor/models/dots_tts/`](../../vllm_omni/model_executor/models/dots_tts/)
|
||||
- Upstream: [rednote-hilab/dots.tts](https://github.com/rednote-hilab/dots.tts)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
This recipe documents one tested 16 GB consumer-GPU configuration. Other
|
||||
vendor sections (ROCm, NPU) and larger-VRAM configurations are welcome as
|
||||
community validation lands.
|
||||
|
||||
## GPU
|
||||
|
||||
### 1 x RTX 5080 16GB (Single GPU, Minimum Recommended)
|
||||
|
||||
dots.tts (~1.7B params across the base LM + DiT + AudioVAE + CAM++
|
||||
speaker encoder, bfloat16) fits comfortably on a single 16 GB GPU. The
|
||||
bundled default config at
|
||||
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
|
||||
(`gpu_memory_utilization: 0.8`, `max_num_seqs: 4`, `enforce_eager: true`,
|
||||
`enable_prefix_caching: false`) loads cleanly with ~5.1 GiB for model
|
||||
weights and ~0.3 GiB peak activation; the remainder of the configured
|
||||
budget is available for KV cache. Total resident footprint at idle is
|
||||
roughly **7-8 GiB / 16 GB** — the only tight spot in the full CUDA-Graph
|
||||
roadmap would be step 8's graph capture (not implemented yet; this
|
||||
release runs `enforce_eager: true`, so it doesn't apply today).
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux (WSL2)
|
||||
- Python: 3.12
|
||||
- Driver / runtime: NVIDIA driver 595.95
|
||||
- torch: 2.11.0+cu130
|
||||
- vLLM: 0.26.0
|
||||
- vLLM-Omni: 0.22.1.dev (current `main`)
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
|
||||
--model rednote-hilab/dots.tts-soar \
|
||||
--text "Hello, this is a test of dots TTS running on vLLM Omni."
|
||||
```
|
||||
|
||||
The deploy config at
|
||||
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
|
||||
is loaded automatically by the model registry (HF `model_type=dots_tts`).
|
||||
Pass `--deploy-config <path>` to override.
|
||||
|
||||
#### Verification
|
||||
|
||||
**T1 — offline zero-shot synthesis**:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
|
||||
--model rednote-hilab/dots.tts-soar \
|
||||
--text "Hello, this is a test of dots TTS running on vLLM Omni."
|
||||
```
|
||||
|
||||
Observed: `output_audio/output.wav`, 3.52 s @ 48 kHz mono. Single
|
||||
one-shot process (init → one `generate()` → exit), so the reported
|
||||
numbers include engine init and first-request warmup, not just steady-state
|
||||
per-step throughput — same caveat as VoxCPM2's recipe. `Inference: 5.52s`,
|
||||
`RTF: 1.569`.
|
||||
|
||||
Weight-loading breakdown from the same run (all tensors matched, no
|
||||
missing/extra keys): 951/951 AudioVAE, 244/244 DiT, 270/270
|
||||
patch_encoder, 198 Qwen2, 938/938 CAM++ speaker encoder.
|
||||
|
||||
Whisper transcription of the output matched the input text with no
|
||||
dropped leading word (confirms the streaming-vocoder patch-boundary fix
|
||||
described in [Known limitations](#known-limitations)).
|
||||
|
||||
#### Notes
|
||||
|
||||
- Output: 48 kHz mono WAV.
|
||||
- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default
|
||||
used throughout this recipe. `rednote-hilab/dots.tts-base` shares the
|
||||
same architecture but is unvalidated in this repo. `rednote-hilab/dots.tts-mf`
|
||||
(MeanFlow, 2-4 step) is not supported — see below.
|
||||
- `enforce_eager: true` and `enable_prefix_caching: false` in the deploy
|
||||
config are load-bearing, not just conservative defaults: with prefix
|
||||
caching enabled, vLLM-Omni's prefix-cache multimodal-output merge path
|
||||
does not preserve the `sparse_audio` marker this model relies on to
|
||||
route audio output correctly, and generation silently truncates to a
|
||||
single ~160 ms patch. Do not override `enable_prefix_caching` for this
|
||||
model until that framework-level gap is fixed.
|
||||
|
||||
## Known limitations
|
||||
|
||||
- **No online serving yet.** dots.tts is not registered in
|
||||
`vllm_omni/entrypoints/openai/tts_adapters/` (the framework other TTS
|
||||
models in this repo use for `/v1/audio/speech`), so `vllm serve
|
||||
rednote-hilab/dots.tts-soar --omni` will start but the OpenAI-compatible
|
||||
speech endpoint does not know how to build request params for it yet.
|
||||
Offline `Omni()` (as used in `end2end.py`) is the only supported path
|
||||
today.
|
||||
- **Voice cloning is not wired.** The CAM++ x-vector speaker encoder
|
||||
weights load and are exercised by `load_weights()`, but nothing in the
|
||||
prompt builder consumes a reference audio yet — generation is zero-shot
|
||||
only. `end2end.py` has no `--ref-audio`/`--ref-text` flags.
|
||||
- **`dots.tts-mf` (MeanFlow, 2-4 step) checkpoint is not supported.** Only
|
||||
the fixed 10-step Euler DiT sampler used by `dots.tts-soar` /
|
||||
`dots.tts-base` is implemented.
|
||||
- **No CUDA graph capture.** The talker runs fully eager. voxcpm2's three
|
||||
captured graphs (base LM decode, CFM solver, VAE decode) have no
|
||||
dots.tts equivalent yet.
|
||||
- **Concurrent requests do not scale.** Each request's 10-step DiT Euler
|
||||
integration runs serially in the side path (no cross-request batching,
|
||||
unlike voxcpm2's `enable_batched_cfm`). A community review of this
|
||||
integration measured no throughput gain at `c=4` concurrent requests
|
||||
versus `c=1`.
|
||||
- **`SamplingParams.seed` does not control audio-generation randomness.**
|
||||
The DiT's flow-matching noise is deterministically derived per-request
|
||||
from a fixed internal seed (reproducible run-to-run, matching
|
||||
voxcpm2's `deterministic_cfm_seed` convention) rather than from the
|
||||
caller-supplied `seed` field — the same limitation voxcpm2 has today.
|
||||
@@ -16,7 +16,7 @@ diffusers==0.38.0
|
||||
safetensors>=0.8.0
|
||||
accelerate==1.12.0
|
||||
soundfile>=0.13.1
|
||||
cache-dit==1.3.0
|
||||
cache-dit==1.5.0
|
||||
tqdm>=4.66.0
|
||||
torchsde>=0.2.6
|
||||
openai-whisper>=20250625
|
||||
@@ -34,3 +34,5 @@ pydantic-settings>=2.0.0
|
||||
# base image does not ship it, so declare it explicitly. The loader also guards
|
||||
# this import so a missing gguf no longer breaks all diffusion model loading.
|
||||
gguf>=0.10.0
|
||||
# Required for running the Cosmos3 model with guardrails enabled (default behavior).
|
||||
cosmos-guardrail>=0.3.1
|
||||
@@ -0,0 +1,92 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Offline Daily-Omni sources: local mirrors must not fall back to the Hub id."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import tarfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (
|
||||
daily_omni_local_qa_json,
|
||||
daily_omni_local_videos_dir,
|
||||
resolve_daily_omni_local_root,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
|
||||
|
||||
|
||||
def _write_qa_json(root: Path) -> Path:
|
||||
qa = root / "qa.json"
|
||||
qa.write_text(json.dumps([{"Question": "q", "Answer": "A", "video_id": "vid0"}]), encoding="utf-8")
|
||||
return qa
|
||||
|
||||
|
||||
def test_resolve_local_root_returns_none_for_hub_id() -> None:
|
||||
assert resolve_daily_omni_local_root("liarliar/Daily-Omni") is None
|
||||
assert resolve_daily_omni_local_root("") is None
|
||||
assert resolve_daily_omni_local_root(None) is None
|
||||
|
||||
|
||||
def test_plain_local_mirror(tmp_path: Path) -> None:
|
||||
_write_qa_json(tmp_path)
|
||||
(tmp_path / "Videos" / "vid0").mkdir(parents=True)
|
||||
|
||||
root = resolve_daily_omni_local_root(str(tmp_path))
|
||||
assert root == tmp_path.resolve()
|
||||
assert daily_omni_local_qa_json(root) == (tmp_path / "qa.json").resolve()
|
||||
assert daily_omni_local_videos_dir(root) == root / "Videos"
|
||||
|
||||
|
||||
def test_hf_cache_dir_resolves_to_snapshot(tmp_path: Path) -> None:
|
||||
cache_dir = tmp_path / "datasets--liarliar--Daily-Omni"
|
||||
snapshot = cache_dir / "snapshots" / "deadbeef"
|
||||
snapshot.mkdir(parents=True)
|
||||
_write_qa_json(snapshot)
|
||||
(cache_dir / "refs").mkdir()
|
||||
(cache_dir / "refs" / "main").write_text("deadbeef", encoding="utf-8")
|
||||
|
||||
root = resolve_daily_omni_local_root(str(cache_dir))
|
||||
assert root == snapshot.resolve()
|
||||
assert daily_omni_local_qa_json(root) == (snapshot / "qa.json").resolve()
|
||||
|
||||
|
||||
def test_hf_cache_dir_without_refs_picks_newest_snapshot(tmp_path: Path) -> None:
|
||||
cache_dir = tmp_path / "datasets--liarliar--Daily-Omni"
|
||||
older = cache_dir / "snapshots" / "aaa"
|
||||
newer = cache_dir / "snapshots" / "bbb"
|
||||
older.mkdir(parents=True)
|
||||
newer.mkdir(parents=True)
|
||||
os.utime(older, (1, 1))
|
||||
os.utime(newer, (10, 10))
|
||||
|
||||
assert resolve_daily_omni_local_root(str(cache_dir)) == newer.resolve()
|
||||
|
||||
|
||||
def test_videos_tar_is_extracted(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HF_HOME", str(tmp_path / "hf_home"))
|
||||
root = tmp_path / "mirror"
|
||||
(root / "Videos" / "vid0").mkdir(parents=True)
|
||||
(root / "Videos" / "vid0" / "vid0_video.mp4").write_bytes(b"fake")
|
||||
with tarfile.open(root / "Videos.tar", "w") as tf:
|
||||
tf.add(root / "Videos", arcname="Videos")
|
||||
# Only the tarball is left in the mirror.
|
||||
for child in (root / "Videos").rglob("*"):
|
||||
if child.is_file():
|
||||
child.unlink()
|
||||
(root / "Videos" / "vid0").rmdir()
|
||||
(root / "Videos").rmdir()
|
||||
|
||||
videos_dir = daily_omni_local_videos_dir(root)
|
||||
assert videos_dir is not None
|
||||
assert (videos_dir / "vid0" / "vid0_video.mp4").is_file()
|
||||
|
||||
# Second call reuses the extraction instead of redoing it.
|
||||
assert daily_omni_local_videos_dir(root) == videos_dir
|
||||
|
||||
|
||||
def test_missing_qa_json_reports_none(tmp_path: Path) -> None:
|
||||
root = resolve_daily_omni_local_root(str(tmp_path))
|
||||
assert daily_omni_local_qa_json(root) is None
|
||||
assert daily_omni_local_videos_dir(root) is None
|
||||
@@ -2,9 +2,8 @@
|
||||
|
||||
The ``minicpm-interleave`` mode is an accuracy-critical protocol: MiniCPM-o only reaches
|
||||
OpenBMB's reported Daily-Omni score when frames and 1s audio segments arrive strictly
|
||||
interleaved, with ``max_slice_nums=1`` / ``use_image_id=False`` and an empty MiniCPM
|
||||
system message (not the Qwen system prompt). These tests pin that contract so a refactor
|
||||
cannot silently regress the score.
|
||||
interleaved, with ``max_slice_nums=1`` / ``use_image_id=False`` and no Qwen system prompt.
|
||||
These tests pin that contract so a refactor cannot silently regress the score.
|
||||
|
||||
vllm stubs are installed by tests/benchmarks/conftest.py before collection.
|
||||
"""
|
||||
@@ -99,22 +98,13 @@ def _stub_extract(monkeypatch, dataset: DailyOmniDataset, counter: list[int]) ->
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("n", "k"), [(10, 4), (100, 64), (1000, 64), (7, 7), (3, 5), (1, 1)])
|
||||
def test_uniform_sample_indices_matches_openbmb_midpoint(n: int, k: int) -> None:
|
||||
"""Pin OpenBMB midpoint-bin sampling (not np.linspace endpoints)."""
|
||||
@pytest.mark.parametrize(("n", "k"), [(10, 4), (100, 64), (7, 7), (3, 5), (1, 1)])
|
||||
def test_uniform_sample_indices_matches_numpy_linspace(n: int, k: int) -> None:
|
||||
got = _uniform_sample_indices(n, k)
|
||||
if n <= k:
|
||||
assert got == list(range(n))
|
||||
else:
|
||||
gap = n / k
|
||||
expected = [int(i * gap + gap / 2) for i in range(k)]
|
||||
assert got == expected
|
||||
# Endpoints differ from linspace for typical Daily-Omni downsample sizes.
|
||||
linspace = [int(i) for i in np.linspace(0, n - 1, k, dtype=int)]
|
||||
if n == 1000 and k == 64:
|
||||
assert got[0] == 7 and got[-1] == 992
|
||||
assert linspace[0] == 0 and linspace[-1] == 999
|
||||
assert got != linspace
|
||||
assert got == [int(i) for i in np.linspace(0, n - 1, k, dtype=int)]
|
||||
|
||||
|
||||
def test_uniform_sample_indices_non_positive_k() -> None:
|
||||
@@ -223,16 +213,14 @@ def test_minicpm_interleave_missing_local_video_returns_none(qa_json, tmp_path)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_minicpm_interleave_keeps_empty_system_message(monkeypatch, qa_json, video_dir) -> None:
|
||||
"""Official MiniCPM omni mode still sends role=system with empty content."""
|
||||
def test_minicpm_interleave_omits_qwen_system_message(monkeypatch, qa_json, video_dir) -> None:
|
||||
ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all")
|
||||
_stub_extract(monkeypatch, ds, [])
|
||||
parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
|
||||
|
||||
messages = ds._build_daily_omni_openai_messages(parts, "q?", {"A": "a"})
|
||||
|
||||
assert [m["role"] for m in messages] == ["system", "user"]
|
||||
assert messages[0]["content"][0]["text"] == ""
|
||||
assert [m["role"] for m in messages] == ["user"]
|
||||
|
||||
|
||||
def test_qwen_pack_mode_keeps_system_message(qa_json, video_dir) -> None:
|
||||
|
||||
@@ -6,7 +6,7 @@ Unit tests for StageConfigFactory and related classes.
|
||||
|
||||
import importlib
|
||||
import warnings
|
||||
from dataclasses import dataclass, fields
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
@@ -34,9 +34,8 @@ from vllm_omni.config.stage_config import (
|
||||
load_deploy_config,
|
||||
merge_pipeline_deploy,
|
||||
pipeline_cfg_resolver,
|
||||
strip_parent_engine_args,
|
||||
)
|
||||
from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys
|
||||
from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -475,7 +474,7 @@ class TestStageConfigFactory:
|
||||
cli_overrides = {
|
||||
"gpu_memory_utilization": 0.9,
|
||||
"model": "some_model", # Internal
|
||||
"stage_configs_path": "/path", # Internal
|
||||
"deploy_config": "/path", # Internal
|
||||
"batch_timeout": 10, # Internal
|
||||
}
|
||||
|
||||
@@ -483,7 +482,7 @@ class TestStageConfigFactory:
|
||||
|
||||
assert overrides["gpu_memory_utilization"] == 0.9
|
||||
assert "model" not in overrides
|
||||
assert "stage_configs_path" not in overrides
|
||||
assert "deploy_config" not in overrides
|
||||
assert "batch_timeout" not in overrides
|
||||
|
||||
def test_per_stage_override_excludes_internal_keys(self):
|
||||
@@ -525,49 +524,6 @@ class TestStageResolutionHelpers:
|
||||
assert "model" not in overrides
|
||||
assert "parallel_config" not in overrides
|
||||
|
||||
def test_strip_parent_engine_args_reports_only_surprising_parent_overrides(self):
|
||||
parent_fields = {f.name: f for f in fields(EngineArgs)}
|
||||
filtered, overridden = strip_parent_engine_args(
|
||||
{
|
||||
"model": "some/model",
|
||||
"stage_configs_path": "/tmp/stages.yaml",
|
||||
"tensor_parallel_size": 4,
|
||||
"worker_extension_cls": "some.Extension",
|
||||
"custom_pipeline_args": {"pipeline_class": "demo.Pipeline"},
|
||||
},
|
||||
parent_fields=parent_fields,
|
||||
keep_keys={"worker_extension_cls"},
|
||||
strip_keys={"stage_configs_path"},
|
||||
no_warn_keys={"model"},
|
||||
)
|
||||
|
||||
assert filtered == {
|
||||
"worker_extension_cls": "some.Extension",
|
||||
"custom_pipeline_args": {"pipeline_class": "demo.Pipeline"},
|
||||
}
|
||||
assert overridden == ["tensor_parallel_size"]
|
||||
|
||||
def test_strip_parent_engine_args_keeps_allowed_media_access_controls(self):
|
||||
parent_fields = {f.name: f for f in fields(EngineArgs)}
|
||||
filtered, overridden = strip_parent_engine_args(
|
||||
{
|
||||
"model": "some/model",
|
||||
"stage_configs_path": "/tmp/stages.yaml",
|
||||
"allowed_local_media_path": "/data/qwentts",
|
||||
"allowed_media_domains": ["example.com"],
|
||||
},
|
||||
parent_fields=parent_fields,
|
||||
keep_keys={"allowed_local_media_path", "allowed_media_domains"},
|
||||
strip_keys={"stage_configs_path"},
|
||||
no_warn_keys={"model"},
|
||||
)
|
||||
|
||||
assert filtered == {
|
||||
"allowed_local_media_path": "/data/qwentts",
|
||||
"allowed_media_domains": ["example.com"],
|
||||
}
|
||||
assert overridden == []
|
||||
|
||||
|
||||
class TestPipelineDiscovery:
|
||||
"""Tests for the central pipeline registry (``OMNI_PIPELINES``)."""
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
import pytest
|
||||
|
||||
import vllm_omni.config.pipeline_registry # noqa: F401 (populate registry)
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES as _PIPELINE_REGISTRY
|
||||
from vllm_omni.config.stage_config import (
|
||||
DeployConfig,
|
||||
StageExecutionType,
|
||||
)
|
||||
from vllm_omni.model_executor.stage_input_processors.forced_aligner import (
|
||||
POOLING_OUTPUT_DECODER_PATH,
|
||||
TIMESTAMPS_MODALITY,
|
||||
)
|
||||
from vllm_omni.utils.forced_aligner import _DEFAULT_CONFIG_PATH, inject_forced_aligner_stage
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_inject_appends_pooling_aligner_stage():
|
||||
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
|
||||
deploy = DeployConfig()
|
||||
n0 = len(pipeline.stages)
|
||||
|
||||
ext_pipeline, ext_deploy = inject_forced_aligner_stage(
|
||||
pipeline, deploy, {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"}
|
||||
)
|
||||
|
||||
# A new pooling stage is appended; the original pipeline is unchanged.
|
||||
assert len(pipeline.stages) == n0
|
||||
assert len(ext_pipeline.stages) == n0 + 1
|
||||
# The original deploy config is not mutated either (double-injection safe).
|
||||
assert len(deploy.stages) == 0
|
||||
assert len(ext_deploy.stages) == 1
|
||||
aligner = ext_pipeline.stages[-1]
|
||||
# Pooling stage reuses LLM_AR (+ runner=pooling), no dedicated LLM_POOLING.
|
||||
assert aligner.execution_type == StageExecutionType.LLM_AR
|
||||
assert aligner.final_output is True
|
||||
assert aligner.final_output_type == TIMESTAMPS_MODALITY
|
||||
assert aligner.input_sources == (n0 - 1,)
|
||||
assert aligner.requires_multimodal_data is True
|
||||
assert aligner.owns_tokenizer is True
|
||||
|
||||
# The deploy stage carries the separate aligner checkpoint + decoder hook.
|
||||
ds = ext_deploy.stages[-1]
|
||||
assert ds.engine_extras["model"] == "/models/Qwen3-ForcedAligner-0.6B"
|
||||
assert ds.engine_extras["runner"] == "pooling"
|
||||
assert ds.engine_extras["pooling_output_decoder"] == POOLING_OUTPUT_DECODER_PATH
|
||||
|
||||
|
||||
def test_inject_from_config_only():
|
||||
# --forced-aligner-config alone (model in the YAML) must also inject.
|
||||
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
|
||||
deploy = DeployConfig()
|
||||
n0 = len(pipeline.stages)
|
||||
|
||||
ext_pipeline, ext_deploy = inject_forced_aligner_stage(
|
||||
pipeline, deploy, {"forced_aligner": None, "forced_aligner_config": str(_DEFAULT_CONFIG_PATH)}
|
||||
)
|
||||
|
||||
assert len(ext_pipeline.stages) == n0 + 1
|
||||
assert ext_pipeline.stages[-1].final_output_type == TIMESTAMPS_MODALITY
|
||||
assert ext_deploy.stages[-1].engine_extras["model"] == "Qwen/Qwen3-ForcedAligner-0.6B"
|
||||
|
||||
|
||||
def test_inject_noop_without_forced_aligner():
|
||||
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
|
||||
deploy = DeployConfig()
|
||||
|
||||
ext_pipeline, ext_deploy = inject_forced_aligner_stage(pipeline, deploy, {})
|
||||
|
||||
assert ext_pipeline is pipeline
|
||||
assert len(ext_deploy.stages) == 0
|
||||
@@ -471,6 +471,7 @@ def test_sub_config_fields_match_structured_scopes():
|
||||
"trust_remote_code",
|
||||
"dtype",
|
||||
"attention_backend",
|
||||
"attention_config",
|
||||
"moe_backend",
|
||||
"hf_overrides",
|
||||
"limit_mm_per_prompt",
|
||||
@@ -513,6 +514,7 @@ def test_sub_config_fields_match_structured_scopes():
|
||||
"enable_prefix_caching",
|
||||
"disable_hybrid_kv_cache_manager",
|
||||
"mm_processor_cache_gb",
|
||||
"mamba_ssm_cache_dtype",
|
||||
}
|
||||
assert {f.name for f in fields(OmniStageSchedulerConfig)} == {
|
||||
"max_num_seqs",
|
||||
|
||||
@@ -27,6 +27,7 @@ _MIXIN_UPDATE_HELPERS = (
|
||||
"_attach_finished_request_sets",
|
||||
"_attach_scheduler_stats",
|
||||
"_capture_omni_connector_output",
|
||||
"_maybe_decode_pooling_output",
|
||||
)
|
||||
|
||||
|
||||
@@ -138,6 +139,7 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
|
||||
pending_stop_after_extraction=set(),
|
||||
waiting_for_transfer_free=set(),
|
||||
_new_prompt_len_snapshot={},
|
||||
_pooling_output_decoder=None,
|
||||
finished_req_ids=set(),
|
||||
finished_req_ids_dict=defaultdict(set),
|
||||
kv_cache_manager=SimpleNamespace(take_events=lambda: None),
|
||||
@@ -267,3 +269,72 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
|
||||
assert output_by_id["bad"].new_token_ids == []
|
||||
assert output_by_id["good"].new_token_ids == [8]
|
||||
np.testing.assert_array_equal(output_by_id["good"].new_logprobs.logprob_token_ids[:, 0], [8])
|
||||
|
||||
|
||||
def _pooling_model_runner_output(pooler_tensor) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
sampled_token_ids=[[]],
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={},
|
||||
pooler_output=[pooler_tensor],
|
||||
num_nans_in_logits=None,
|
||||
kv_connector_output=None,
|
||||
cudagraph_stats=None,
|
||||
req_id_to_index={"req": 0},
|
||||
routed_experts=None,
|
||||
)
|
||||
|
||||
|
||||
def _make_pooling_stub(decoder) -> tuple[_Request, SimpleNamespace]:
|
||||
request = _Request("req")
|
||||
request.sampling_params = SimpleNamespace(num_logprobs=None)
|
||||
request.pooling_params = SimpleNamespace()
|
||||
scheduler = _make_scheduler_stub([request])
|
||||
scheduler._pooling_output_decoder = decoder
|
||||
scheduler.vllm_config = SimpleNamespace(model_config=SimpleNamespace(hf_config=None))
|
||||
_bind_request_lifecycle(scheduler, update_request=lambda req, tokens: (tokens, False))
|
||||
return request, scheduler
|
||||
|
||||
|
||||
def test_pooling_decode_success_emits_decoded_payload() -> None:
|
||||
import torch
|
||||
|
||||
decoded = torch.tensor([[1, 2]], dtype=torch.int32)
|
||||
request, scheduler = _make_pooling_stub(lambda logits, req, hf: decoded)
|
||||
|
||||
scheduler_output = SimpleNamespace(
|
||||
num_scheduled_tokens={"req": 1},
|
||||
scheduled_spec_decode_tokens={},
|
||||
num_invalid_spec_tokens=0,
|
||||
)
|
||||
outputs = OmniARScheduler.update_from_output(
|
||||
scheduler, scheduler_output, _pooling_model_runner_output(torch.zeros(2, 4))
|
||||
)
|
||||
(output,) = outputs[0].outputs
|
||||
|
||||
assert output.finish_reason is FinishReason.STOP
|
||||
assert output.pooling_output is decoded
|
||||
|
||||
|
||||
def test_pooling_decode_failure_finishes_request_with_error() -> None:
|
||||
import torch
|
||||
|
||||
def raising_decoder(logits, req, hf):
|
||||
raise ValueError("boom")
|
||||
|
||||
request, scheduler = _make_pooling_stub(raising_decoder)
|
||||
|
||||
scheduler_output = SimpleNamespace(
|
||||
num_scheduled_tokens={"req": 1},
|
||||
scheduled_spec_decode_tokens={},
|
||||
num_invalid_spec_tokens=0,
|
||||
)
|
||||
outputs = OmniARScheduler.update_from_output(
|
||||
scheduler, scheduler_output, _pooling_model_runner_output(torch.zeros(2, 4))
|
||||
)
|
||||
(output,) = outputs[0].outputs
|
||||
|
||||
# Decoder failure is a request error (500), not an empty success.
|
||||
assert output.finish_reason is FinishReason.ERROR
|
||||
assert output.pooling_output is None
|
||||
assert "pooling output decode failed" in str(request.stop_reason)
|
||||
@@ -120,6 +120,9 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase):
|
||||
("MiMoAudioModel", "code2wav"),
|
||||
("Qwen3TTSCode2Wav", "code2wav"),
|
||||
("MiniCPMO45Code2Wav", "code2wav"),
|
||||
# nemotron_voicechat: sync-path RVQ-VAE decode waits for the
|
||||
# talker's full [frames, 31] code-stack payload.
|
||||
("NemotronVoiceChatCode2Wav", "code2wav"),
|
||||
("CosyVoice3Model", "cosyvoice3_code2wav"),
|
||||
# audex: sync-path TTS decoder (async_chunk: false deploys) and
|
||||
# the TTA XCodec1 stage (always sync full-payload).
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_minicpmo_4_5_ready",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100",
|
||||
"npu": "A3"
|
||||
},
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"core_model",
|
||||
"omni"
|
||||
],
|
||||
"server_params": {
|
||||
"model": "openbmb/MiniCPM-o-4_5",
|
||||
"extra_cli_args": [
|
||||
"--deploy-config",
|
||||
"vllm_omni/deploy/minicpmo_4_5.yaml",
|
||||
"--trust-remote-code"
|
||||
]
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"dataset_name": "seed-tts",
|
||||
"dataset_path": "zhaochenyang20/seed-tts-eval",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"num_prompts": [
|
||||
10,
|
||||
40
|
||||
],
|
||||
"max_concurrency": [
|
||||
1,
|
||||
4
|
||||
],
|
||||
"no_oversample": true,
|
||||
"trust_remote_code": true,
|
||||
"seed_tts_locale": "en",
|
||||
"extra_body": {
|
||||
"modalities": [
|
||||
"text",
|
||||
"audio"
|
||||
],
|
||||
"chat_template_kwargs": {
|
||||
"enable_thinking": false,
|
||||
"use_tts_template": true
|
||||
}
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"request_throughput": [
|
||||
0.7803,
|
||||
1.1008
|
||||
],
|
||||
"mean_ttft_ms": [
|
||||
166.9764,
|
||||
307.7133
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
1361.1682,
|
||||
3800.6332
|
||||
],
|
||||
"mean_audio_ttfp_ms": [
|
||||
635.8099,
|
||||
1615.6591
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.3151,
|
||||
0.8568
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -125,7 +125,7 @@ def _extract_images(output: OmniRequestOutput) -> list:
|
||||
return output.images
|
||||
# When the output comes from the orchestrator pipeline, images may be
|
||||
# nested inside request_output.
|
||||
inner = getattr(output, "request_output", None)
|
||||
inner = output
|
||||
if inner is not None and hasattr(inner, "images") and inner.images:
|
||||
return inner.images
|
||||
return []
|
||||
|
||||
+2
-2
@@ -56,10 +56,10 @@ def test_teacache(model_name: str):
|
||||
# Extract images from request_output['images']
|
||||
first_output = outputs[0]
|
||||
assert first_output.final_output_type == "image"
|
||||
if not hasattr(first_output, "request_output") or not first_output.request_output:
|
||||
if not isinstance(first_output, OmniRequestOutput) or not first_output:
|
||||
raise ValueError("No request_output found in OmniRequestOutput")
|
||||
|
||||
req_out = first_output.request_output
|
||||
req_out = first_output
|
||||
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
|
||||
raise ValueError("Invalid request_output structure or missing 'images' key")
|
||||
|
||||
|
||||
@@ -131,7 +131,7 @@ def _run_inference(
|
||||
elapsed_ms = (time.time() - start) * 1000
|
||||
|
||||
return InferenceResult(
|
||||
images=outputs[0].request_output.images,
|
||||
images=outputs[0].images,
|
||||
elapsed_ms=elapsed_ms,
|
||||
)
|
||||
finally:
|
||||
|
||||
@@ -49,10 +49,10 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
|
||||
def _extract_single_image(outputs) -> Image.Image:
|
||||
first_output = outputs[0]
|
||||
assert first_output.final_output_type == "image"
|
||||
if not hasattr(first_output, "request_output") or not first_output.request_output:
|
||||
if not isinstance(first_output, OmniRequestOutput) or not first_output:
|
||||
raise ValueError("No request_output found in OmniRequestOutput")
|
||||
|
||||
req_out = first_output.request_output
|
||||
req_out = first_output
|
||||
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
|
||||
raise ValueError("Invalid request_output structure or missing 'images' key")
|
||||
|
||||
|
||||
@@ -231,6 +231,128 @@ def test_reshard_from_trimmed_width_pads_invalid_columns(monkeypatch: pytest.Mon
|
||||
assert torch.equal(out.flatten(), torch.tensor([4.0, 0.0]))
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize("split_dim", ["height", "width"])
|
||||
@pytest.mark.parametrize(
|
||||
("full_extent", "world_size"),
|
||||
[
|
||||
pytest.param(5, 4, id="four-ranks-empty-tail"),
|
||||
pytest.param(20, 8, id="production-eight-ranks-empty-tail"),
|
||||
pytest.param(39, 2, id="two-ranks-partial-tail"),
|
||||
],
|
||||
)
|
||||
def test_reshard_from_trimmed_extent_clamps_and_reconstructs(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
split_dim: str,
|
||||
full_extent: int,
|
||||
world_size: int,
|
||||
):
|
||||
local_extent = (full_extent + world_size - 1) // world_size
|
||||
shape = (1, 1, 1, full_extent, 1) if split_dim == "height" else (1, 1, 1, 1, full_extent)
|
||||
x = torch.arange(full_extent, dtype=torch.float32).reshape(shape)
|
||||
outputs = []
|
||||
|
||||
for rank in range(world_size):
|
||||
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group, rank=rank: (rank, world_size))
|
||||
token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
|
||||
wan_spatial_shard.SpatialShardContext(
|
||||
input_extent=full_extent,
|
||||
local_input_extent=local_extent,
|
||||
split_dim=split_dim,
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
)
|
||||
)
|
||||
try:
|
||||
out = wan_spatial_shard.reshard_from_trimmed_extent(
|
||||
x,
|
||||
local_extent=local_extent,
|
||||
split_dim=split_dim,
|
||||
group=object(),
|
||||
)
|
||||
finally:
|
||||
wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
|
||||
|
||||
start = rank * local_extent
|
||||
valid_extent = min(local_extent, max(0, full_extent - start))
|
||||
expected = torch.zeros(local_extent, dtype=torch.float32)
|
||||
if valid_extent:
|
||||
expected[:valid_extent] = torch.arange(start, start + valid_extent, dtype=torch.float32)
|
||||
assert out.shape[_spatial_dim_for_test(split_dim)] == local_extent
|
||||
assert torch.equal(out.flatten(), expected)
|
||||
outputs.append(out)
|
||||
|
||||
reconstructed = torch.cat(outputs, dim=_spatial_dim_for_test(split_dim))
|
||||
reconstructed = reconstructed.narrow(_spatial_dim_for_test(split_dim), 0, full_extent)
|
||||
assert torch.equal(reconstructed, x)
|
||||
|
||||
|
||||
def _spatial_dim_for_test(split_dim: str) -> int:
|
||||
return -2 if split_dim == "height" else -1
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_reshard_from_trimmed_extent_rejects_rank_mismatch(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 2))
|
||||
x = torch.zeros((1, 1, 1, 1, 4), dtype=torch.float32)
|
||||
token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
|
||||
wan_spatial_shard.SpatialShardContext(
|
||||
input_extent=4,
|
||||
local_input_extent=2,
|
||||
split_dim="width",
|
||||
rank=0,
|
||||
world_size=2,
|
||||
)
|
||||
)
|
||||
try:
|
||||
with pytest.raises(RuntimeError, match="group_rank=1, context_rank=0"):
|
||||
wan_spatial_shard.reshard_from_trimmed_extent(
|
||||
x,
|
||||
local_extent=2,
|
||||
split_dim="width",
|
||||
group=object(),
|
||||
)
|
||||
finally:
|
||||
wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_attention_wrapper_rejects_spatial_extent_change(monkeypatch: pytest.MonkeyPatch):
|
||||
class ShrinkingWanAttentionBlock(torch.nn.Module):
|
||||
def forward(self, x):
|
||||
return x[..., :-1]
|
||||
|
||||
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 2))
|
||||
monkeypatch.setattr(
|
||||
wan_spatial_shard,
|
||||
"all_gather_along_dim",
|
||||
lambda x, *, group, dim, dst=None: torch.cat([x, x], dim=dim),
|
||||
)
|
||||
module = ShrinkingWanAttentionBlock()
|
||||
wan_spatial_shard._patch_attention_block(module, group=object(), split_dim="width")
|
||||
x = torch.zeros((1, 1, 1, 1, 2), dtype=torch.float32)
|
||||
token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
|
||||
wan_spatial_shard.SpatialShardContext(
|
||||
input_extent=4,
|
||||
local_input_extent=2,
|
||||
split_dim="width",
|
||||
rank=0,
|
||||
world_size=2,
|
||||
)
|
||||
)
|
||||
try:
|
||||
with pytest.raises(
|
||||
RuntimeError,
|
||||
match=r"local_extent=2, gathered_extent=4, trimmed_extent=4, output_extent=3",
|
||||
):
|
||||
module(x)
|
||||
finally:
|
||||
wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_halo_exchange_single_rank_noop(monkeypatch: pytest.MonkeyPatch):
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion]
|
||||
|
||||
_HEAD_DIM = 128
|
||||
_ROTARY_DIM = 96
|
||||
_EPS = 1e-5
|
||||
|
||||
|
||||
def _reference(q, k, q_weight, k_weight, rope_table):
|
||||
q = F.rms_norm(q, (_HEAD_DIM,), q_weight, _EPS)
|
||||
k = F.rms_norm(k, (_HEAD_DIM,), k_weight, _EPS)
|
||||
half = _ROTARY_DIM // 2
|
||||
cos = rope_table[..., :half].unsqueeze(1)
|
||||
sin = rope_table[..., half:].unsqueeze(1)
|
||||
|
||||
def apply(x):
|
||||
first = x[..., :half]
|
||||
second = x[..., half:_ROTARY_DIM]
|
||||
return torch.cat(
|
||||
(
|
||||
first * cos - second * sin,
|
||||
second * cos + first * sin,
|
||||
x[..., _ROTARY_DIM:],
|
||||
),
|
||||
dim=-1,
|
||||
)
|
||||
|
||||
return apply(q), apply(k)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
|
||||
@pytest.mark.parametrize("seq_len", [1, 257, 1024])
|
||||
def test_fused_qk_norm_rope_matches_bf16_reference(seq_len):
|
||||
from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
|
||||
fused_qk_norm_rope,
|
||||
)
|
||||
|
||||
torch.manual_seed(17)
|
||||
heads = 14
|
||||
qkv = torch.randn(
|
||||
seq_len,
|
||||
heads * _HEAD_DIM * 3,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
q = qkv[:, : heads * _HEAD_DIM].view(seq_len, heads, _HEAD_DIM)
|
||||
k = qkv[:, heads * _HEAD_DIM : 2 * heads * _HEAD_DIM].view(
|
||||
seq_len,
|
||||
heads,
|
||||
_HEAD_DIM,
|
||||
)
|
||||
q_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
|
||||
k_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
|
||||
freqs = torch.randn(seq_len, _ROTARY_DIM // 2, device="cuda")
|
||||
rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16)
|
||||
|
||||
expected_q, expected_k = _reference(
|
||||
q,
|
||||
k,
|
||||
q_weight,
|
||||
k_weight,
|
||||
rope_table,
|
||||
)
|
||||
actual_q, actual_k = fused_qk_norm_rope(
|
||||
q,
|
||||
k,
|
||||
q_weight,
|
||||
k_weight,
|
||||
rope_table,
|
||||
_EPS,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
|
||||
torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)
|
||||
@@ -29,10 +29,10 @@ def test_diffusion_model(tmp_path: Path, omni_runner):
|
||||
raise ValueError("Empty outputs from Omni.generate()")
|
||||
first_output = outputs[0]
|
||||
assert first_output.final_output_type == "image"
|
||||
if not hasattr(first_output, "request_output") or not first_output.request_output:
|
||||
if not isinstance(first_output, OmniRequestOutput) or not first_output:
|
||||
raise ValueError("No request_output found in OmniRequestOutput")
|
||||
|
||||
req_out = first_output.request_output
|
||||
req_out = first_output
|
||||
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
|
||||
raise ValueError("Invalid request_output structure or missing 'images' key")
|
||||
return req_out.images
|
||||
|
||||
@@ -23,7 +23,7 @@ def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1):
|
||||
return pipeline
|
||||
|
||||
|
||||
def test_prepare_video_latents_samples_directly_in_packed_token_space():
|
||||
def test_prepare_video_latents_matches_official_values_and_token_major_layout():
|
||||
pipeline = _make_pipeline(LTX2Pipeline)
|
||||
pipeline.vae_spatial_compression_ratio = 8
|
||||
pipeline.vae_temporal_compression_ratio = 8
|
||||
@@ -44,6 +44,7 @@ def test_prepare_video_latents_samples_directly_in_packed_token_space():
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
assert actual.stride()[1:] == (1, actual.shape[1])
|
||||
|
||||
|
||||
def test_prepare_audio_latents_samples_directly_in_packed_token_space():
|
||||
@@ -81,6 +82,42 @@ def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
|
||||
assert original_num_frames == 1
|
||||
assert padded_num_frames == 2
|
||||
assert latents.shape == (1, 2, 128)
|
||||
torch.testing.assert_close(latents[:, 1:], torch.zeros_like(latents[:, 1:]))
|
||||
|
||||
|
||||
def test_prepare_audio_latents_request_rng_is_invariant_to_sp_padding():
|
||||
pipeline_sp1 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=1)
|
||||
pipeline_sp4 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
|
||||
generator_sp1 = torch.Generator().manual_seed(42)
|
||||
generator_sp4 = torch.Generator().manual_seed(42)
|
||||
|
||||
latents_sp1, _, _ = pipeline_sp1.prepare_audio_latents(
|
||||
batch_size=1,
|
||||
num_channels_latents=2,
|
||||
num_mel_bins=8,
|
||||
audio_latent_length=3,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
generator=generator_sp1,
|
||||
)
|
||||
latents_sp4, _, _ = pipeline_sp4.prepare_audio_latents(
|
||||
batch_size=1,
|
||||
num_channels_latents=2,
|
||||
num_mel_bins=8,
|
||||
audio_latent_length=3,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
generator=generator_sp4,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(latents_sp4[:, :3], latents_sp1, rtol=0, atol=0)
|
||||
torch.testing.assert_close(latents_sp4[:, 3:], torch.zeros_like(latents_sp4[:, 3:]))
|
||||
torch.testing.assert_close(
|
||||
torch.randn(8, generator=generator_sp4),
|
||||
torch.randn(8, generator=generator_sp1),
|
||||
rtol=0,
|
||||
atol=0,
|
||||
)
|
||||
|
||||
|
||||
def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
|
||||
@@ -137,7 +174,9 @@ def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
|
||||
assert original_num_frames == 10
|
||||
assert padded_num_frames == 12
|
||||
assert padded.shape == (1, 12, 8)
|
||||
torch.testing.assert_close(padded, latent_ops.pack_audio_latents(latents))
|
||||
packed = latent_ops.pack_audio_latents(latents)
|
||||
torch.testing.assert_close(padded[:, :10], packed[:, :10])
|
||||
torch.testing.assert_close(padded[:, 10:], torch.zeros_like(padded[:, 10:]))
|
||||
|
||||
|
||||
def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
|
||||
|
||||
@@ -12,6 +12,19 @@ import torch
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker):
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_guidance import velocity_from_x0
|
||||
|
||||
sigma = mocker.MagicMock()
|
||||
sigma.to.return_value.item.return_value = 0.5
|
||||
|
||||
actual = velocity_from_x0(torch.tensor([1.0]), torch.tensor([0.0]), sigma)
|
||||
|
||||
sigma.to.assert_called_once_with(torch.float32)
|
||||
sigma.to.return_value.item.assert_called_once_with()
|
||||
torch.testing.assert_close(actual, torch.tensor([2.0]))
|
||||
|
||||
|
||||
class TestCFGParallelHelpers:
|
||||
"""Test LTX-2.3 CFG helper math without loading model weights."""
|
||||
|
||||
@@ -199,6 +212,7 @@ class TestCFGParallelHelpers:
|
||||
),
|
||||
attention_kwargs=None,
|
||||
audio_scheduler=SimpleNamespace(sigmas=torch.stack([audio_sigma])),
|
||||
original_audio_num_frames=state.audio.shape[1],
|
||||
)
|
||||
|
||||
actual_video, actual_audio = executor.predict_parallel_guidance(
|
||||
|
||||
@@ -0,0 +1,483 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Unit tests for the fixed LTX refinement-phase adapter."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from safetensors.torch import save_file
|
||||
|
||||
from vllm_omni.diffusion.models.ltx2 import ltx2_phase_adapter
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_adapter_parser import (
|
||||
AdapterTarget,
|
||||
parse_ltx_adapter,
|
||||
)
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_components import LTX23_TWO_STAGE_COMPONENT_PROFILE
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_phase_adapter import LTXPhaseAdapterRuntime
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
|
||||
LTX2_TWO_STAGE_RECIPE,
|
||||
LTX_DISTILLED_ADAPTER_SLOT,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _target(
|
||||
lora_a,
|
||||
lora_b,
|
||||
*,
|
||||
source_module="proj",
|
||||
module="proj",
|
||||
target_slice=None,
|
||||
):
|
||||
return AdapterTarget(
|
||||
source_module=source_module,
|
||||
module=module,
|
||||
slice=target_slice,
|
||||
a_key="a",
|
||||
b_key="b",
|
||||
a_shape=tuple(lora_a.shape),
|
||||
b_shape=tuple(lora_b.shape),
|
||||
)
|
||||
|
||||
|
||||
def _wrap(
|
||||
layer,
|
||||
lora_a,
|
||||
lora_b,
|
||||
*,
|
||||
source_module="proj",
|
||||
module="proj",
|
||||
target_slice=None,
|
||||
layer_fused=False,
|
||||
):
|
||||
target = _target(
|
||||
lora_a,
|
||||
lora_b,
|
||||
source_module=source_module,
|
||||
module=module,
|
||||
target_slice=target_slice,
|
||||
)
|
||||
wrapper = ltx2_phase_adapter._PhaseAdapterLinear(
|
||||
layer,
|
||||
[target],
|
||||
adapter_dtype=lora_a.dtype,
|
||||
layer_fused=layer_fused,
|
||||
)
|
||||
wrapper.load_target(target, lora_a, lora_b)
|
||||
wrapper.set_enabled(True)
|
||||
return wrapper
|
||||
|
||||
|
||||
class _LoRAAttention(torch.nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.to_qkv = torch.nn.Identity()
|
||||
|
||||
|
||||
class _LoRATransformer(torch.nn.Module):
|
||||
stacked_params_mapping = ((".attn1.to_qkv", ".attn1.to_q", "q"),)
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.proj_in = torch.nn.Linear(4, 6, bias=False)
|
||||
block = torch.nn.Module()
|
||||
block.attn1 = _LoRAAttention()
|
||||
self.transformer_blocks = torch.nn.ModuleList([block])
|
||||
|
||||
|
||||
def test_ltx_adapter_parser_maps_official_and_packed_module_names(tmp_path):
|
||||
adapter_path = tmp_path / "distilled.safetensors"
|
||||
save_file(
|
||||
{
|
||||
"diffusion_model.text_embedding_projection.aggregate_embed.lora_A.weight": torch.randn(2, 4),
|
||||
"diffusion_model.patchify_proj.lora_A.weight": torch.randn(2, 4),
|
||||
"diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2),
|
||||
"diffusion_model.transformer_blocks.0.attn1.to_q.lora_A.weight": torch.randn(2, 4),
|
||||
"diffusion_model.transformer_blocks.0.attn1.to_q.lora_B.weight": torch.randn(4, 2),
|
||||
},
|
||||
adapter_path,
|
||||
)
|
||||
|
||||
manifest = parse_ltx_adapter(_LoRATransformer(), str(adapter_path))
|
||||
assert {(target.source_module, target.module, target.slice) for target in manifest.targets} == {
|
||||
("proj_in", "proj_in", None),
|
||||
("transformer_blocks.0.attn1.to_q", "transformer_blocks.0.attn1.to_qkv", "q"),
|
||||
}
|
||||
|
||||
|
||||
def test_ltx_adapter_parser_rejects_unpaired_b_tensor(tmp_path):
|
||||
adapter_path = tmp_path / "distilled.safetensors"
|
||||
save_file(
|
||||
{"diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2)},
|
||||
adapter_path,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Missing paired LoRA A tensors"):
|
||||
parse_ltx_adapter(_LoRATransformer(), str(adapter_path))
|
||||
|
||||
|
||||
def _pipeline(*, dtype, quantization_config):
|
||||
return SimpleNamespace(
|
||||
pipeline_recipe=LTX2_TWO_STAGE_RECIPE,
|
||||
component_profile=LTX23_TWO_STAGE_COMPONENT_PROFILE,
|
||||
transformer=object(),
|
||||
od_config=SimpleNamespace(
|
||||
model="unused",
|
||||
lora_path=None,
|
||||
dtype=dtype,
|
||||
quantization_config=quantization_config,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("dtype", "quantization_config", "layer_fused"),
|
||||
[
|
||||
(torch.float32, object(), False),
|
||||
(torch.bfloat16, None, True),
|
||||
],
|
||||
)
|
||||
def test_ltx_phase_adapter_selects_execution_from_quantization(
|
||||
monkeypatch,
|
||||
dtype,
|
||||
quantization_config,
|
||||
layer_fused,
|
||||
):
|
||||
pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config)
|
||||
manifest = object()
|
||||
captured = {}
|
||||
|
||||
def resolve_artifact(*_args, **kwargs):
|
||||
assert not kwargs
|
||||
return "adapter"
|
||||
|
||||
def parse(transformer, path):
|
||||
assert transformer is pipeline.transformer
|
||||
assert path == "adapter"
|
||||
return manifest
|
||||
|
||||
class Runtime:
|
||||
def __init__(self, transformer, parsed_manifest, **kwargs):
|
||||
captured.update(transformer=transformer, manifest=parsed_manifest, **kwargs)
|
||||
|
||||
monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", resolve_artifact)
|
||||
monkeypatch.setattr(ltx2_phase_adapter, "parse_ltx_adapter", parse)
|
||||
monkeypatch.setattr(ltx2_phase_adapter, "LTXPhaseAdapterRuntime", Runtime)
|
||||
|
||||
runtime = ltx2_phase_adapter.build_ltx_phase_adapter(pipeline)
|
||||
|
||||
assert isinstance(runtime, Runtime)
|
||||
assert captured == {
|
||||
"transformer": pipeline.transformer,
|
||||
"manifest": manifest,
|
||||
"dtype": dtype,
|
||||
"layer_fused": layer_fused,
|
||||
}
|
||||
|
||||
|
||||
def test_ltx_phase_adapter_rejects_non_bfloat16_layer_fused(monkeypatch):
|
||||
pipeline = _pipeline(dtype=torch.float16, quantization_config=None)
|
||||
monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", lambda *_args, **_kwargs: "adapter")
|
||||
|
||||
with pytest.raises(ValueError, match="bfloat16"):
|
||||
ltx2_phase_adapter.build_ltx_phase_adapter(pipeline)
|
||||
|
||||
|
||||
def test_ltx_phase_adapter_keeps_one_transformer_and_switches_a_fixed_slot(tmp_path):
|
||||
class TinyTransformer(torch.nn.Module):
|
||||
stacked_params_mapping = ()
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.proj = torch.nn.Linear(2, 2, bias=False)
|
||||
|
||||
adapter_path = tmp_path / "distilled.safetensors"
|
||||
save_file(
|
||||
{
|
||||
"diffusion_model.proj.lora_A.weight": torch.tensor([[1.0, 2.0]]),
|
||||
"diffusion_model.proj.lora_B.weight": torch.tensor([[3.0], [4.0]]),
|
||||
},
|
||||
adapter_path,
|
||||
)
|
||||
transformer = TinyTransformer()
|
||||
runtime = LTXPhaseAdapterRuntime(
|
||||
transformer,
|
||||
parse_ltx_adapter(transformer, str(adapter_path)),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
assert runtime.base_parameter_name("proj.weight") == "proj.base_layer.weight"
|
||||
assert {name for name, _ in transformer.named_parameters()} == {"proj.base_layer.weight"}
|
||||
with torch.no_grad():
|
||||
transformer.proj.base_layer.weight.copy_(torch.eye(2))
|
||||
runtime.finalize()
|
||||
|
||||
x = torch.tensor([[1.0, 1.0]])
|
||||
runtime.activate(None)
|
||||
torch.testing.assert_close(transformer.proj(x), torch.tensor([[1.0, 1.0]]))
|
||||
runtime.activate(LTX_DISTILLED_ADAPTER_SLOT)
|
||||
torch.testing.assert_close(transformer.proj(x), torch.tensor([[10.0, 13.0]]))
|
||||
|
||||
|
||||
def test_ltx_layer_fused_adapter_matches_official_bf16_weight_fusion():
|
||||
layer = torch.nn.Linear(4, 3, bias=True, dtype=torch.bfloat16)
|
||||
base_weight = torch.tensor(
|
||||
[
|
||||
[0.5, -0.25, 0.125, 0.75],
|
||||
[-0.375, 0.625, 0.25, -0.5],
|
||||
[0.875, -0.125, -0.75, 0.375],
|
||||
],
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
lora_a = torch.tensor(
|
||||
[
|
||||
[0.1015625, -0.203125, 0.3046875, -0.40625],
|
||||
[0.5078125, -0.609375, 0.7109375, -0.8125],
|
||||
],
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
lora_b = torch.tensor(
|
||||
[[0.9140625, -0.1171875], [-0.21875, 0.3203125], [0.421875, -0.5234375]],
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
with torch.no_grad():
|
||||
layer.weight.copy_(base_weight)
|
||||
layer.bias.copy_(torch.tensor([0.125, -0.25, 0.5], dtype=torch.bfloat16))
|
||||
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
|
||||
|
||||
expected_weight = torch.matmul(lora_b, lora_a)
|
||||
expected_weight.add_(base_weight)
|
||||
x = torch.tensor([[0.6015625, -0.703125, 0.8046875, -0.90625]], dtype=torch.bfloat16)
|
||||
expected = torch.nn.functional.linear(x, expected_weight, layer.bias)
|
||||
|
||||
assert torch.equal(wrapper(x), expected)
|
||||
assert torch.equal(layer.weight, base_weight)
|
||||
|
||||
|
||||
def test_ltx_layer_fused_row_parallel_uses_one_collective(monkeypatch):
|
||||
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
||||
|
||||
class FakeRow(torch.nn.Module):
|
||||
input_size = 4
|
||||
output_size = 2
|
||||
input_size_per_partition = 2
|
||||
input_is_parallel = True
|
||||
tp_rank = 0
|
||||
tp_size = 2
|
||||
reduce_results = True
|
||||
skip_bias_add = False
|
||||
return_bias = False
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.weight = torch.nn.Parameter(torch.eye(2, dtype=torch.bfloat16))
|
||||
self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16))
|
||||
self.quant_method = UnquantizedLinearMethod()
|
||||
|
||||
monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", FakeRow)
|
||||
reductions = []
|
||||
|
||||
def all_reduce(value):
|
||||
reductions.append(value.clone())
|
||||
return value + 1
|
||||
|
||||
monkeypatch.setattr(ltx2_phase_adapter, "tensor_model_parallel_all_reduce", all_reduce)
|
||||
layer = FakeRow()
|
||||
lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16)
|
||||
lora_b = torch.tensor([[0.25], [-0.5]], dtype=torch.bfloat16)
|
||||
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
|
||||
|
||||
expected_weight = torch.matmul(lora_b, lora_a[:, :2])
|
||||
expected_weight.add_(layer.weight)
|
||||
x = torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16)
|
||||
expected_local = torch.nn.functional.linear(x, expected_weight, layer.bias)
|
||||
|
||||
actual = wrapper(x)
|
||||
assert len(reductions) == 1
|
||||
assert torch.equal(reductions[0], expected_local)
|
||||
assert torch.equal(actual, expected_local + 1)
|
||||
|
||||
|
||||
def test_ltx_layer_fused_column_parallel_uses_local_weight_shard(monkeypatch):
|
||||
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
||||
|
||||
class FakeColumn(torch.nn.Module):
|
||||
input_size = 3
|
||||
output_size = 4
|
||||
output_size_per_partition = 2
|
||||
tp_rank = 1
|
||||
tp_size = 2
|
||||
gather_output = False
|
||||
skip_bias_add = False
|
||||
return_bias = False
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.weight = torch.nn.Parameter(
|
||||
torch.tensor([[0.5, -0.25, 0.125], [-0.375, 0.625, 0.25]], dtype=torch.bfloat16)
|
||||
)
|
||||
self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16))
|
||||
self.quant_method = UnquantizedLinearMethod()
|
||||
|
||||
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn)
|
||||
layer = FakeColumn()
|
||||
base_weight = layer.weight.detach().clone()
|
||||
lora_a = torch.tensor([[0.5, -0.25, 0.75]], dtype=torch.bfloat16)
|
||||
lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16)
|
||||
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
|
||||
|
||||
expected_weight = torch.matmul(lora_b[2:], lora_a)
|
||||
expected_weight.add_(base_weight)
|
||||
x = torch.tensor([[1.0, -0.5, 0.25]], dtype=torch.bfloat16)
|
||||
expected = torch.nn.functional.linear(x, expected_weight, layer.bias)
|
||||
|
||||
assert torch.equal(wrapper(x), expected)
|
||||
assert torch.equal(layer.weight, base_weight)
|
||||
|
||||
|
||||
def test_ltx_layer_fused_packed_qkv_updates_only_the_target_slice(monkeypatch):
|
||||
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
||||
|
||||
class FakeColumn(torch.nn.Module):
|
||||
pass
|
||||
|
||||
class FakeQKV(FakeColumn):
|
||||
input_size = 4
|
||||
head_size = 2
|
||||
v_head_size = 2
|
||||
total_num_heads = 2
|
||||
total_num_kv_heads = 2
|
||||
num_heads = 1
|
||||
num_kv_heads = 1
|
||||
num_kv_head_replicas = 1
|
||||
tp_rank = 1
|
||||
tp_size = 2
|
||||
gather_output = False
|
||||
skip_bias_add = True
|
||||
return_bias = True
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.weight = torch.nn.Parameter(
|
||||
torch.arange(24, dtype=torch.float32).reshape(6, 4).to(torch.bfloat16) / 16
|
||||
)
|
||||
self.bias = torch.nn.Parameter(torch.arange(6, dtype=torch.float32).to(torch.bfloat16) / 8)
|
||||
self.quant_method = UnquantizedLinearMethod()
|
||||
|
||||
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn)
|
||||
monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", FakeQKV)
|
||||
layer = FakeQKV()
|
||||
base_weight = layer.weight.detach().clone()
|
||||
lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16)
|
||||
lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16)
|
||||
wrapper = _wrap(
|
||||
layer,
|
||||
lora_a,
|
||||
lora_b,
|
||||
source_module="to_k",
|
||||
module="to_qkv",
|
||||
target_slice="k",
|
||||
layer_fused=True,
|
||||
)
|
||||
|
||||
expected_weight = base_weight.clone()
|
||||
expected_k = torch.matmul(lora_b[2:], lora_a)
|
||||
expected_k.add_(base_weight[2:4])
|
||||
expected_weight[2:4].copy_(expected_k)
|
||||
x = torch.tensor([[1.0, -0.5, 0.25, -0.125]], dtype=torch.bfloat16)
|
||||
expected_output = torch.nn.functional.linear(x, expected_weight)
|
||||
|
||||
actual_output, returned_bias = wrapper(x)
|
||||
assert torch.equal(actual_output, expected_output)
|
||||
assert returned_bias is layer.bias
|
||||
assert torch.equal(layer.weight, base_weight)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("layer_factory", "target_slice", "expected_a", "expected_b", "expected_output"),
|
||||
[
|
||||
(
|
||||
lambda: type(
|
||||
"FakeColumn",
|
||||
(torch.nn.Module,),
|
||||
{
|
||||
"input_size": 4,
|
||||
"output_size": 4,
|
||||
"output_size_per_partition": 2,
|
||||
"tp_rank": 1,
|
||||
"gather_output": False,
|
||||
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)),
|
||||
},
|
||||
)(),
|
||||
None,
|
||||
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
|
||||
torch.tensor([[1.0], [2.0], [3.0], [4.0]]),
|
||||
torch.tensor([[12963.0, 17284.0]]),
|
||||
),
|
||||
(
|
||||
lambda: type(
|
||||
"FakeRow",
|
||||
(torch.nn.Module,),
|
||||
{
|
||||
"input_size": 4,
|
||||
"output_size": 2,
|
||||
"input_size_per_partition": 2,
|
||||
"tp_rank": 1,
|
||||
"reduce_results": False,
|
||||
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)),
|
||||
},
|
||||
)(),
|
||||
None,
|
||||
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
|
||||
torch.tensor([[2.0], [3.0]]),
|
||||
torch.tensor([[8600.0, 12900.0]]),
|
||||
),
|
||||
(
|
||||
lambda: type(
|
||||
"FakeQKV",
|
||||
(torch.nn.Module,),
|
||||
{
|
||||
"input_size": 4,
|
||||
"head_size": 2,
|
||||
"v_head_size": 2,
|
||||
"total_num_heads": 2,
|
||||
"total_num_kv_heads": 2,
|
||||
"num_heads": 1,
|
||||
"num_kv_heads": 1,
|
||||
"num_kv_head_replicas": 1,
|
||||
"tp_rank": 1,
|
||||
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 6)),
|
||||
},
|
||||
)(),
|
||||
"k",
|
||||
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
|
||||
torch.tensor([[1.0], [2.0], [3.0], [4.0]]),
|
||||
torch.tensor([[0.0, 0.0, 12963.0, 17284.0, 0.0, 0.0]]),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_ltx_phase_adapter_uses_local_tp_lora_slices(
|
||||
monkeypatch,
|
||||
layer_factory,
|
||||
target_slice,
|
||||
expected_a,
|
||||
expected_b,
|
||||
expected_output,
|
||||
):
|
||||
layer = layer_factory()
|
||||
layer.register_buffer("device_anchor", torch.empty(0))
|
||||
layer_type = type(layer)
|
||||
if target_slice == "k":
|
||||
monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", layer_type)
|
||||
elif "Column" in layer_type.__name__:
|
||||
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", layer_type)
|
||||
else:
|
||||
monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", layer_type)
|
||||
|
||||
wrapper = _wrap(layer, expected_a, expected_b, target_slice=target_slice)
|
||||
torch.testing.assert_close(wrapper(torch.tensor([[1.0, 2.0, 3.0, 4.0]])), expected_output)
|
||||
expected_local_a_width = getattr(layer, "input_size_per_partition", layer.input_size)
|
||||
assert wrapper.pieces[0].lora_a.shape[-1] == expected_local_a_width
|
||||
@@ -13,13 +13,21 @@ import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
from vllm_omni.diffusion.models.ltx2 import ltx2_components
|
||||
from vllm_omni.diffusion.models.ltx2 import (
|
||||
ltx2_components,
|
||||
ltx2_transformer,
|
||||
)
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_components import (
|
||||
LTX2_COMPONENT_PROFILE,
|
||||
LTX2_DISTILLED_COMPONENT_PROFILE,
|
||||
LTX2_TWO_STAGE_COMPONENT_PROFILE,
|
||||
LTX23_COMPONENT_PROFILE,
|
||||
LTX23_DISTILLED_COMPONENT_PROFILE,
|
||||
LTX23_TWO_STAGE_COMPONENT_PROFILE,
|
||||
_install_connector_attention,
|
||||
detect_ltx_model_version,
|
||||
resolve_ltx_artifact,
|
||||
resolve_ltx_component_profile,
|
||||
)
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
|
||||
@@ -33,6 +41,7 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
|
||||
LTXGuidancePlan,
|
||||
LTXGuidanceSpec,
|
||||
LTXModalityGuidance,
|
||||
_repeat_batch,
|
||||
build_perturbation_kwargs,
|
||||
combine_guided_x0,
|
||||
)
|
||||
@@ -40,18 +49,29 @@ from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
|
||||
LTX2_DISTILLED_TWO_STAGE_RECIPE,
|
||||
LTX2_ONE_STAGE_RECIPE,
|
||||
LTX2_TWO_STAGE_RECIPE,
|
||||
LTX23_DISTILLED_TWO_STAGE_RECIPE,
|
||||
LTX23_ONE_STAGE_RECIPE,
|
||||
LTX23_TWO_STAGE_RECIPE,
|
||||
LTX_DEFAULT_NEGATIVE_PROMPT,
|
||||
LTX_POSITIVE_ONLY_RECIPE,
|
||||
resolve_ltx_pipeline_recipe,
|
||||
)
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_request import (
|
||||
LTXRequestInputs,
|
||||
validate_ltx_checkpoint,
|
||||
validate_pipeline_request,
|
||||
)
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs, validate_pipeline_request
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
|
||||
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
|
||||
LTX2I2VDMD2Pipeline,
|
||||
LTX2Pipeline,
|
||||
LTX2T2VDMD2Pipeline,
|
||||
)
|
||||
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline
|
||||
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
|
||||
LTX2DistilledPipeline,
|
||||
LTX2TwoStagePipeline,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -104,10 +124,22 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries():
|
||||
)
|
||||
|
||||
assert issubclass(LTX2Pipeline, LTXRuntime)
|
||||
assert issubclass(LTX2TwoStagePipeline, LTXRuntime)
|
||||
assert issubclass(LTX2DistilledPipeline, LTXRuntime)
|
||||
assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE
|
||||
assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE
|
||||
assert LTX2TwoStagePipeline.component_profile is LTX2_TWO_STAGE_COMPONENT_PROFILE
|
||||
assert LTX2TwoStagePipeline.pipeline_recipe is LTX2_TWO_STAGE_RECIPE
|
||||
assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE
|
||||
for pipeline_cls, profile in (
|
||||
(LTX2Pipeline, LTX2_COMPONENT_PROFILE),
|
||||
(LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE),
|
||||
(LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE),
|
||||
):
|
||||
assert pipeline_cls._dit_modules == list(profile.dit_modules)
|
||||
assert pipeline_cls._encoder_modules == list(profile.encoder_modules)
|
||||
assert pipeline_cls._vae_modules == list(profile.vae_modules)
|
||||
assert pipeline_cls._resident_modules == list(profile.resident_modules)
|
||||
assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
|
||||
assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
|
||||
|
||||
@@ -123,6 +155,72 @@ 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):
|
||||
filename = "ltx-sidecar.safetensors"
|
||||
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 / "model"), "Lightricks/test", filename) == "/cache/ltx-sidecar.safetensors"
|
||||
)
|
||||
assert calls == [{"repo_id": "Lightricks/test", "filename": filename}]
|
||||
|
||||
|
||||
def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
|
||||
filename = "ltx-sidecar.safetensors"
|
||||
expected = tmp_path / filename
|
||||
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)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("expected_kind", "scheduler_config"),
|
||||
[
|
||||
("regular", {"use_dynamic_shifting": True, "shift_terminal": 0.1}),
|
||||
("distilled", {"use_dynamic_shifting": False, "shift_terminal": None}),
|
||||
],
|
||||
)
|
||||
def test_ltx_checkpoint_validation_accepts_matching_scheduler_metadata(expected_kind, scheduler_config):
|
||||
validate_ltx_checkpoint(
|
||||
scheduler_config,
|
||||
expected_kind=expected_kind,
|
||||
pipeline_name="TestLTXPipeline",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("expected_kind", "scheduler_config", "error"),
|
||||
[
|
||||
("regular", {"use_dynamic_shifting": False, "shift_terminal": None}, "regular non-distilled"),
|
||||
("distilled", {"use_dynamic_shifting": True, "shift_terminal": 0.1}, "merged distilled"),
|
||||
],
|
||||
)
|
||||
def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expected_kind, scheduler_config, error):
|
||||
with pytest.raises(ValueError, match=error):
|
||||
validate_ltx_checkpoint(
|
||||
scheduler_config,
|
||||
expected_kind=expected_kind,
|
||||
pipeline_name="TestLTXPipeline",
|
||||
)
|
||||
|
||||
|
||||
def test_ltx_rejects_advanced_uaa_before_component_initialization():
|
||||
od_config = SimpleNamespace(
|
||||
model="unused",
|
||||
parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"):
|
||||
LTX2Pipeline(od_config=od_config)
|
||||
|
||||
|
||||
def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch):
|
||||
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
|
||||
|
||||
@@ -144,16 +242,104 @@ def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, m
|
||||
assert pipe.reports_stage_durations
|
||||
|
||||
|
||||
def test_ltx_one_stage_entry_exposes_batch_image_and_distributed_decode_contracts():
|
||||
def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_path, monkeypatch):
|
||||
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
|
||||
|
||||
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]}))
|
||||
initialized = False
|
||||
|
||||
def stub_components(*_args, **_kwargs):
|
||||
nonlocal initialized
|
||||
initialized = True
|
||||
|
||||
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
|
||||
|
||||
with pytest.raises(ValueError, match="only one-stage LTX recipes"):
|
||||
LTX2TwoStagePipeline(
|
||||
od_config=SimpleNamespace(
|
||||
model=str(tmp_path),
|
||||
cache_backend="cache_dit",
|
||||
)
|
||||
)
|
||||
|
||||
assert not initialized
|
||||
|
||||
|
||||
def test_ltx23_checkpoint_selects_version_specific_two_stage_profiles(tmp_path, monkeypatch):
|
||||
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
|
||||
|
||||
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]}))
|
||||
|
||||
def stub_components(pipe, od_config):
|
||||
pipe.od_config = od_config
|
||||
pipe.vae_spatial_compression_ratio = 32
|
||||
|
||||
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
|
||||
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
|
||||
monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda *_args: SimpleNamespace())
|
||||
pipe = LTX2TwoStagePipeline(
|
||||
od_config=SimpleNamespace(
|
||||
model=str(tmp_path),
|
||||
enable_diffusion_pipeline_profiler=False,
|
||||
lora_path=None,
|
||||
model_config={},
|
||||
)
|
||||
)
|
||||
|
||||
assert pipe.model_version == "2.3"
|
||||
assert pipe.component_profile is LTX23_TWO_STAGE_COMPONENT_PROFILE
|
||||
assert pipe.pipeline_recipe is LTX23_TWO_STAGE_RECIPE
|
||||
|
||||
|
||||
def test_ltx23_checkpoint_selects_full_distilled_profile_and_recipe(tmp_path, monkeypatch):
|
||||
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
|
||||
|
||||
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]}))
|
||||
|
||||
def stub_components(pipe, od_config):
|
||||
pipe.od_config = od_config
|
||||
pipe.vae_spatial_compression_ratio = 32
|
||||
|
||||
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
|
||||
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
|
||||
|
||||
pipe = LTX2DistilledPipeline(
|
||||
od_config=SimpleNamespace(
|
||||
model=str(tmp_path),
|
||||
enable_diffusion_pipeline_profiler=False,
|
||||
)
|
||||
)
|
||||
|
||||
assert pipe.model_version == "2.3"
|
||||
assert pipe.component_profile is LTX23_DISTILLED_COMPONENT_PROFILE
|
||||
assert pipe.pipeline_recipe is LTX23_DISTILLED_TWO_STAGE_RECIPE
|
||||
assert pipe._phase_adapter is None
|
||||
assert pipe.preserve_sp_padded_audio_duration
|
||||
assert pipe.reports_stage_durations
|
||||
|
||||
|
||||
def test_ltx_entries_expose_batch_image_and_distributed_decode_contracts():
|
||||
assert LTX2Pipeline.supports_request_batch
|
||||
assert LTX2Pipeline.support_image_input
|
||||
assert LTX2Pipeline.unified_text_image_entry
|
||||
assert LTX2Pipeline.distributed_video_decode
|
||||
assert not LTX2TwoStagePipeline.supports_request_batch
|
||||
assert LTX2TwoStagePipeline.support_image_input
|
||||
assert LTX2TwoStagePipeline.unified_text_image_entry
|
||||
assert not LTX2DistilledPipeline.supports_request_batch
|
||||
assert LTX2DistilledPipeline.support_image_input
|
||||
assert LTX2DistilledPipeline.unified_text_image_entry
|
||||
|
||||
|
||||
def test_ltx_two_stage_resolves_request_embedded_images():
|
||||
pipe = object.__new__(LTX2TwoStagePipeline)
|
||||
request_inputs = SimpleNamespace(latents=None)
|
||||
image = torch.zeros(3, 8, 8)
|
||||
req = SimpleNamespace(prompts=[{"prompt": "image conditioned", "multi_modal_data": {"image": image}}])
|
||||
|
||||
assert pipe._resolve_request_image(req, None, request_inputs) is image
|
||||
|
||||
|
||||
def test_ltx_dmd2_entries_retain_their_task_boundaries():
|
||||
assert not LTX2T2VDMD2Pipeline.support_image_input
|
||||
assert LTX2I2VDMD2Pipeline.support_image_input
|
||||
@@ -330,6 +516,15 @@ def test_ltx_guidance_uses_official_close_to_disabled_semantics():
|
||||
assert LTXGuidancePlan.build(LTXGuidanceSpec(video=guidance)).names == ("cond",)
|
||||
|
||||
|
||||
def test_ltx_positive_only_guidance_preserves_token_major_layout():
|
||||
tensor = torch.arange(24).view(1, 3, 8).transpose(1, 2).contiguous().transpose(1, 2)
|
||||
|
||||
repeated = _repeat_batch(tensor, 1)
|
||||
|
||||
assert repeated is tensor
|
||||
assert repeated.stride() == tensor.stride()
|
||||
|
||||
|
||||
def test_ltx_guidance_rescale_is_invariant_to_request_batching():
|
||||
cond = torch.tensor(
|
||||
[
|
||||
@@ -380,6 +575,32 @@ def test_ltx_guidance_rescale_handles_zero_variance_prediction():
|
||||
torch.testing.assert_close(actual, torch.zeros_like(cond))
|
||||
|
||||
|
||||
def test_ltx_audio_guidance_rescale_ignores_sp_padding_tokens():
|
||||
logical_cond = torch.tensor([[[1.0, 2.0], [3.0, 5.0]]])
|
||||
logical_uncond = torch.tensor([[[0.5, -1.0], [2.0, 1.0]]])
|
||||
guidance = LTXModalityGuidance(cfg_scale=3.0, rescale_scale=1.0)
|
||||
expected = combine_guided_x0(
|
||||
cond=logical_cond,
|
||||
uncond_text=logical_uncond,
|
||||
uncond_perturbed=0.0,
|
||||
uncond_modality=0.0,
|
||||
guidance=guidance,
|
||||
)
|
||||
padded_cond = torch.cat([logical_cond, torch.tensor([[[1000.0, -1000.0]]])], dim=1)
|
||||
padded_uncond = torch.cat([logical_uncond, torch.tensor([[[-500.0, 500.0]]])], dim=1)
|
||||
|
||||
actual = combine_guided_x0(
|
||||
cond=padded_cond,
|
||||
uncond_text=padded_uncond,
|
||||
uncond_perturbed=0.0,
|
||||
uncond_modality=0.0,
|
||||
guidance=guidance,
|
||||
rescale_token_count=2,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual[:, :2], expected)
|
||||
|
||||
|
||||
def test_ltx_guidance_builds_per_sample_stg_and_modality_masks():
|
||||
plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance)
|
||||
perturbations = build_perturbation_kwargs(plan, batch_size=2, reference=torch.ones(8, 1, 1))
|
||||
@@ -547,11 +768,19 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
|
||||
|
||||
def test_ltx2_distilled_pipeline_shares_recipe_runtime():
|
||||
assert issubclass(LTX2DistilledPipeline, LTXI2VConditioningMixin)
|
||||
assert issubclass(LTX2DistilledPipeline, LTX2Pipeline)
|
||||
assert not issubclass(LTX2DistilledPipeline, LTX2Pipeline)
|
||||
assert LTX2DistilledPipeline._run_recipe is LTXRuntime._run_recipe
|
||||
assert "_forward_request" not in LTX2DistilledPipeline.__dict__
|
||||
|
||||
|
||||
def test_ltx_two_stage_entry_is_a_thin_pipeline_variant():
|
||||
assert "_run_recipe" not in LTX2TwoStagePipeline.__dict__
|
||||
assert "_forward_request" not in LTX2TwoStagePipeline.__dict__
|
||||
assert "load_weights" not in LTX2TwoStagePipeline.__dict__
|
||||
assert issubclass(LTX2TwoStagePipeline, LTXI2VConditioningMixin)
|
||||
assert not issubclass(LTX2TwoStagePipeline, LTX2Pipeline)
|
||||
|
||||
|
||||
def test_ltx_variants_share_denoise_loop_and_i2v_conditioning():
|
||||
assert "_denoise_loop" not in LTX2Pipeline.__dict__
|
||||
assert issubclass(LTX2Pipeline, LTXI2VConditioningMixin)
|
||||
@@ -594,7 +823,8 @@ def test_denoise_executor_owns_progress_and_interrupt():
|
||||
torch.testing.assert_close(state.audio, torch.tensor(11.0))
|
||||
|
||||
|
||||
def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
|
||||
@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
|
||||
def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
|
||||
request_inputs = LTXRequestInputs(
|
||||
prompt="prompt",
|
||||
negative_prompt="",
|
||||
@@ -671,12 +901,15 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
|
||||
return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
|
||||
|
||||
prompt_context_sentinel = prompt_context
|
||||
pipeline = object.__new__(LTX2DistilledPipeline)
|
||||
pipeline = object.__new__(pipeline_cls)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.vae_spatial_compression_ratio = 32
|
||||
pipeline.vae_temporal_compression_ratio = 8
|
||||
pipeline.latent_upsampler = FakeUpsampler()
|
||||
activated_slots = []
|
||||
if pipeline_cls is LTX2TwoStagePipeline:
|
||||
pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append)
|
||||
object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
|
||||
object.__setattr__(pipeline, "run_phase", run_phase)
|
||||
object.__setattr__(pipeline, "decode_phase", decode_phase)
|
||||
@@ -686,14 +919,18 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
|
||||
|
||||
assert len(phase_calls) == 2
|
||||
assert phase_calls[1][2] is prompt_context_sentinel
|
||||
if pipeline_cls is LTX2TwoStagePipeline:
|
||||
assert activated_slots == [None, "ltx_distilled"]
|
||||
torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0))
|
||||
torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0))
|
||||
expected_audio = 2.0 if pipeline_cls is LTX2TwoStagePipeline else 4.0
|
||||
torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), expected_audio))
|
||||
|
||||
|
||||
def test_ltx2_distilled_stage2_reapplies_final_resolution_i2v_conditioning():
|
||||
@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
|
||||
def test_ltx_two_stage_stage2_reapplies_final_resolution_i2v_conditioning(pipeline_cls):
|
||||
from vllm_omni.diffusion.models.ltx2 import ltx2_latents
|
||||
|
||||
pipeline = object.__new__(LTX2DistilledPipeline)
|
||||
pipeline = object.__new__(pipeline_cls)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.vae_spatial_compression_ratio = 32
|
||||
pipeline.vae_temporal_compression_ratio = 8
|
||||
@@ -750,6 +987,18 @@ def test_ltx2_distilled_two_stage_recipe_is_fixed_positive_only():
|
||||
assert (LTX2_DISTILLED_TWO_STAGE_RECIPE.height, LTX2_DISTILLED_TWO_STAGE_RECIPE.width) == (1024, 1536)
|
||||
|
||||
|
||||
def test_ltx23_full_distilled_route_uses_merged_weights_without_lora():
|
||||
assert resolve_ltx_component_profile("distilled_two_stage", "2.3") is LTX23_DISTILLED_COMPONENT_PROFILE
|
||||
assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.3") is LTX23_DISTILLED_TWO_STAGE_RECIPE
|
||||
assert LTX23_DISTILLED_TWO_STAGE_RECIPE is LTX2_DISTILLED_TWO_STAGE_RECIPE
|
||||
assert all(phase.adapter_slot is None for phase in LTX23_DISTILLED_TWO_STAGE_RECIPE.phases)
|
||||
assert LTX23_DISTILLED_COMPONENT_PROFILE.distilled_lora_filename is None
|
||||
assert LTX23_DISTILLED_COMPONENT_PROFILE.vocoder_cls is LTX23_COMPONENT_PROFILE.vocoder_cls
|
||||
assert LTX23_DISTILLED_COMPONENT_PROFILE.latent_upsampler_filename == (
|
||||
"ltx-2.3-spatial-upscaler-x2-1.1.safetensors"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("request_inputs", "error"),
|
||||
[
|
||||
@@ -958,6 +1207,98 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting():
|
||||
torch.testing.assert_close(timesteps, expected[:-1] * scheduler.config.num_train_timesteps)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("recipe", "steps", "stg_block"),
|
||||
[
|
||||
(LTX2_TWO_STAGE_RECIPE, 40, 29),
|
||||
(LTX23_TWO_STAGE_RECIPE, 30, 28),
|
||||
],
|
||||
)
|
||||
def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, steps, stg_block):
|
||||
stage1, stage2 = recipe.phases
|
||||
assert (recipe.width, recipe.height) == (1536, 1024)
|
||||
assert recipe.num_inference_steps == steps
|
||||
assert stage1.spatial_downscale == 2
|
||||
assert stage1.adapter_slot is None
|
||||
assert stage1.guidance.video.stg_blocks == (stg_block,)
|
||||
assert stage1.guidance.audio.stg_blocks == (stg_block,)
|
||||
assert stage2.input_transform == "spatial_upsample"
|
||||
assert stage2.adapter_slot == "ltx_distilled"
|
||||
assert stage2.guidance == LTXGuidanceSpec.positive_only()
|
||||
assert stage2.num_inference_steps == 3
|
||||
assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0)
|
||||
# Official Stage 2 refines video only; its audio result is discarded.
|
||||
assert (recipe.video_output_phase, recipe.audio_output_phase) == (1, 0)
|
||||
|
||||
|
||||
def test_ltx_two_stage_entries_select_the_official_adapter_slots():
|
||||
phase_switches = []
|
||||
ordinary_pipeline = object.__new__(LTX2TwoStagePipeline)
|
||||
ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append)
|
||||
for phase in LTX2_TWO_STAGE_RECIPE.phases:
|
||||
ordinary_pipeline._enter_phase(phase)
|
||||
assert phase_switches == [None, "ltx_distilled"]
|
||||
|
||||
distilled_pipeline = object.__new__(LTX2DistilledPipeline)
|
||||
for phase in LTX2_DISTILLED_TWO_STAGE_RECIPE.phases:
|
||||
distilled_pipeline._enter_phase(phase)
|
||||
|
||||
|
||||
def test_ltx23_gate_projection_declares_its_global_column_parallel_layout(monkeypatch):
|
||||
class FakeColumn(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_size,
|
||||
output_size,
|
||||
*,
|
||||
bias,
|
||||
gather_output,
|
||||
return_bias,
|
||||
quant_config,
|
||||
prefix,
|
||||
):
|
||||
super().__init__()
|
||||
self.input_size = input_size
|
||||
self.output_size = output_size
|
||||
self.bias_enabled = bias
|
||||
self.gather_output = gather_output
|
||||
self.return_bias = return_bias
|
||||
self.quant_config = quant_config
|
||||
self.prefix = prefix
|
||||
|
||||
class FakeRow(torch.nn.Module):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__()
|
||||
|
||||
class FakeAttention(torch.nn.Module):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__()
|
||||
|
||||
monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumn)
|
||||
monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRow)
|
||||
monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeAttention)
|
||||
monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2)
|
||||
|
||||
attention = ltx2_transformer.LTX2Attention(
|
||||
query_dim=16,
|
||||
heads=4,
|
||||
dim_head=4,
|
||||
qk_norm="rms_norm",
|
||||
apply_gated_attention=True,
|
||||
pack_qkv=False,
|
||||
prefix="transformer_blocks.0.attn1",
|
||||
)
|
||||
|
||||
gate = attention.to_gate_logits
|
||||
assert isinstance(gate, FakeColumn)
|
||||
assert gate.input_size == 16
|
||||
assert gate.output_size == 4
|
||||
assert gate.gather_output is False
|
||||
assert gate.return_bias is False
|
||||
assert gate.quant_config is None
|
||||
assert gate.prefix == "transformer_blocks.0.attn1.to_gate_logits"
|
||||
|
||||
|
||||
class TestLTXRequestParsing:
|
||||
def test_request_resolves_official_and_explicit_negative_prompts_per_sample(self):
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
@@ -1432,6 +1773,7 @@ class TestPipelineComponents:
|
||||
|
||||
def test_distilled_profile_keeps_upsampler_in_managed_resident_components(self):
|
||||
assert LTX2_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler")
|
||||
assert LTX23_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler")
|
||||
|
||||
|
||||
class TestLTX23DecodeConditioning:
|
||||
@@ -1477,6 +1819,11 @@ class TestRegistryIntegration:
|
||||
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
|
||||
|
||||
assert _DIFFUSION_MODELS["LTX2Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX2Pipeline")
|
||||
assert _DIFFUSION_MODELS["LTX2TwoStagePipeline"] == (
|
||||
"ltx2",
|
||||
"pipeline_ltx2_two_stage",
|
||||
"LTX2TwoStagePipeline",
|
||||
)
|
||||
assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == (
|
||||
"ltx2",
|
||||
"pipeline_ltx2_two_stage",
|
||||
@@ -1488,6 +1835,9 @@ class TestRegistryIntegration:
|
||||
"LTX23ImageToVideoPipeline",
|
||||
"LTX2TwoStagesPipeline",
|
||||
"LTX2ImageToVideoTwoStagesPipeline",
|
||||
"LTX2ImageToVideoTwoStagePipeline",
|
||||
"LTX23TwoStagePipeline",
|
||||
"LTX23ImageToVideoTwoStagePipeline",
|
||||
}
|
||||
assert removed_entries.isdisjoint(_DIFFUSION_MODELS)
|
||||
|
||||
@@ -1497,6 +1847,7 @@ class TestRegistryIntegration:
|
||||
|
||||
expected = [
|
||||
"LTX2Pipeline",
|
||||
"LTX2TwoStagePipeline",
|
||||
"LTX2DistilledPipeline",
|
||||
"LTX2T2VDMD2Pipeline",
|
||||
"LTX2I2VDMD2Pipeline",
|
||||
@@ -1509,6 +1860,7 @@ class TestRegistryIntegration:
|
||||
"model_class_name",
|
||||
[
|
||||
"LTX2Pipeline",
|
||||
"LTX2TwoStagePipeline",
|
||||
"LTX2DistilledPipeline",
|
||||
"LTX2T2VDMD2Pipeline",
|
||||
"LTX2I2VDMD2Pipeline",
|
||||
|
||||
@@ -1,13 +1,20 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel, _make_rms_norm
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import (
|
||||
LTX2AudioVideoAttnProcessor,
|
||||
LTX2VideoTransformer3DModel,
|
||||
_LTX2ParallelAttention,
|
||||
_make_rms_norm,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
@@ -66,3 +73,83 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
|
||||
assert video_timestep.expected_dims == audio_timestep.expected_dims == 2
|
||||
assert not video_timestep.split_output
|
||||
assert not audio_timestep.split_output
|
||||
|
||||
|
||||
def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):
|
||||
monkeypatch.setattr(LTX2AudioVideoAttnProcessor, "_is_sp_enabled", staticmethod(lambda: True))
|
||||
processor = LTX2AudioVideoAttnProcessor()
|
||||
hidden_states = torch.zeros(1, 3, 4)
|
||||
encoder_hidden_states = torch.zeros(1, 4, 4)
|
||||
key_padding_mask = torch.tensor([[True, True, True, False]])
|
||||
|
||||
actual = processor._prepare_attention_mask(
|
||||
SimpleNamespace(),
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
key_padding_mask,
|
||||
batch_size=1,
|
||||
sequence_length=4,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual, key_padding_mask[:, None, None, :])
|
||||
|
||||
|
||||
def _fake_ltx_parallel_attention():
|
||||
calls = []
|
||||
|
||||
class Backend:
|
||||
@staticmethod
|
||||
def get_name():
|
||||
return "FLASH_ATTN"
|
||||
|
||||
@staticmethod
|
||||
def supports_attention_mask():
|
||||
return True
|
||||
|
||||
def run(name):
|
||||
def forward(query, _key, _value, _metadata):
|
||||
calls.append(name)
|
||||
return torch.zeros_like(query)
|
||||
|
||||
return forward
|
||||
|
||||
layer = object.__new__(_LTX2ParallelAttention)
|
||||
nn.Module.__init__(layer)
|
||||
layer.attn_backend = Backend
|
||||
layer.attention = SimpleNamespace(forward=run("native"))
|
||||
layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa"))
|
||||
layer.backend_pref = "FLASH_ATTN"
|
||||
return layer, calls
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)])
|
||||
def test_ltx_masked_cross_attention_uses_sdpa(query_length, key_length):
|
||||
layer, calls = _fake_ltx_parallel_attention()
|
||||
query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16)
|
||||
key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16)
|
||||
key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]])
|
||||
metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
|
||||
|
||||
layer._run_local_attention(query, key, value, metadata)
|
||||
|
||||
assert calls == ["sdpa"]
|
||||
|
||||
|
||||
def test_ltx_masked_self_attention_uses_sdpa():
|
||||
layer, calls = _fake_ltx_parallel_attention()
|
||||
query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
|
||||
key_padding_mask = torch.tensor([[True, True, True, False]])
|
||||
metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
|
||||
|
||||
layer._run_local_attention(query, key, value, metadata)
|
||||
|
||||
assert calls == ["sdpa"]
|
||||
|
||||
|
||||
def test_ltx_unmasked_cross_attention_keeps_native_backend():
|
||||
layer, calls = _fake_ltx_parallel_attention()
|
||||
query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
|
||||
|
||||
layer._run_local_attention(query, key, value, None)
|
||||
|
||||
assert calls == ["native"]
|
||||
@@ -0,0 +1,26 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLRMSNorm
|
||||
|
||||
|
||||
def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None:
|
||||
"""Qwen keeps BF16 gamma while native fallback accumulates in FP32."""
|
||||
eps = 1e-6
|
||||
norm = MiniMaxH3Qwen3VLRMSNorm(hidden_size=4, eps=eps, dtype=torch.bfloat16)
|
||||
norm.weight.data.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0], dtype=torch.bfloat16))
|
||||
x = torch.tensor(
|
||||
[[1.0, -2.0, 3.0, -4.0], [-5.0, 6.0, -7.0, 8.0]],
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
x_fp32 = x.float()
|
||||
expected = (x_fp32 * torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + eps) * norm.weight.float()).to(x.dtype)
|
||||
|
||||
assert isinstance(norm, RMSNorm)
|
||||
assert norm.weight.dtype == torch.bfloat16
|
||||
assert set(norm.state_dict()) == {"weight"}
|
||||
torch.testing.assert_close(norm.forward_native(x), expected, atol=0, rtol=0)
|
||||
Loaded 100 of 306 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user