Compare commits
34
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 hsliuustc0106
|
5390783232 | ||
|
hsliuustc0106
|
fe7f5ee7bb | ||
|
hsliuustc0106
|
83a153bf70 | ||
|
hsliuustc0106
|
901a8db21f | ||
|
hsliuustc0106
|
2d5e427f88 | ||
|
Samit
|
956f53b278 | ||
|
Ziming Huang
|
d2aa9cf08a | ||
|
|
2d980133fe | ||
|
|
e2b0ee4b87 | ||
|
Samit
|
a7bf4050de | ||
|
LiBai
|
9225039d17 | ||
|
Ting FU
|
0e8e630c5b | ||
|
Zhengyuan Su (苏政渊)andClaude Opus 4.6
|
92c788e5fd | ||
|
汪志鹏
|
9d87229a30 | ||
|
pjh4993
|
ed7a448cd3 | ||
|
wangyu
|
0edc356fc8 | ||
|
Nick CaoandClaude
|
3bd8a5239e | ||
|
zhumingjue138
|
e6f88f7c5a | ||
|
Zhengyuan Su (苏政渊)andClaude Opus 4.6
|
c3f1042061 | ||
|
Nick CaoandClaude
|
149b9f179a | ||
|
vveerrggandvveerrgg
|
2c6c07c4f6 | ||
|
Peiqi Yin
|
fcda835f4d | ||
|
Dnoob
|
7e7efdd1ab | ||
|
Jinheng
|
aefa2ee45c | ||
|
R0CKSTAR
|
fb3c6bd9b1 | ||
|
|
c3c736d488 | ||
|
Lancer
|
8609bc8ed9 | ||
|
Yangshen DengandClaude Opus 4.6
|
1cd5210440 | ||
|
Y. FisherandCanlin Guo
|
ec082add35 | ||
|
|
cb6a8739c1 | ||
|
skfandCanlin Guo
|
6433847249 | ||
|
Yuanheng Zhao
|
8a55d3d01f | ||
|
Kevin H. Luu
|
b246617fe5 | ||
|
Kevin H. Luu
|
bc5e945544 |
No files matched your search
@@ -4,6 +4,7 @@
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
from dataclasses import asdict, dataclass
|
||||
from datetime import datetime
|
||||
@@ -11,8 +12,6 @@ from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import quote
|
||||
|
||||
import re
|
||||
|
||||
|
||||
def normalize_package_name(name: str) -> str:
|
||||
"""Normalize package name per PEP 503."""
|
||||
|
||||
@@ -0,0 +1,367 @@
|
||||
# Nightly diffusion GPU tests — appended to the main nightly build via
|
||||
# buildkite-agent pipeline upload .buildkite/test-nightly-diffusion.yml
|
||||
# from test-nightly.yml (step key: nightly-diffusion-model-test). Top-level groups are
|
||||
# foldable in the Buildkite UI (Other / Wan / Qwen-Image).
|
||||
steps:
|
||||
- group: ":card_index_dividers: Other Model Test"
|
||||
key: nightly-other-model-test-group
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion · Other · Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not test_qwen_image" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion · Other · Function Test with L4"
|
||||
timeout_in_minutes: 60
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: ":full_moon: Diffusion · Other · Doc Test"
|
||||
timeout_in_minutes: 60
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- pytest -s -v tests/examples/online_serving/test_text_to_image.py tests/examples/offline_inference/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- group: ":card_index_dividers: Wan Series Model Test"
|
||||
key: nightly-wan-model-test-group
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion · Wan · Function Test"
|
||||
timeout_in_minutes: 90
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py -m "advanced_model" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion · Wan · Accuracy Test"
|
||||
key: nightly-wan22-i2v-accuracy
|
||||
timeout_in_minutes: 180
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --run-level advanced_model
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- group: ":card_index_dividers: Qwen-Image Series Model Test"
|
||||
key: nightly-qwen-image-edit-group
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion · Qwen-Image · Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen_image*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion · Qwen-Image · GEBench Accuracy Test"
|
||||
key: nightly-gebench-accuracy
|
||||
timeout_in_minutes: 60
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion · Qwen-Image · GEdit-Bench Accuracy Test"
|
||||
key: nightly-gedit-bench-accuracy
|
||||
timeout_in_minutes: 60
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
|
||||
value: "120"
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion · Qwen-Image · Perf Test"
|
||||
key: nightly-qwen-image-performance
|
||||
timeout_in_minutes: 180
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/benchmark_results_*.json"
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 4
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
+183
-388
@@ -1,228 +1,199 @@
|
||||
steps:
|
||||
- label: ":full_moon: Omni Model Test with H100"
|
||||
timeout_in_minutes: 90
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
pytest -s -v \
|
||||
tests/examples/ \
|
||||
tests/e2e/online_serving/test_*_expansion.py \
|
||||
-m "advanced_model and H100 and omni" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
# Group: collapses under one heading in the Buildkite UI; child steps still run in parallel.
|
||||
- group: ":card_index_dividers: Omni Model Test"
|
||||
key: nightly-omni-test-group
|
||||
steps:
|
||||
- label: ":full_moon: Omni · Function Test with H100"
|
||||
timeout_in_minutes: 90
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and H100 and omni" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Omni Model Test with L4"
|
||||
timeout_in_minutes: 90
|
||||
- label: ":full_moon: Omni · Function Test with L4"
|
||||
timeout_in_minutes: 90
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: ":full_moon: Omni · Doc Test with L4"
|
||||
timeout_in_minutes: 90
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/examples/ -m "advanced_model and omni and L4" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: ":full_moon: Omni · Doc Test with H100"
|
||||
timeout_in_minutes: 90
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/examples/ -m "advanced_model and omni and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Omni · Perf Test"
|
||||
key: nightly-omni-performance
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# Dynamically appends steps from test-nightly-diffusion.yml into this build (same mechanism as
|
||||
# pipeline.yml → test-ready.yml / test-merge.yml / test-nightly.yml). Foldable groups stay in the
|
||||
# uploaded YAML (Other / Wan / Qwen-Image).
|
||||
- label: ":card_index_dividers: Diffusion Model Test"
|
||||
key: nightly-diffusion-model-test
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/examples/ tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly-diffusion.yml
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
- label: ":full_moon: Diffusion Model Test with H100"
|
||||
- label: ":bar_chart: Testcase Statistics"
|
||||
key: nightly-testcase-statistics
|
||||
timeout_in_minutes: 120
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion Model (Wan) Test with H100"
|
||||
timeout_in_minutes: 90
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py -m "advanced_model" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion Model Test"
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
|
||||
- label: ":full_moon: Doc Example Code Test with H100"
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- pytest -s -v tests/examples/online_serving/test_text_to_image.py tests/examples/offline_inference/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Omni Model Perf Test & Testcase Statistics with H100"
|
||||
key: nightly-omni-performance
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
- python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
|
||||
agents:
|
||||
@@ -259,189 +230,13 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: GEBench Accuracy Test with H100"
|
||||
key: nightly-gebench-accuracy
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: GEdit-Bench Accuracy Test with H100"
|
||||
key: nightly-gedit-bench-accuracy
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
|
||||
value: "120"
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Wan22 I2V Accuracy Test with H100"
|
||||
key: nightly-wan22-i2v-accuracy
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --run-level advanced_model
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion Perf Test with H100"
|
||||
key: nightly-qwen-image-performance
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/benchmark_results_*.json"
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 4
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# No need to run this step for PRs with label nightly-test
|
||||
- label: ":email: Nightly Collection & Email"
|
||||
key: nightly-perf-distribution
|
||||
depends_on:
|
||||
- nightly-omni-performance
|
||||
- nightly-qwen-image-performance
|
||||
- nightly-testcase-statistics
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- pip install openpyxl
|
||||
|
||||
+18
-21
@@ -16,11 +16,10 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Voxtral TTS CUDA Unit Test"
|
||||
timeout_in_minutes: 10
|
||||
- label: "CUDA Unit Test with single card"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- "timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py"
|
||||
- timeout 10m pytest -v -s -m 'core_model and cuda and L4 and not distributed_cuda' --ignore=tests/e2e --ignore=tests/engine/test_async_omni_engine_abort.py --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
@@ -33,6 +32,22 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "CUDA Unit Test with multi cards"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- timeout 10m pytest -v -s -m 'core_model and cuda and L4 and distributed_cuda' --ignore=tests/e2e --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
@@ -152,24 +167,6 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
|
||||
- label: "Engine Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
|
||||
@@ -0,0 +1,534 @@
|
||||
---
|
||||
name: add-diffusion-model
|
||||
description: Add a new diffusion model (text-to-image, text-to-video, image-to-video, text-to-audio, image editing) to vLLM-Omni, including Cache-DiT acceleration and parallelism support (TP, SP/USP, CFG-Parallel, HSDP). Use when integrating a new diffusion model, porting a diffusers pipeline or a custom model repo to vllm-omni, creating a new DiT transformer adapter, adding diffusion model support, or enabling multi-GPU parallelism and cache acceleration for an existing model.
|
||||
---
|
||||
|
||||
# Adding a Diffusion Model to vLLM-Omni
|
||||
|
||||
## Overview
|
||||
|
||||
This skill guides you through adding a new diffusion model to vLLM-Omni. The model may come from HuggingFace Diffusers (structured pipeline) or from a private/custom repo. The workflow differs significantly depending on the source.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Before starting, determine:
|
||||
|
||||
1. **Model category**: Text-to-Image, Text-to-Video, Image-to-Video, Image Editing, Text-to-Audio, or Omni
|
||||
2. **Reference source**: Diffusers pipeline, custom repo, or a combination
|
||||
3. **Model HuggingFace ID** or local checkpoint path
|
||||
4. **Architecture**: Scheduler, text encoder, VAE, transformer/backbone
|
||||
|
||||
## Step 0: Classify the Migration Path
|
||||
|
||||
Check the model's HF repo for `model_index.json`. This determines your path:
|
||||
|
||||
| Scenario | How to identify | Migration path |
|
||||
|----------|----------------|----------------|
|
||||
| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip to Step 5 (test) and Step 7 (docs) |
|
||||
| **Diffusers-based** | Has standard `model_index.json` with `_diffusers_version`, subfolders for `transformer/`, `vae/`, etc. | Follow **Path A** below |
|
||||
| **Custom/private repo** | No diffusers `model_index.json`, weights in non-standard format, custom model code in a separate git repo | Follow **Path B** below |
|
||||
| **Hybrid** | Has some diffusers components (VAE) but custom transformer/fusion | Mix of Path A and Path B |
|
||||
|
||||
## Path A: Diffusers-Based Model
|
||||
|
||||
For models with a standard diffusers layout. See [references/transformer-adaptation.md](references/transformer-adaptation.md) for detailed code patterns.
|
||||
|
||||
### A1. Analyze `model_index.json`
|
||||
|
||||
Identify components: `transformer`, `scheduler`, `vae`, `text_encoder`, `tokenizer`.
|
||||
|
||||
### A2. Create model directory
|
||||
|
||||
```
|
||||
vllm_omni/diffusion/models/your_model_name/
|
||||
├── __init__.py
|
||||
├── pipeline_your_model.py
|
||||
└── your_model_transformer.py
|
||||
```
|
||||
|
||||
### A3. Adapt transformer
|
||||
|
||||
1. Copy from diffusers source. Remove mixins (`ModelMixin`, `ConfigMixin`, `AttentionModuleMixin`).
|
||||
2. Replace attention with `vllm_omni.diffusion.attention.layer.Attention` (QKV shape: `[B, seq, heads, head_dim]`).
|
||||
3. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`.
|
||||
4. Add `load_weights()` method mapping diffusers weight names to vllm-omni names.
|
||||
5. Add class attributes: `_repeated_blocks`, `_layerwise_offload_blocks_attr`.
|
||||
|
||||
### A4. Adapt pipeline
|
||||
|
||||
Inherit from `nn.Module`. The key contract:
|
||||
|
||||
```python
|
||||
class YourPipeline(nn.Module):
|
||||
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
|
||||
# Load VAE, text encoder, tokenizer via from_pretrained()
|
||||
# Instantiate transformer (weights loaded later via weights_sources)
|
||||
self.weights_sources = [
|
||||
DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path=od_config.model, subfolder="transformer",
|
||||
prefix="transformer.", fall_back_to_pt=True)]
|
||||
|
||||
def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
|
||||
# Encode prompt → prepare latents → denoise loop → VAE decode
|
||||
return DiffusionOutput(output=output)
|
||||
|
||||
def load_weights(self, weights):
|
||||
return AutoWeightsLoader(self).load_weights(weights)
|
||||
```
|
||||
|
||||
Add post/pre-process functions in the same pipeline file. Register them in `registry.py`.
|
||||
|
||||
### A5. Register, test, docs → continue at Step 4 below.
|
||||
|
||||
---
|
||||
|
||||
## Path B: Custom/Private Repo Model
|
||||
|
||||
For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3.
|
||||
|
||||
### B1. Understand the reference repo
|
||||
|
||||
Study the original model's code to identify:
|
||||
- **Model architecture files** (transformers, fusion modules, embeddings)
|
||||
- **Weight format** (safetensors, `.pth`, custom checkpoint structure)
|
||||
- **Weight loading helpers** (custom init functions, checkpoint loaders)
|
||||
- **Pre/post-processing** (image/audio transforms, tokenization, VAE encode/decode)
|
||||
- **External dependencies** (packages not on PyPI)
|
||||
- **Config format** (JSON config files, hardcoded dicts)
|
||||
|
||||
### B2. Decide what lives WHERE
|
||||
|
||||
This is the key design decision for custom models. Follow these placement rules:
|
||||
|
||||
| Code type | Where to place | Example |
|
||||
|-----------|---------------|---------|
|
||||
| **Pipeline orchestration** (init, forward, denoise loop) | `vllm_omni/diffusion/models/<name>/pipeline_<name>.py` | Always required |
|
||||
| **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models/<name>/<name>_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` |
|
||||
| **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models/<name>/` as separate files | `autoencoder.py`, `fusion.py` |
|
||||
| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | `dreamid_omni` package via git clone |
|
||||
| **Hardcoded model configs** | Module-level dicts in pipeline file | `VIDEO_CONFIG`, `AUDIO_CONFIG` dicts |
|
||||
| **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_dreamid_omni.py` |
|
||||
| **Custom `model_index.json`** | Generated by download script, placed at model root | Minimal: `{"_class_name": "YourPipeline", ...}` |
|
||||
|
||||
### B3. Handle external dependencies
|
||||
|
||||
If the model's code lives in a separate git repo:
|
||||
|
||||
**Option 1: Import with graceful fallback** (recommended for models with external utils)
|
||||
|
||||
```python
|
||||
try:
|
||||
from external_model.utils import init_vae, load_checkpoint
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Failed to import from dependency 'external_model'. "
|
||||
"Please run the download script first."
|
||||
)
|
||||
```
|
||||
|
||||
**Option 2: Port the code directly** (preferred when feasible)
|
||||
|
||||
Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and adapt them. This avoids external dependencies. BAGEL does this — `autoencoder.py` and `bagel_transformer.py` are ported directly.
|
||||
|
||||
**Decision criteria**: Port if the code is self-contained and won't diverge. Use external deps if the model repo is actively maintained and the code is complex.
|
||||
|
||||
### B4. Handle custom weight loading
|
||||
|
||||
Custom models have two patterns for weight loading:
|
||||
|
||||
**Pattern 1: Bypass standard loader** (DreamID-Omni style)
|
||||
|
||||
When the original model has complex custom init functions that load weights in `__init__`:
|
||||
|
||||
```python
|
||||
class CustomPipeline(nn.Module):
|
||||
def __init__(self, *, od_config, prefix=""):
|
||||
super().__init__()
|
||||
model = od_config.model
|
||||
# Load everything eagerly in __init__ using custom helpers
|
||||
self.vae = custom_init_vae(model, device=self.device)
|
||||
self.text_encoder = custom_init_text_encoder(model, device=self.device)
|
||||
self.transformer = CustomFusionModel(CONFIG)
|
||||
load_custom_checkpoint(self.transformer,
|
||||
checkpoint_path=os.path.join(model, "model.safetensors"))
|
||||
# NO weights_sources defined — bypasses standard loader
|
||||
|
||||
def load_weights(self, weights):
|
||||
pass # No-op — all weights loaded in __init__
|
||||
```
|
||||
|
||||
**Pattern 2: Use standard loader with custom `load_weights`** (BAGEL style)
|
||||
|
||||
When weights are in safetensors format but need name remapping:
|
||||
|
||||
```python
|
||||
class CustomPipeline(nn.Module):
|
||||
def __init__(self, *, od_config, prefix=""):
|
||||
super().__init__()
|
||||
# Instantiate model architecture without weights
|
||||
self.bagel = BagelModel(config)
|
||||
self.vae = AutoEncoder(ae_params)
|
||||
|
||||
# Point loader at the safetensors in the model root
|
||||
self.weights_sources = [
|
||||
DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path=od_config.model,
|
||||
subfolder=None, # weights at root, not in subfolder
|
||||
prefix="",
|
||||
fall_back_to_pt=False,
|
||||
)
|
||||
]
|
||||
|
||||
def load_weights(self, weights):
|
||||
# Custom name remapping for non-diffusers weight names
|
||||
params = dict(self.named_parameters())
|
||||
loaded = set()
|
||||
for name, tensor in weights:
|
||||
# Remap original weight names to vllm-omni module names
|
||||
name = self._remap_weight_name(name)
|
||||
if name in params:
|
||||
default_weight_loader(params[name], tensor)
|
||||
loaded.add(name)
|
||||
return loaded
|
||||
```
|
||||
|
||||
### B5. Create the `model_index.json`
|
||||
|
||||
Custom models need a `model_index.json` at the model root for vllm-omni to discover them. For custom models, this is minimal:
|
||||
|
||||
```json
|
||||
{
|
||||
"_class_name": "YourModelPipeline",
|
||||
"custom_key": "path/to/custom_weights.safetensors"
|
||||
}
|
||||
```
|
||||
|
||||
The `_class_name` must match a key in `_DIFFUSION_MODELS` in `registry.py`. Additional keys are model-specific (accessed via `od_config.model_config`).
|
||||
|
||||
If the model's weights come from multiple HF repos, write a **download script** that:
|
||||
1. Downloads from each repo
|
||||
2. Assembles into a single directory
|
||||
3. Generates `model_index.json`
|
||||
4. Installs any external dependencies (git clone + `.pth` file)
|
||||
|
||||
Place at: `examples/offline_inference/<name>/download_<name>.py`
|
||||
|
||||
### B6. Handle multi-modal inputs
|
||||
|
||||
If the model accepts images, audio, or other multi-modal inputs, implement the protocol classes from `vllm_omni/diffusion/models/interface.py`:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.models.interface import SupportImageInput, SupportAudioInput
|
||||
|
||||
class MyPipeline(nn.Module, SupportImageInput, SupportAudioInput):
|
||||
# Protocol markers — the engine uses these to enable proper input routing
|
||||
pass
|
||||
```
|
||||
|
||||
Preprocessing for custom models is typically done **inside `forward()`** rather than via registered pre-process functions, since the logic is often tightly coupled to the model.
|
||||
|
||||
### B7. Continue at Step 4 below.
|
||||
|
||||
---
|
||||
|
||||
## Common Steps (Both Paths)
|
||||
|
||||
### Step 4: Register Model in registry.py
|
||||
|
||||
Edit `vllm_omni/diffusion/registry.py`:
|
||||
|
||||
```python
|
||||
_DIFFUSION_MODELS = {
|
||||
"YourModelPipeline": ("your_model_name", "pipeline_your_model", "YourModelPipeline"),
|
||||
}
|
||||
_DIFFUSION_POST_PROCESS_FUNCS = {
|
||||
"YourModelPipeline": "get_your_model_post_process_func", # if applicable
|
||||
}
|
||||
_DIFFUSION_PRE_PROCESS_FUNCS = {
|
||||
"YourModelPipeline": "get_your_model_pre_process_func", # if applicable
|
||||
}
|
||||
```
|
||||
|
||||
The registry key is the `_class_name` from `model_index.json`. The tuple is `(folder_name, module_file, class_name)`.
|
||||
|
||||
Create `__init__.py` exporting the pipeline class and any factory functions.
|
||||
|
||||
### Step 5: Run, Test, Debug
|
||||
|
||||
Use the appropriate existing example script:
|
||||
|
||||
| Category | Script |
|
||||
|----------|--------|
|
||||
| Text-to-Image | `examples/offline_inference/text_to_image/text_to_image.py` |
|
||||
| Text-to-Video | `examples/offline_inference/text_to_video/text_to_video.py` |
|
||||
| Image-to-Video | `examples/offline_inference/image_to_video/image_to_video.py` |
|
||||
| Image-to-Image | `examples/offline_inference/image_to_image/image_edit.py` |
|
||||
| Text-to-Audio | `examples/offline_inference/text_to_audio/text_to_audio.py` |
|
||||
|
||||
For custom/Omni models that don't fit these categories, create a dedicated example script.
|
||||
|
||||
**Validation**: No errors, output is meaningful, quality matches reference implementation.
|
||||
|
||||
See [references/troubleshooting.md](references/troubleshooting.md) for common errors.
|
||||
|
||||
### Step 6: Add Example Scripts
|
||||
|
||||
For Omni or custom models, create:
|
||||
- `examples/offline_inference/your_model_name/` — offline script + README
|
||||
- `examples/online_serving/your_model_name/` — server script + client
|
||||
- Download script if weights require assembly from multiple sources
|
||||
|
||||
### Step 7: Update Documentation
|
||||
|
||||
Required updates:
|
||||
1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table
|
||||
2. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
|
||||
3. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
|
||||
4. `examples/offline_inference/xxx/README.md` — offline example docs
|
||||
5. `examples/online_serve/xxx/README.md` — online serve docs
|
||||
|
||||
### Step 8: Add E2E Tests (Recommended)
|
||||
|
||||
Create `tests/e2e/online_serving/test_your_model_expansion.py`.
|
||||
|
||||
### Step 9: Add Cache-DiT Acceleration
|
||||
|
||||
Cache-DiT accelerates inference by caching intermediate computation results across denoising steps. After your model is working correctly on a single GPU, add cache-dit support.
|
||||
|
||||
See [references/cache-dit-patterns.md](references/cache-dit-patterns.md) for detailed code patterns.
|
||||
|
||||
#### 9a. Determine your model type
|
||||
|
||||
| Model Type | Description | Action |
|
||||
|------------|-------------|--------|
|
||||
| **Standard single-transformer** | One transformer with one `ModuleList` of blocks | No code needed — `CacheDiTBackend` auto-detects via `enable_cache_for_dit()` |
|
||||
| **Multi-block-list** | One transformer with multiple block lists (e.g., `transformer_blocks` + `single_transformer_blocks`) | Write custom enabler with `BlockAdapter` |
|
||||
| **Dual-transformer** | Two transformers (e.g., high-noise + low-noise) | Write custom enabler with `BlockAdapter` wrapping both |
|
||||
|
||||
#### 9b. Standard models — verify automatic support
|
||||
|
||||
For standard single-transformer models, test directly:
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="your-model-name",
|
||||
cache_backend="cache_dit",
|
||||
cache_config={
|
||||
"Fn_compute_blocks": 1,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4,
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
Check logs for "Cache-dit enabled successfully on xxx". If it works, skip to Step 9e.
|
||||
|
||||
#### 9c. Custom architectures — write a custom enabler
|
||||
|
||||
For multi-block-list or dual-transformer models, write a custom enabler function:
|
||||
|
||||
```python
|
||||
from cache_dit import BlockAdapter, ForwardPattern, ParamsModifier, DBCacheConfig
|
||||
|
||||
def enable_cache_for_your_model(pipeline, cache_config):
|
||||
db_cache_config = DBCacheConfig(
|
||||
num_inference_steps=None,
|
||||
Fn_compute_blocks=cache_config.Fn_compute_blocks,
|
||||
Bn_compute_blocks=cache_config.Bn_compute_blocks,
|
||||
max_warmup_steps=cache_config.max_warmup_steps,
|
||||
max_cached_steps=cache_config.max_cached_steps,
|
||||
max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
|
||||
residual_diff_threshold=cache_config.residual_diff_threshold,
|
||||
)
|
||||
|
||||
cache_dit.enable_cache(
|
||||
BlockAdapter(
|
||||
transformer=pipeline.transformer,
|
||||
blocks=[
|
||||
pipeline.transformer.transformer_blocks,
|
||||
pipeline.transformer.single_transformer_blocks,
|
||||
],
|
||||
forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
|
||||
params_modifiers=[ParamsModifier(...)],
|
||||
),
|
||||
cache_config=db_cache_config,
|
||||
)
|
||||
|
||||
def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
|
||||
cache_dit.refresh_context(
|
||||
pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
|
||||
)
|
||||
return refresh_cache_context
|
||||
```
|
||||
|
||||
#### 9d. Register the custom enabler
|
||||
|
||||
Add your enabler to `CUSTOM_DIT_ENABLERS` in `vllm_omni/diffusion/cache/cache_dit_backend.py`:
|
||||
|
||||
```python
|
||||
CUSTOM_DIT_ENABLERS = {
|
||||
"Wan22Pipeline": enable_cache_for_wan22,
|
||||
"LongCatImagePipeline": enable_cache_for_longcat_image,
|
||||
"YourModelPipeline": enable_cache_for_your_model, # Add here
|
||||
}
|
||||
```
|
||||
|
||||
#### 9e. Test Cache-DiT
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="your-model-name",
|
||||
cache_backend="cache_dit",
|
||||
cache_config={
|
||||
"Fn_compute_blocks": 1, "Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4, "residual_diff_threshold": 0.24,
|
||||
}
|
||||
)
|
||||
images = omni.generate("a beautiful landscape",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=50))
|
||||
```
|
||||
|
||||
**Verify**: 1) logs show cache enabled, 2) 1.5-2x speedup, 3) output quality acceptable vs baseline.
|
||||
|
||||
If quality degrades, lower `residual_diff_threshold` (try 0.12-0.18) or increase `max_warmup_steps` (try 6-8).
|
||||
|
||||
---
|
||||
|
||||
### Step 10: Add Parallelism Support
|
||||
|
||||
After the model works on a single GPU, add multi-GPU parallelism. Add each type incrementally, testing after each addition.
|
||||
|
||||
See [references/parallelism-patterns.md](references/parallelism-patterns.md) for detailed code patterns and API reference.
|
||||
|
||||
**Recommended order**: TP → SP/USP → CFG Parallel → HSDP
|
||||
|
||||
#### 10a. Tensor Parallelism (TP)
|
||||
|
||||
Shards DiT linear layers across GPUs. Requires code changes in the transformer.
|
||||
|
||||
**What to change in the transformer**:
|
||||
1. Replace `nn.Linear` with `ColumnParallelLinear` / `RowParallelLinear` / `QKVParallelLinear`
|
||||
2. Update `load_weights()` to handle QKV fusion with `stacked_params_mapping`
|
||||
3. Use `self.to_qkv.num_heads` (local heads) instead of total heads for split sizes
|
||||
|
||||
```python
|
||||
from vllm.model_executor.layers.linear import (
|
||||
QKVParallelLinear, RowParallelLinear, ColumnParallelLinear,
|
||||
)
|
||||
|
||||
# Attention: QKV → RowParallel output
|
||||
self.to_qkv = QKVParallelLinear(dim, head_dim, num_heads, num_kv_heads)
|
||||
self.to_out = RowParallelLinear(dim, dim, input_is_parallel=True)
|
||||
|
||||
# FFN: ColumnParallel → RowParallel
|
||||
self.w1 = ColumnParallelLinear(dim, ffn_dim)
|
||||
self.w2 = RowParallelLinear(ffn_dim, dim, input_is_parallel=True)
|
||||
```
|
||||
|
||||
**Constraints**: `num_heads % tp_size == 0` and `num_kv_heads % tp_size == 0`.
|
||||
|
||||
**Test**: `--tensor-parallel-size 2`
|
||||
|
||||
#### 10b. Sequence Parallelism (SP / USP)
|
||||
|
||||
Splits sequence tokens across GPUs. Non-intrusive via `_sp_plan` on the transformer class — no changes to `forward()`.
|
||||
|
||||
**What to change in the transformer**:
|
||||
|
||||
Add `_sp_plan` class attribute:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.distributed.sp_plan import (
|
||||
SequenceParallelInput, SequenceParallelOutput,
|
||||
)
|
||||
|
||||
class YourTransformer(nn.Module):
|
||||
_sp_plan = {
|
||||
"blocks.0": {
|
||||
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
|
||||
},
|
||||
"proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
|
||||
}
|
||||
```
|
||||
|
||||
If inline tensor ops (e.g., `torch.cat`) exist between shard/gather points, extract them into `nn.Module` submodules so hooks can intercept them.
|
||||
|
||||
For RoPE that needs splitting, add an entry for the RoPE module with `split_output=True`.
|
||||
|
||||
**Test**: `--ulysses-degree 2` (offline) or `--usp 2` (online serving)
|
||||
|
||||
#### 10c. CFG Parallel
|
||||
|
||||
Distributes positive/negative CFG branches across 2 GPUs. Requires the pipeline to inherit `CFGParallelMixin`.
|
||||
|
||||
**What to change in the pipeline**:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
|
||||
class YourPipeline(nn.Module, CFGParallelMixin):
|
||||
def diffuse(self, ...) -> torch.Tensor:
|
||||
for i, t in enumerate(timesteps):
|
||||
positive_kwargs = {...}
|
||||
negative_kwargs = {...} if do_true_cfg else None
|
||||
noise_pred = self.predict_noise_maybe_with_cfg(
|
||||
do_true_cfg=do_true_cfg, true_cfg_scale=cfg_scale,
|
||||
positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs,
|
||||
)
|
||||
latents = self.scheduler_step_maybe_with_cfg(
|
||||
noise_pred, t, latents, do_true_cfg
|
||||
)
|
||||
return latents
|
||||
```
|
||||
|
||||
Override `predict_noise()` if your transformer call is non-standard. Override `combine_cfg_noise()` for multi-output models (e.g., video + audio).
|
||||
|
||||
**Constraint**: Exactly 2 GPUs. Only for models using classifier-free guidance.
|
||||
|
||||
**Test**: `--cfg-parallel-size 2`
|
||||
|
||||
#### 10d. HSDP (Hybrid Sharded Data Parallel)
|
||||
|
||||
Shards transformer weights via PyTorch FSDP2 to reduce per-GPU VRAM. No code changes to the forward pass — just add a class attribute.
|
||||
|
||||
**What to change in the transformer**:
|
||||
|
||||
```python
|
||||
class YourTransformer(nn.Module):
|
||||
@staticmethod
|
||||
def _is_transformer_block(name: str, module) -> bool:
|
||||
return "blocks" in name and name.split(".")[-1].isdigit()
|
||||
|
||||
_hsdp_shard_conditions = [_is_transformer_block]
|
||||
```
|
||||
|
||||
**Constraint**: Cannot combine with TP. For standalone HSDP, set `hsdp_shard_size` explicitly.
|
||||
|
||||
**Test**: `--use-hsdp` or `DiffusionParallelConfig(use_hsdp=True)`
|
||||
|
||||
#### 10e. Update parallelism documentation
|
||||
|
||||
After adding parallelism support, update:
|
||||
1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table
|
||||
2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch)
|
||||
|
||||
---
|
||||
|
||||
## Iterative Development Tips
|
||||
|
||||
1. **Start minimal**: Basic generation first, no parallelism/caching
|
||||
2. **Use `--enforce-eager`**: Disable torch.compile during debugging
|
||||
3. **Use small models**: Test with smaller variants first
|
||||
4. **Check tensor shapes**: Most errors are reshape mismatches in attention
|
||||
5. **Add features incrementally**: Single GPU → TP → SP → CFG → HSDP → Cache-DiT
|
||||
6. **For custom models**: Get the model running with the original code first, then progressively replace components with vllm-omni equivalents
|
||||
7. **Cache-DiT before parallelism tuning**: Cache-DiT is lossy — verify quality at baseline before combining with parallelism
|
||||
8. **Combine lossless + lossy**: e.g., TP + SP + Cache-DiT for maximum throughput
|
||||
|
||||
## Reference Files
|
||||
|
||||
- [Transformer Adaptation](references/transformer-adaptation.md) — porting transformers from diffusers
|
||||
- [Custom Model Patterns](references/custom-model-patterns.md) — patterns for non-diffusers models
|
||||
- [Parallelism Patterns](references/parallelism-patterns.md) — TP, SP/USP, CFG parallel, HSDP implementation details
|
||||
- [Cache-DiT Patterns](references/cache-dit-patterns.md) — cache-dit acceleration for standard and custom architectures
|
||||
- [Troubleshooting](references/troubleshooting.md) — common errors and fixes
|
||||
@@ -0,0 +1,254 @@
|
||||
# Cache-DiT Patterns Reference
|
||||
|
||||
## Overview
|
||||
|
||||
Cache-DiT accelerates Diffusion Transformers by caching intermediate computation results across denoising steps. Adjacent steps produce similar features, so redundant computations can be skipped.
|
||||
|
||||
Three caching strategies:
|
||||
- **DBCache**: Dynamic block-level caching — selectively computes or caches transformer blocks based on residual differences
|
||||
- **TaylorSeer**: Calibration-based prediction using Taylor expansion to estimate block outputs
|
||||
- **SCM** (Step Computation Masking): Dynamic step skipping based on configurable policies
|
||||
|
||||
**Typical speedup**: 1.5-2.5x depending on model and configuration.
|
||||
|
||||
**Official docs**: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cache_dit
|
||||
|
||||
## Architecture
|
||||
|
||||
vLLM-Omni integrates cache-dit through `CacheDiTBackend`:
|
||||
|
||||
| Component | Purpose |
|
||||
|-----------|---------|
|
||||
| `CacheDiTBackend` | Unified backend — auto-selects enabler (standard or custom) |
|
||||
| `enable_cache_for_dit()` | Default enabler for standard single-transformer models |
|
||||
| `CUSTOM_DIT_ENABLERS` dict | Registry of custom enablers keyed by pipeline class name |
|
||||
| `BlockAdapter` | Wraps complex architectures (multi-block-list or multi-transformer) |
|
||||
| `ForwardPattern` | Specifies block forward signature: `Pattern_0`, `Pattern_1`, `Pattern_2` |
|
||||
| `ParamsModifier` | Per-transformer or per-block-list config customization |
|
||||
| `DBCacheConfig` | Configuration for DBCache parameters |
|
||||
| `cache_dit.refresh_context()` | Updates cache context when `num_inference_steps` changes |
|
||||
|
||||
**Source files**:
|
||||
- `vllm_omni/diffusion/cache/cache_dit_backend.py` — `CacheDiTBackend`, enablers, `CUSTOM_DIT_ENABLERS`
|
||||
- `vllm_omni/diffusion/cache/` — cache backend implementations
|
||||
|
||||
## Standard Models: Automatic Support
|
||||
|
||||
Most DiT models follow this pattern:
|
||||
- Single transformer with one `nn.ModuleList` of blocks
|
||||
- Standard forward signature
|
||||
- Compatible with cache-dit's automatic detection
|
||||
|
||||
**Examples**: Qwen-Image, Z-Image, FLUX
|
||||
|
||||
No code changes needed. `CacheDiTBackend` automatically uses `enable_cache_for_dit()`:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_backend="cache_dit",
|
||||
cache_config={
|
||||
"Fn_compute_blocks": 1,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4,
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
What happens automatically:
|
||||
|
||||
```python
|
||||
def enable_cache_for_dit(pipeline, cache_config):
|
||||
db_cache_config = DBCacheConfig(
|
||||
num_inference_steps=None,
|
||||
Fn_compute_blocks=cache_config.Fn_compute_blocks,
|
||||
Bn_compute_blocks=cache_config.Bn_compute_blocks,
|
||||
max_warmup_steps=cache_config.max_warmup_steps,
|
||||
max_cached_steps=cache_config.max_cached_steps,
|
||||
max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
|
||||
residual_diff_threshold=cache_config.residual_diff_threshold,
|
||||
)
|
||||
|
||||
cache_dit.enable_cache(pipeline.transformer, cache_config=db_cache_config)
|
||||
|
||||
def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
|
||||
cache_dit.refresh_context(
|
||||
pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
|
||||
)
|
||||
return refresh_cache_context
|
||||
```
|
||||
|
||||
## Custom Architectures: Writing Custom Enablers
|
||||
|
||||
### When you need a custom enabler
|
||||
|
||||
- Model has multiple block lists in one transformer (e.g., `transformer_blocks` + `single_transformer_blocks`)
|
||||
- Model has two transformers (e.g., high-noise + low-noise like Wan2.2)
|
||||
- Model uses non-standard block forward signature
|
||||
|
||||
### Pattern 1: Multi-Block-List (LongCat-Image style)
|
||||
|
||||
Single transformer with two block lists:
|
||||
|
||||
```python
|
||||
import cache_dit
|
||||
from cache_dit import BlockAdapter, ForwardPattern, ParamsModifier, DBCacheConfig
|
||||
|
||||
def enable_cache_for_your_model(pipeline, cache_config):
|
||||
db_cache_config = DBCacheConfig(
|
||||
num_inference_steps=None,
|
||||
Fn_compute_blocks=cache_config.Fn_compute_blocks,
|
||||
Bn_compute_blocks=cache_config.Bn_compute_blocks,
|
||||
max_warmup_steps=cache_config.max_warmup_steps,
|
||||
max_cached_steps=cache_config.max_cached_steps,
|
||||
max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
|
||||
residual_diff_threshold=cache_config.residual_diff_threshold,
|
||||
)
|
||||
|
||||
cache_dit.enable_cache(
|
||||
BlockAdapter(
|
||||
transformer=pipeline.transformer,
|
||||
blocks=[
|
||||
pipeline.transformer.transformer_blocks,
|
||||
pipeline.transformer.single_transformer_blocks,
|
||||
],
|
||||
forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
|
||||
params_modifiers=[ParamsModifier(...)],
|
||||
),
|
||||
cache_config=db_cache_config,
|
||||
)
|
||||
|
||||
def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
|
||||
cache_dit.refresh_context(
|
||||
pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
|
||||
)
|
||||
return refresh_cache_context
|
||||
```
|
||||
|
||||
For single transformer with multiple block lists, `refresh_context` works the same as standard models — call it once on the transformer.
|
||||
|
||||
### Pattern 2: Dual-Transformer (Wan2.2 style)
|
||||
|
||||
Two transformers with separate configs:
|
||||
|
||||
```python
|
||||
def enable_cache_for_dual_transformer(pipeline, cache_config):
|
||||
db_cache_config = DBCacheConfig(...)
|
||||
|
||||
cache_dit.enable_cache(
|
||||
BlockAdapter(
|
||||
transformer=[pipeline.transformer, pipeline.transformer_2],
|
||||
blocks=[pipeline.transformer.blocks, pipeline.transformer_2.blocks],
|
||||
forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2],
|
||||
params_modifiers=[
|
||||
ParamsModifier(...), # Config for transformer 1
|
||||
ParamsModifier(...), # Config for transformer 2
|
||||
],
|
||||
),
|
||||
cache_config=db_cache_config,
|
||||
)
|
||||
|
||||
def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
|
||||
high_steps, low_steps = _split_inference_steps(num_inference_steps)
|
||||
cache_dit.refresh_context(
|
||||
pipeline.transformer, num_inference_steps=high_steps, verbose=verbose
|
||||
)
|
||||
cache_dit.refresh_context(
|
||||
pipeline.transformer_2, num_inference_steps=low_steps, verbose=verbose
|
||||
)
|
||||
return refresh_cache_context
|
||||
```
|
||||
|
||||
Key difference: `refresh_context` must be called on **each transformer separately** with its own step count.
|
||||
|
||||
### Choosing the ForwardPattern
|
||||
|
||||
| Pattern | Block forward signature | Example models |
|
||||
|---------|------------------------|----------------|
|
||||
| `Pattern_0` | `block(hidden_states, **kwargs)` → residual added inside block | Default |
|
||||
| `Pattern_1` | `block(hidden_states, **kwargs)` → returns `(hidden_states, ...)` tuple | FLUX-style single blocks |
|
||||
| `Pattern_2` | `block(hidden_states, **kwargs)` → `(hidden_states, ...)` with different residual pattern | Wan2.2 blocks |
|
||||
|
||||
Inspect your block's `forward()` return type and residual connection pattern to choose the right one. See [Cache-DiT API Reference](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/) for details.
|
||||
|
||||
## Registering Custom Enablers
|
||||
|
||||
Add your enabler to `CUSTOM_DIT_ENABLERS` in `vllm_omni/diffusion/cache/cache_dit_backend.py`:
|
||||
|
||||
```python
|
||||
CUSTOM_DIT_ENABLERS = {
|
||||
"Wan22Pipeline": enable_cache_for_wan22,
|
||||
"LongCatImagePipeline": enable_cache_for_longcat_image,
|
||||
"YourModelPipeline": enable_cache_for_your_model,
|
||||
}
|
||||
```
|
||||
|
||||
The key must match `pipeline.__class__.__name__`.
|
||||
|
||||
## Configuration Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `Fn_compute_blocks` | 1 | Number of blocks to always compute at the front |
|
||||
| `Bn_compute_blocks` | 0 | Number of blocks to always compute at the back |
|
||||
| `max_warmup_steps` | 4 | Steps to run without caching at the beginning |
|
||||
| `max_cached_steps` | — | Max total cached steps |
|
||||
| `max_continuous_cached_steps` | — | Max consecutive cached steps |
|
||||
| `residual_diff_threshold` | 0.24 | Threshold for deciding whether to cache a block |
|
||||
|
||||
### Tuning for quality vs speed
|
||||
|
||||
| Goal | Adjustments |
|
||||
|------|-------------|
|
||||
| **More speed, acceptable quality loss** | Higher `residual_diff_threshold` (0.24-0.4), lower `max_warmup_steps` (2-4) |
|
||||
| **Better quality, less speed** | Lower `residual_diff_threshold` (0.12-0.18), higher `max_warmup_steps` (6-8), lower `max_continuous_cached_steps` (2) |
|
||||
|
||||
## Testing
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
omni = Omni(
|
||||
model="your-model-name",
|
||||
cache_backend="cache_dit",
|
||||
cache_config={
|
||||
"Fn_compute_blocks": 1,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4,
|
||||
"residual_diff_threshold": 0.24,
|
||||
}
|
||||
)
|
||||
images = omni.generate(
|
||||
"a beautiful landscape",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=50),
|
||||
)
|
||||
```
|
||||
|
||||
CLI (online serving):
|
||||
|
||||
```bash
|
||||
vllm serve your-model --omni --port 8098 \
|
||||
--cache-backend cache_dit \
|
||||
--cache-config '{"Fn_compute_blocks": 1, "Bn_compute_blocks": 0, "max_warmup_steps": 4}'
|
||||
```
|
||||
|
||||
**Verification checklist**:
|
||||
1. Logs show "Cache-dit enabled successfully on xxx"
|
||||
2. Performance: 1.5-2x speedup vs no cache
|
||||
3. Quality: compare output with `cache_backend=None`
|
||||
|
||||
## Excluded Models
|
||||
|
||||
Models listed in `_NO_CACHE_ACCELERATION` in `vllm_omni/diffusion/registry.py` do not support cache-dit (e.g., `NextStep11Pipeline`, `StableDiffusionPipeline`). Check this set before attempting to enable cache-dit.
|
||||
|
||||
## Reference Implementations
|
||||
|
||||
| Model | Path | Notes |
|
||||
|-------|------|-------|
|
||||
| Standard DiT | `cache_dit_backend.py::enable_cache_for_dit` | Default enabler, automatic |
|
||||
| Wan2.2 | `cache_dit_backend.py::enable_cache_for_wan22` | Dual-transformer, auto-detects mode |
|
||||
| LongCat | `cache_dit_backend.py::enable_cache_for_longcat_image` | Multi-block-list |
|
||||
| BAGEL | `cache_dit_backend.py::enable_cache_for_bagel` | Complex omni model |
|
||||
@@ -0,0 +1,273 @@
|
||||
# Custom Model Patterns Reference
|
||||
|
||||
Patterns for adding models that don't come from the standard diffusers pipeline format.
|
||||
|
||||
## Directory Structure Comparison
|
||||
|
||||
### Diffusers-based model (e.g., Wan2.2)
|
||||
|
||||
```
|
||||
vllm_omni/diffusion/models/wan2_2/
|
||||
├── __init__.py # Exports pipeline + transformer + helpers
|
||||
├── pipeline_wan2_2.py # Pipeline: loads components via from_pretrained()
|
||||
├── pipeline_wan2_2_i2v.py # Variant pipeline for image-to-video
|
||||
└── wan2_2_transformer.py # Transformer: ported from diffusers, uses Attention layer
|
||||
```
|
||||
|
||||
The transformer is loaded separately via `weights_sources` + `load_weights()`. Non-transformer components (VAE, text encoder) are loaded in `__init__` via `from_pretrained()`.
|
||||
|
||||
### Custom model with external deps (e.g., DreamID-Omni)
|
||||
|
||||
```
|
||||
vllm_omni/diffusion/models/dreamid_omni/
|
||||
├── __init__.py # Exports pipeline only
|
||||
├── pipeline_dreamid_omni.py # Pipeline: loads ALL weights in __init__ via custom helpers
|
||||
├── fusion.py # Custom fusion architecture (video + audio cross-attention)
|
||||
└── wan2_2.py # Re-implemented Wan backbone with split API
|
||||
|
||||
examples/offline_inference/x_to_video_audio/
|
||||
└── download_dreamid_omni.py # Downloads weights from 3 HF repos + clones code repo
|
||||
```
|
||||
|
||||
All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency (`dreamid_omni` package) imported with try/except.
|
||||
|
||||
### Custom model with ported code (e.g., BAGEL)
|
||||
|
||||
```
|
||||
vllm_omni/diffusion/models/bagel/
|
||||
├── __init__.py
|
||||
├── pipeline_bagel.py # Pipeline: instantiates models, uses weights_sources
|
||||
├── bagel_transformer.py # Full LLM backbone (Qwen2-MoT) ported into vllm-omni
|
||||
└── autoencoder.py # Custom VAE ported from original repo
|
||||
```
|
||||
|
||||
Model code is fully ported (no external dependency). Uses `weights_sources` and `load_weights()` with custom name remapping to handle non-diffusers safetensors format.
|
||||
|
||||
## Weight Loading Patterns
|
||||
|
||||
### Pattern 1: Standard diffusers flow (Wan2.2, Z-Image, FLUX)
|
||||
|
||||
```
|
||||
init → create transformer (empty) → set weights_sources → [loader calls load_weights()]
|
||||
```
|
||||
|
||||
- `weights_sources` points to safetensors in HF subfolder (e.g., `transformer/`)
|
||||
- `load_weights()` receives `(name, tensor)` pairs from the loader
|
||||
- Name remapping handles diffusers→vllm-omni differences (QKV fusion, Sequential index removal)
|
||||
|
||||
### Pattern 2: Custom safetensors at root (BAGEL)
|
||||
|
||||
```
|
||||
init → create all models (empty) → set weights_sources(subfolder=None) → [loader calls load_weights()]
|
||||
```
|
||||
|
||||
- `weights_sources` points to **root** of model directory, not a subfolder
|
||||
- Weights have non-diffusers names (e.g., `bagel.language_model.model.layers.0.self_attn.q_proj.weight`)
|
||||
- `load_weights()` does heavy name normalization
|
||||
|
||||
```python
|
||||
self.weights_sources = [
|
||||
DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path=od_config.model,
|
||||
subfolder=None, # root directory
|
||||
prefix="", # no prefix stripping
|
||||
fall_back_to_pt=False,
|
||||
)
|
||||
]
|
||||
```
|
||||
|
||||
### Pattern 3: Fully custom loading (DreamID-Omni)
|
||||
|
||||
```
|
||||
init → load ALL weights eagerly via custom helpers → load_weights() = no-op
|
||||
```
|
||||
|
||||
- No `weights_sources` attribute — standard loader finds nothing to iterate
|
||||
- Custom init functions (e.g., `init_wan_vae_2_2()`, `load_fusion_checkpoint()`) handle downloading and loading
|
||||
- `load_weights()` is `pass`
|
||||
- Weights may come from multiple HF repos in different formats (`.pth`, `.safetensors`)
|
||||
|
||||
Use this when:
|
||||
- The original model has complex, well-tested loading code you don't want to rewrite
|
||||
- Weights span multiple HF repos
|
||||
- Weight format is non-standard (e.g., a single `.pth` file, not sharded safetensors)
|
||||
|
||||
## model_index.json for Custom Models
|
||||
|
||||
Standard diffusers `model_index.json`:
|
||||
```json
|
||||
{
|
||||
"_class_name": "WanPipeline",
|
||||
"_diffusers_version": "0.35.0.dev0",
|
||||
"scheduler": ["diffusers", "UniPCMultistepScheduler"],
|
||||
"transformer": ["diffusers", "WanTransformer3DModel"],
|
||||
"vae": ["diffusers", "AutoencoderKLWan"]
|
||||
}
|
||||
```
|
||||
|
||||
Custom model `model_index.json` (minimal):
|
||||
```json
|
||||
{
|
||||
"_class_name": "DreamIDOmniPipeline",
|
||||
"fusion": "DreamID-Omni/dreamid_omni.safetensors"
|
||||
}
|
||||
```
|
||||
|
||||
The only **required** field is `_class_name` — it must match a key in `_DIFFUSION_MODELS` in `registry.py`. Other fields are model-specific and accessible via `od_config.model_config` dict.
|
||||
|
||||
## External Dependency Management
|
||||
|
||||
### Git clone + .pth injection (DreamID-Omni pattern)
|
||||
|
||||
```python
|
||||
def download_dependency():
|
||||
CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
with open(LOCK_FILE, "w") as f:
|
||||
fcntl.flock(f, fcntl.LOCK_EX)
|
||||
if not DEPENDENCY_DIR.exists():
|
||||
subprocess.run([
|
||||
"git", "clone", "--depth", "1",
|
||||
REPO_URL, "--branch", BRANCH,
|
||||
str(DEPENDENCY_DIR)
|
||||
], check=True)
|
||||
fcntl.flock(f, fcntl.LOCK_UN)
|
||||
|
||||
# Add to Python path via .pth file
|
||||
site_packages = Path(site.getsitepackages()[0])
|
||||
pth_file = site_packages / "vllm_omni_dependency.pth"
|
||||
pth_file.write_text(str(DEPENDENCY_DIR))
|
||||
```
|
||||
|
||||
### Direct port (BAGEL pattern)
|
||||
|
||||
Copy essential files from the original repo into `vllm_omni/diffusion/models/<name>/`. Adapt imports to use vllm-omni utilities. Benefits: no external dependency, no git clone step. Drawback: must maintain the ported code.
|
||||
|
||||
## Multi-Modal Input/Output Protocols
|
||||
|
||||
Custom models that handle images, audio, or video I/O should implement protocol classes:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.models.interface import (
|
||||
SupportImageInput, # Model accepts image input
|
||||
SupportAudioInput, # Model accepts audio input
|
||||
SupportAudioOutput, # Model produces audio output
|
||||
)
|
||||
|
||||
class MyPipeline(nn.Module, SupportImageInput, SupportAudioInput, SupportAudioOutput):
|
||||
pass # Protocol markers enable proper engine routing
|
||||
```
|
||||
|
||||
The engine checks `isinstance(pipeline, SupportImageInput)` at startup to configure input validation and warmup behavior.
|
||||
|
||||
## Hardcoded Config vs Config Files
|
||||
|
||||
Diffusers models use `config.json` in each subfolder. Custom models often use:
|
||||
|
||||
**Module-level config dicts** (DreamID-Omni):
|
||||
```python
|
||||
VIDEO_CONFIG = {
|
||||
"patch_size": [1, 2, 2], "model_type": "ti2v",
|
||||
"dim": 3072, "ffn_dim": 14336, "num_heads": 24, "num_layers": 30, ...
|
||||
}
|
||||
```
|
||||
|
||||
**Loaded from custom JSON** (BAGEL):
|
||||
```python
|
||||
cfg_path = os.path.join(model_path, "config.json")
|
||||
with open(cfg_path) as f:
|
||||
bagel_cfg = json.load(f)
|
||||
vae_cfg = bagel_cfg.get("vae_config", {})
|
||||
```
|
||||
|
||||
## Custom Architecture Patterns
|
||||
|
||||
### Split forward API (DreamID-Omni)
|
||||
|
||||
When a fusion model needs to interleave blocks from two backbones:
|
||||
|
||||
```python
|
||||
class WanModel(nn.Module):
|
||||
def prepare_transformer_block_kwargs(self, x, t, context, ...):
|
||||
# Patch embed, time embed, text embed, RoPE
|
||||
return x, e, kwargs
|
||||
|
||||
def post_transformer_block_out(self, x, grid_sizes, e):
|
||||
# Output projection, unpatchify
|
||||
return output
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
raise NotImplementedError # Fusion model handles block iteration
|
||||
```
|
||||
|
||||
The `FusionModel` then iterates blocks in lock-step:
|
||||
```python
|
||||
for video_block, audio_block in zip(self.video_model.blocks, self.audio_model.blocks):
|
||||
video_out = video_block(video_hidden, ...)
|
||||
audio_out = audio_block(audio_hidden, ...)
|
||||
# Cross-attend between modalities
|
||||
video_out = cross_attention(video_out, audio_out)
|
||||
audio_out = cross_attention(audio_out, video_out)
|
||||
```
|
||||
|
||||
### LLM-as-denoiser (BAGEL)
|
||||
|
||||
When the backbone is a language model that also does diffusion:
|
||||
|
||||
```python
|
||||
class BagelModel(nn.Module):
|
||||
def __init__(self):
|
||||
self.language_model = Qwen2MoTForCausalLM(config)
|
||||
self.vit_model = SiglipVisionModel(vit_config)
|
||||
```
|
||||
|
||||
The LLM processes both text tokens and latent image tokens in a single forward pass, using KV caching for the text portion.
|
||||
|
||||
## Pre/Post Processing for Custom Models
|
||||
|
||||
Custom models typically handle pre/post processing **inside `forward()`** rather than via registered functions, because the logic is tightly coupled:
|
||||
|
||||
```python
|
||||
def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
|
||||
# Inline preprocessing
|
||||
image = self._load_and_resize_image(req.prompts[0].get("multi_modal_data", {}).get("image"))
|
||||
image_latent = self._vae_encode(image)
|
||||
|
||||
# ... denoising loop ...
|
||||
|
||||
# Inline postprocessing
|
||||
pil_image = self._decode_to_pil(latents)
|
||||
return DiffusionOutput(output=[pil_image])
|
||||
```
|
||||
|
||||
If pre/post functions are not registered in `_DIFFUSION_PRE_PROCESS_FUNCS` / `_DIFFUSION_POST_PROCESS_FUNCS`, the engine simply skips those steps.
|
||||
|
||||
## Download Script Template
|
||||
|
||||
```python
|
||||
# examples/offline_inference/<name>/download_<name>.py
|
||||
from huggingface_hub import snapshot_download
|
||||
import json, os
|
||||
|
||||
def main(output_dir):
|
||||
# Download model weights from HF
|
||||
snapshot_download(repo_id="org/model-weights", local_dir=os.path.join(output_dir, "weights"))
|
||||
|
||||
# Download additional components if from separate repos
|
||||
snapshot_download(repo_id="org/vae-weights", local_dir=os.path.join(output_dir, "vae"),
|
||||
allow_patterns=["*.safetensors"])
|
||||
|
||||
# Generate model_index.json
|
||||
config = {"_class_name": "YourPipeline", "custom_key": "weights/model.safetensors"}
|
||||
with open(os.path.join(output_dir, "model_index.json"), "w") as f:
|
||||
json.dump(config, f, indent=2)
|
||||
|
||||
# Install external code dependency (if needed)
|
||||
download_dependency()
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--output-dir", default="./your_model")
|
||||
args = parser.parse_args()
|
||||
main(args.output_dir)
|
||||
```
|
||||
@@ -0,0 +1,571 @@
|
||||
# Parallelism Patterns Reference
|
||||
|
||||
## Overview
|
||||
|
||||
vLLM-Omni supports multiple parallelism strategies for diffusion models. Each targets a different bottleneck:
|
||||
|
||||
| Strategy | Splits | Best For | Constraint |
|
||||
|----------|--------|----------|------------|
|
||||
| Tensor Parallel (TP) | Model layers across GPUs | Latency reduction, large models | Requires fast GPU interconnect, `num_heads % tp == 0` |
|
||||
| Sequence Parallel (SP/USP) | Sequence tokens across GPUs | Long sequences (video, high-res) | Near-linear scaling |
|
||||
| CFG Parallel | Positive/negative CFG branches | Models using classifier-free guidance | Exactly 2 GPUs |
|
||||
| HSDP | Weight shards via FSDP2 | VRAM reduction | Cannot combine with TP |
|
||||
| VAE Patch Parallel | VAE decode spatial tiles | Large VAE outputs | Auto-enables tiling |
|
||||
|
||||
**Recommended integration order**: TP → SP → CFG Parallel → HSDP
|
||||
|
||||
**Official design docs**:
|
||||
- TP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/tensor_parallel
|
||||
- SP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/sequence_parallel
|
||||
- CFG: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cfg_parallel
|
||||
- HSDP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/hsdp
|
||||
|
||||
---
|
||||
|
||||
## Tensor Parallelism (TP)
|
||||
|
||||
Replace standard `nn.Linear` with vLLM's parallel linear layers. This is the most invasive change but provides direct VRAM savings and compute speedup.
|
||||
|
||||
### Layer replacement rules
|
||||
|
||||
| Pattern | vLLM Layer | When to Use |
|
||||
|---------|-----------|-------------|
|
||||
| Fan-out (first in FFN) | `ColumnParallelLinear` | Projection that splits output across ranks |
|
||||
| Fan-in (second in FFN) | `RowParallelLinear` | Projection that gathers across ranks |
|
||||
| QKV projection | `QKVParallelLinear` | Fused Q/K/V for self-attention |
|
||||
| Single Q or K or V | `ColumnParallelLinear` | Separate projections (cross-attention) |
|
||||
| Attention output | `RowParallelLinear` | Output projection after attention |
|
||||
| Must not shard | `ReplicatedLinear` | Layers that must stay replicated |
|
||||
|
||||
### MLP Block (Up-Down Pattern)
|
||||
|
||||
```python
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear, RowParallelLinear,
|
||||
)
|
||||
|
||||
class TPFeedForward(nn.Module):
|
||||
def __init__(self, dim, ffn_dim):
|
||||
super().__init__()
|
||||
self.fc1 = ColumnParallelLinear(dim, ffn_dim, bias=False, return_bias=False)
|
||||
self.fc2 = RowParallelLinear(
|
||||
ffn_dim, dim, bias=False,
|
||||
input_is_parallel=True, # Input already sharded from fc1
|
||||
return_bias=False,
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x, _ = self.fc1(x)
|
||||
x = torch.nn.functional.gelu(x)
|
||||
x, _ = self.fc2(x)
|
||||
return x
|
||||
```
|
||||
|
||||
### Attention Block (QKV-Out Pattern)
|
||||
|
||||
```python
|
||||
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
|
||||
class TPSelfAttention(nn.Module):
|
||||
def __init__(self, dim, num_heads, num_kv_heads=None):
|
||||
super().__init__()
|
||||
num_kv_heads = num_kv_heads or num_heads
|
||||
self.head_dim = dim // num_heads
|
||||
|
||||
self.to_qkv = QKVParallelLinear(
|
||||
hidden_size=dim,
|
||||
head_size=self.head_dim,
|
||||
total_num_heads=num_heads,
|
||||
total_num_kv_heads=num_kv_heads,
|
||||
bias=False,
|
||||
return_bias=False,
|
||||
)
|
||||
self.to_out = RowParallelLinear(
|
||||
dim, dim, bias=False,
|
||||
input_is_parallel=True,
|
||||
return_bias=False,
|
||||
)
|
||||
self.attn = Attention(
|
||||
num_heads=self.to_qkv.num_heads, # Local heads per GPU
|
||||
head_size=self.head_dim,
|
||||
softmax_scale=1.0 / (self.head_dim ** 0.5),
|
||||
causal=False,
|
||||
num_kv_heads=self.to_qkv.num_kv_heads, # Local KV heads per GPU
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
qkv, _ = self.to_qkv(x)
|
||||
q, k, v = qkv.split(
|
||||
[self.to_qkv.num_heads * self.head_dim,
|
||||
self.to_qkv.num_kv_heads * self.head_dim,
|
||||
self.to_qkv.num_kv_heads * self.head_dim],
|
||||
dim=-1,
|
||||
)
|
||||
B, S, _ = x.shape
|
||||
q = q.view(B, S, self.to_qkv.num_heads, self.head_dim)
|
||||
k = k.view(B, S, self.to_qkv.num_kv_heads, self.head_dim)
|
||||
v = v.view(B, S, self.to_qkv.num_kv_heads, self.head_dim)
|
||||
out = self.attn(q, k, v)
|
||||
out = out.reshape(B, S, -1)
|
||||
out, _ = self.to_out(out)
|
||||
return out
|
||||
```
|
||||
|
||||
### QKV Fusion in load_weights
|
||||
|
||||
When you fuse separate Q/K/V into `QKVParallelLinear`, map diffusers' separate weight names:
|
||||
|
||||
```python
|
||||
stacked_params_mapping = [
|
||||
("to_qkv", "to_q", "q"),
|
||||
("to_qkv", "to_k", "k"),
|
||||
("to_qkv", "to_v", "v"),
|
||||
]
|
||||
|
||||
def load_weights(self, weights):
|
||||
params = dict(self.named_parameters())
|
||||
loaded = set()
|
||||
for name, tensor in weights:
|
||||
for fused_name, orig_name, shard_id in stacked_params_mapping:
|
||||
if orig_name in name:
|
||||
name = name.replace(orig_name, fused_name)
|
||||
param = params[name]
|
||||
param.weight_loader(param, tensor, shard_id)
|
||||
loaded.add(name)
|
||||
break
|
||||
else:
|
||||
if name in params:
|
||||
param = params[name]
|
||||
if hasattr(param, "weight_loader"):
|
||||
param.weight_loader(param, tensor)
|
||||
else:
|
||||
default_weight_loader(param, tensor)
|
||||
loaded.add(name)
|
||||
return loaded
|
||||
```
|
||||
|
||||
### RMSNorm with TP
|
||||
|
||||
When RMSNorm sits between TP-sharded dimensions, use `DistributedRMSNorm` — it computes global RMS via all-reduce across TP ranks. See the Wan2.2 implementation for the pattern.
|
||||
|
||||
### TP Constraints
|
||||
|
||||
- `num_heads % tp_size == 0`
|
||||
- `num_kv_heads % tp_size == 0`
|
||||
- Use `self.to_qkv.num_heads` (local per-GPU count), not total heads, for split sizes
|
||||
|
||||
### Testing TP
|
||||
|
||||
```bash
|
||||
python text_to_image.py --model Your-org/your-model \
|
||||
--tensor-parallel-size 2 --output "tp_test.png"
|
||||
```
|
||||
|
||||
**Verify**: speedup, memory reduction proportional to TP size, quality matches single-GPU.
|
||||
|
||||
### Reference implementations
|
||||
|
||||
| Model | Path |
|
||||
|-------|------|
|
||||
| Z-Image | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` |
|
||||
| FLUX | `vllm_omni/diffusion/models/flux/flux_transformer.py` |
|
||||
| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` |
|
||||
|
||||
---
|
||||
|
||||
## Sequence Parallelism (SP / USP)
|
||||
|
||||
SP splits sequence tokens across GPUs using Ulysses (all-to-all) or Ring (P2P) communication. It is applied non-intrusively via the `_sp_plan` dict — no changes to `forward()` logic.
|
||||
|
||||
### Approach 1: Non-Intrusive `_sp_plan` (Recommended)
|
||||
|
||||
The framework automatically registers hooks to shard inputs and gather outputs at `nn.Module` boundaries.
|
||||
|
||||
#### Step 1: Identify module boundaries
|
||||
|
||||
Find where tensors need sharding/gathering:
|
||||
|
||||
```python
|
||||
class MyTransformer(nn.Module):
|
||||
def __init__(self):
|
||||
self.patch_embed = PatchEmbed() # Before blocks
|
||||
self.pos_embed = RoPE() # RoPE may need splitting
|
||||
self.blocks = nn.ModuleList([...]) # Blocks process sharded x
|
||||
self.norm_out = LayerNorm()
|
||||
self.proj_out = Linear() # Gather after this
|
||||
|
||||
def forward(self, x):
|
||||
x = self.patch_embed(x)
|
||||
pos = self.pos_embed(x)
|
||||
for block in self.blocks:
|
||||
x = block(x, pos)
|
||||
x = self.norm_out(x)
|
||||
return self.proj_out(x)
|
||||
```
|
||||
|
||||
#### Step 2: Handle inline operations
|
||||
|
||||
`_sp_plan` hooks only work at `nn.Module` boundaries. Inline ops like `torch.cat()` must be extracted into submodules:
|
||||
|
||||
```python
|
||||
# BAD: Inline — hooks can't intercept
|
||||
unified = torch.cat([x, cap_feats], dim=1)
|
||||
|
||||
# GOOD: Extract into submodule
|
||||
class UnifiedPrepare(nn.Module):
|
||||
def forward(self, x, cap_feats):
|
||||
return torch.cat([x, cap_feats], dim=1)
|
||||
|
||||
self.unified_prepare = UnifiedPrepare()
|
||||
unified = self.unified_prepare(x, cap_feats)
|
||||
```
|
||||
|
||||
Common cases: `torch.cat()`, `pad_sequence()`, `tensor.reshape()`, complex preprocessing.
|
||||
|
||||
#### Step 3: Write `_sp_plan`
|
||||
|
||||
**Pattern 1: Shard at first block, gather at output** (most common)
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.distributed.sp_plan import (
|
||||
SequenceParallelInput, SequenceParallelOutput,
|
||||
)
|
||||
|
||||
class StandardTransformer(nn.Module):
|
||||
_sp_plan = {
|
||||
"blocks.0": {
|
||||
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
|
||||
},
|
||||
"proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
|
||||
}
|
||||
```
|
||||
|
||||
**Pattern 2: Shard RoPE outputs separately**
|
||||
|
||||
```python
|
||||
class TransformerWithRoPE(nn.Module):
|
||||
_sp_plan = {
|
||||
"rope": {
|
||||
0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
|
||||
1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
|
||||
},
|
||||
"blocks.0": {
|
||||
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
|
||||
},
|
||||
"proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
|
||||
}
|
||||
```
|
||||
|
||||
**Pattern 3: Dual-stream (shard image, replicate text)**
|
||||
|
||||
```python
|
||||
class DualStreamTransformer(nn.Module):
|
||||
_sp_plan = {
|
||||
"rope_preparer": {
|
||||
2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
|
||||
3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
|
||||
},
|
||||
"transformer_blocks.0": {
|
||||
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
|
||||
},
|
||||
"proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
|
||||
}
|
||||
```
|
||||
|
||||
### API Reference
|
||||
|
||||
**SequenceParallelInput**:
|
||||
|
||||
| Parameter | Type | Description |
|
||||
|-----------|------|-------------|
|
||||
| `split_dim` | int | Dimension to split (usually 1 for sequence) |
|
||||
| `expected_dims` | int/None | Expected tensor rank for validation |
|
||||
| `split_output` | bool | `False`: shard input params; `True`: shard output tensors |
|
||||
| `auto_pad` | bool | Auto-pad if sequence not divisible by world_size |
|
||||
|
||||
**SequenceParallelOutput**:
|
||||
|
||||
| Parameter | Type | Description |
|
||||
|-----------|------|-------------|
|
||||
| `gather_dim` | int | Dimension to gather (usually 1 for sequence) |
|
||||
| `expected_dims` | int/None | Expected tensor rank for validation |
|
||||
|
||||
**Module naming**:
|
||||
|
||||
| Key | Meaning |
|
||||
|-----|---------|
|
||||
| `"blocks.0"` | First element of ModuleList |
|
||||
| `"blocks.*"` | All elements of ModuleList |
|
||||
| `"rope"` | Named submodule |
|
||||
|
||||
**Dictionary value types**:
|
||||
|
||||
| Key type | split_output | Description |
|
||||
|----------|-------------|-------------|
|
||||
| `"param_name"` (str) | False | Shard input parameter by name |
|
||||
| `0, 1, ...` (int) | True | Shard output tuple by index |
|
||||
|
||||
### Approach 2: Intrusive Modification (Complex Cases)
|
||||
|
||||
For dynamic sharding logic that can't be expressed via `_sp_plan`:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
|
||||
|
||||
def forward(self, hidden_states, ...):
|
||||
if self.parallel_config.sequence_parallel_size > 1:
|
||||
hidden_states = sp_shard(hidden_states, dim=1)
|
||||
for block in self.blocks:
|
||||
hidden_states = block(hidden_states)
|
||||
if self.parallel_config.sequence_parallel_size > 1:
|
||||
hidden_states = sp_gather(hidden_states, dim=1)
|
||||
return hidden_states
|
||||
```
|
||||
|
||||
Use intrusive modification as a last resort — `_sp_plan` is preferred for maintainability.
|
||||
|
||||
### UAA Mode (Experimental)
|
||||
|
||||
`ulysses_mode="advanced_uaa"` handles arbitrary sequence lengths and head counts that aren't divisible by `ulysses_degree`. Uses variable all-to-all split sizes and temporary head padding.
|
||||
|
||||
### Combining SP methods
|
||||
|
||||
Ulysses and Ring can be combined: `ulysses_degree × ring_degree = total SP GPUs`.
|
||||
|
||||
```python
|
||||
DiffusionParallelConfig(ulysses_degree=2, ring_degree=2) # 4 GPUs total
|
||||
```
|
||||
|
||||
### Testing SP
|
||||
|
||||
```bash
|
||||
# Offline
|
||||
python text_to_image.py --model Your-model --ulysses-degree 2
|
||||
|
||||
# Online serving
|
||||
vllm serve Your-model --omni --usp 2
|
||||
```
|
||||
|
||||
### Reference implementations
|
||||
|
||||
| Model | Path |
|
||||
|-------|------|
|
||||
| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` |
|
||||
| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` |
|
||||
| Z-Image | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` |
|
||||
|
||||
---
|
||||
|
||||
## CFG Parallelism
|
||||
|
||||
Distributes positive/negative Classifier-Free Guidance branches across 2 GPUs.
|
||||
|
||||
### Implementation
|
||||
|
||||
Inherit `CFGParallelMixin` and implement `diffuse()`:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
|
||||
class YourPipeline(nn.Module, CFGParallelMixin):
|
||||
def diffuse(self, latents, timesteps, prompt_embeds, negative_embeds,
|
||||
do_true_cfg, true_cfg_scale, **kwargs):
|
||||
for i, t in enumerate(timesteps):
|
||||
positive_kwargs = {
|
||||
"hidden_states": latents,
|
||||
"encoder_hidden_states": prompt_embeds,
|
||||
"timestep": t,
|
||||
}
|
||||
negative_kwargs = {
|
||||
"hidden_states": latents,
|
||||
"encoder_hidden_states": negative_embeds,
|
||||
"timestep": t,
|
||||
} if do_true_cfg else None
|
||||
|
||||
noise_pred = self.predict_noise_maybe_with_cfg(
|
||||
do_true_cfg=do_true_cfg,
|
||||
true_cfg_scale=true_cfg_scale,
|
||||
positive_kwargs=positive_kwargs,
|
||||
negative_kwargs=negative_kwargs,
|
||||
)
|
||||
latents = self.scheduler_step_maybe_with_cfg(
|
||||
noise_pred, t, latents, do_true_cfg
|
||||
)
|
||||
return latents
|
||||
```
|
||||
|
||||
### Customization hooks
|
||||
|
||||
| Method | Override when |
|
||||
|--------|-------------|
|
||||
| `predict_noise()` | Non-standard transformer call (e.g., dual-transformer like Wan2.2) |
|
||||
| `cfg_normalize_function()` | Custom normalization (e.g., LongCat with clamping) |
|
||||
| `combine_cfg_noise()` | Multi-output models (e.g., video + audio: CFG on video, positive-only on audio) |
|
||||
|
||||
**Custom predict_noise** (Wan2.2 — selects active transformer):
|
||||
|
||||
```python
|
||||
def predict_noise(self, current_model=None, **kwargs):
|
||||
if current_model is None:
|
||||
current_model = self.transformer
|
||||
return current_model(**kwargs)[0]
|
||||
```
|
||||
|
||||
**Custom combine_cfg_noise** (multi-output):
|
||||
|
||||
```python
|
||||
def combine_cfg_noise(self, positive_pred, negative_pred, scale, normalize):
|
||||
video_pos, audio_pos = positive_pred
|
||||
video_neg, audio_neg = negative_pred
|
||||
video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize)
|
||||
return (video_combined, audio_pos)
|
||||
```
|
||||
|
||||
### Composite scheduler for multi-output
|
||||
|
||||
When each output has its own schedule:
|
||||
|
||||
```python
|
||||
class VideoAudioScheduler:
|
||||
def __init__(self, video_scheduler, audio_scheduler):
|
||||
self.video_scheduler = video_scheduler
|
||||
self.audio_scheduler = audio_scheduler
|
||||
|
||||
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
|
||||
video_out = self.video_scheduler.step(
|
||||
noise_pred[0], t[0], latents[0], return_dict=False, generator=generator
|
||||
)[0]
|
||||
audio_out = self.audio_scheduler.step(
|
||||
noise_pred[1], t[1], latents[1], return_dict=False, generator=generator
|
||||
)[0]
|
||||
return ((video_out, audio_out),)
|
||||
```
|
||||
|
||||
### Testing CFG Parallel
|
||||
|
||||
```bash
|
||||
python text_to_image.py --model Your-model \
|
||||
--cfg-parallel-size 2 --cfg-scale 4.0 \
|
||||
--negative-prompt "ugly, unclear"
|
||||
```
|
||||
|
||||
**Constraint**: `guidance_scale > 1.0` and negative prompt must be provided.
|
||||
|
||||
### Reference implementations
|
||||
|
||||
| Model | Path |
|
||||
|-------|------|
|
||||
| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/cfg_parallel.py` |
|
||||
| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py` |
|
||||
| Mixin base | `vllm_omni/diffusion/distributed/cfg_parallel.py` |
|
||||
|
||||
---
|
||||
|
||||
## HSDP (Hybrid Sharded Data Parallel)
|
||||
|
||||
Shards model weights across GPUs using PyTorch FSDP2. Reduces per-GPU VRAM without changing computation.
|
||||
|
||||
### Implementation
|
||||
|
||||
Add `_hsdp_shard_conditions` to the transformer class:
|
||||
|
||||
```python
|
||||
class YourTransformer(nn.Module):
|
||||
@staticmethod
|
||||
def _is_transformer_block(name: str, module) -> bool:
|
||||
return "blocks" in name and name.split(".")[-1].isdigit()
|
||||
|
||||
_hsdp_shard_conditions = [_is_transformer_block]
|
||||
```
|
||||
|
||||
For MoE models, add additional conditions:
|
||||
|
||||
```python
|
||||
class MoETransformer(nn.Module):
|
||||
@staticmethod
|
||||
def _is_transformer_block(name, module):
|
||||
return "blocks" in name and name.split(".")[-1].isdigit()
|
||||
|
||||
@staticmethod
|
||||
def _is_moe_expert(name, module):
|
||||
return "experts" in name and name.split(".")[-1].isdigit()
|
||||
|
||||
_hsdp_shard_conditions = [_is_transformer_block, _is_moe_expert]
|
||||
```
|
||||
|
||||
A module is sharded if **any** condition returns `True`.
|
||||
|
||||
### Constraints
|
||||
|
||||
- Cannot combine with Tensor Parallelism
|
||||
- For standalone HSDP (no other parallelism), `hsdp_shard_size` must be specified explicitly
|
||||
- Can combine with SP: HSDP reduces memory while SP distributes sequence
|
||||
|
||||
### Testing HSDP
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
|
||||
parallel_config = DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=8)
|
||||
omni = Omni(model="your-model", parallel_config=parallel_config)
|
||||
```
|
||||
|
||||
Or CLI:
|
||||
|
||||
```bash
|
||||
vllm serve Your-model --omni --use-hsdp
|
||||
```
|
||||
|
||||
**Verify**: logs show "HSDP Inference: replicate_size=..., shard_size=..." and "Sharded N modules + root". Check VRAM reduction.
|
||||
|
||||
### Reference implementations
|
||||
|
||||
| Model | Path |
|
||||
|-------|------|
|
||||
| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` |
|
||||
| HSDP Core | `vllm_omni/diffusion/distributed/hsdp.py` |
|
||||
|
||||
---
|
||||
|
||||
## VAE Patch Parallelism
|
||||
|
||||
Shards VAE decode spatially across ranks using tiling:
|
||||
|
||||
```bash
|
||||
python text_to_image.py --model Your-model --vae-patch-parallel-size 4
|
||||
```
|
||||
|
||||
Auto-enables `--vae-use-tiling`. Uses `DistributedAutoencoderKLWan` or similar distributed VAE. Set `vae_patch_parallel_size` in `DiffusionParallelConfig`.
|
||||
|
||||
---
|
||||
|
||||
## Combining Parallelism Methods
|
||||
|
||||
Common multi-GPU recipes:
|
||||
|
||||
```bash
|
||||
# 4 GPUs: CFG (2) × Ulysses (2)
|
||||
python text_to_image.py --model Qwen/Qwen-Image \
|
||||
--cfg-parallel-size 2 --ulysses-degree 2
|
||||
|
||||
# 8 GPUs: Ulysses (4) × Ring (2) + VAE patch (8)
|
||||
python text_to_video.py --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--ulysses-degree 4 --ring-degree 2 --vae-patch-parallel-size 8
|
||||
|
||||
# 2 GPUs: HSDP + Ulysses (cannot combine HSDP with TP)
|
||||
vllm serve Your-model --omni --use-hsdp --usp 2
|
||||
```
|
||||
|
||||
## Discovering Parallelism Support
|
||||
|
||||
Check which parallelism methods a model supports:
|
||||
|
||||
| Check | How |
|
||||
|-------|-----|
|
||||
| **Ulysses / Ring SP** | Transformer defines `_sp_plan`. Search: `grep -r '_sp_plan' vllm_omni/diffusion/models/` |
|
||||
| **CFG Parallel** | Pipeline inherits `CFGParallelMixin`. Search: `grep -r 'CFGParallelMixin' vllm_omni/diffusion/models/` |
|
||||
| **TP** | Uses `ColumnParallelLinear` / `QKVParallelLinear`. Search: `grep -r 'ParallelLinear\|QKVParallel' vllm_omni/diffusion/models/<model>/` |
|
||||
| **HSDP** | Transformer defines `_hsdp_shard_conditions`. Search: `grep -r '_hsdp_shard_conditions' vllm_omni/diffusion/models/` |
|
||||
|
||||
The canonical per-model support table is in `docs/user_guide/diffusion/parallelism_acceleration.md`.
|
||||
@@ -0,0 +1,218 @@
|
||||
# Transformer Adaptation Reference
|
||||
|
||||
## Adapting a Diffusers Transformer to vLLM-Omni
|
||||
|
||||
### Step-by-step Checklist
|
||||
|
||||
1. Copy the transformer class from diffusers source
|
||||
2. Remove all mixin classes — inherit only from `nn.Module`
|
||||
3. Replace attention dispatch with `vllm_omni.diffusion.attention.layer.Attention`
|
||||
4. Replace logger with `vllm.logger.init_logger`
|
||||
5. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`
|
||||
6. Remove training-only code (gradient checkpointing, dropout)
|
||||
7. Add `load_weights()` method for weight loading from safetensors
|
||||
8. Add class-level attributes for acceleration features
|
||||
|
||||
### Mixin Removal
|
||||
|
||||
Remove these diffusers mixins (and their imports):
|
||||
|
||||
```python
|
||||
# Remove all of these:
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.attention_processor import AttentionModuleMixin
|
||||
from diffusers.loaders import PeftAdapterMixin, FromOriginalModelMixin
|
||||
|
||||
# Replace:
|
||||
class MyTransformer(ModelMixin, ConfigMixin, AttentionModuleMixin):
|
||||
# With:
|
||||
class MyTransformer(nn.Module):
|
||||
```
|
||||
|
||||
Also remove `@register_to_config` decorators from `__init__`.
|
||||
|
||||
### Attention Replacement
|
||||
|
||||
The vLLM-Omni `Attention` layer wraps backend selection (FlashAttention, SDPA, SageAttn, etc.) and supports sequence parallelism hooks.
|
||||
|
||||
**QKV tensor shape must be `[batch, seq_len, num_heads, head_dim]`.**
|
||||
|
||||
#### Self-Attention Pattern
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
|
||||
class SelfAttentionBlock(nn.Module):
|
||||
def __init__(self, dim, num_heads):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = dim // num_heads
|
||||
|
||||
self.to_q = nn.Linear(dim, dim)
|
||||
self.to_k = nn.Linear(dim, dim)
|
||||
self.to_v = nn.Linear(dim, dim)
|
||||
self.to_out = nn.Linear(dim, dim)
|
||||
|
||||
self.attn = Attention(
|
||||
num_heads=num_heads,
|
||||
head_size=self.head_dim,
|
||||
softmax_scale=1.0 / (self.head_dim ** 0.5),
|
||||
causal=False,
|
||||
num_kv_heads=num_heads,
|
||||
)
|
||||
|
||||
def forward(self, x, attn_mask=None):
|
||||
B, S, _ = x.shape
|
||||
q = self.to_q(x).view(B, S, self.num_heads, self.head_dim)
|
||||
k = self.to_k(x).view(B, S, self.num_heads, self.head_dim)
|
||||
v = self.to_v(x).view(B, S, self.num_heads, self.head_dim)
|
||||
|
||||
attn_metadata = AttentionMetadata(attn_mask=attn_mask)
|
||||
out = self.attn(q, k, v, attn_metadata=attn_metadata)
|
||||
out = out.reshape(B, S, -1)
|
||||
return self.to_out(out)
|
||||
```
|
||||
|
||||
#### Fused QKV with TP (Advanced)
|
||||
|
||||
For tensor parallelism, use vLLM's parallel linear layers:
|
||||
|
||||
```python
|
||||
from vllm.model_executor.layers.linear import (
|
||||
QKVParallelLinear, RowParallelLinear
|
||||
)
|
||||
|
||||
class TPSelfAttention(nn.Module):
|
||||
def __init__(self, dim, num_heads):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = dim // num_heads
|
||||
|
||||
self.to_qkv = QKVParallelLinear(
|
||||
hidden_size=dim,
|
||||
head_size=self.head_dim,
|
||||
total_num_heads=num_heads,
|
||||
total_num_kv_heads=num_heads,
|
||||
)
|
||||
self.to_out = RowParallelLinear(dim, dim)
|
||||
|
||||
self.attn = Attention(
|
||||
num_heads=num_heads,
|
||||
head_size=self.head_dim,
|
||||
softmax_scale=1.0 / (self.head_dim ** 0.5),
|
||||
causal=False,
|
||||
num_kv_heads=num_heads,
|
||||
)
|
||||
```
|
||||
|
||||
### Logger Replacement
|
||||
|
||||
```python
|
||||
# Replace:
|
||||
from diffusers.utils import logging
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
# With:
|
||||
from vllm.logger import init_logger
|
||||
logger = init_logger(__name__)
|
||||
```
|
||||
|
||||
### Custom Layers from vLLM-Omni
|
||||
|
||||
Available utility layers:
|
||||
|
||||
```python
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
|
||||
from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
|
||||
```
|
||||
|
||||
### Config Support
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
class MyTransformer(nn.Module):
|
||||
def __init__(self, *, od_config=None, num_layers=28, hidden_size=3072, **kwargs):
|
||||
super().__init__()
|
||||
self.od_config = od_config
|
||||
self.parallel_config = od_config.parallel_config if od_config else None
|
||||
# ... build layers
|
||||
```
|
||||
|
||||
The transformer config values come from `model_index.json` → `config.json` in the transformer subfolder. The pipeline uses `get_transformer_config_kwargs(od_config.tf_model_config, TransformerClass)` to filter config keys to match the `__init__` signature.
|
||||
|
||||
### Weight Loading
|
||||
|
||||
The `load_weights` method receives an iterable of `(name, tensor)` from safetensors files, with the prefix (e.g., `"transformer."`) already stripped by the loader.
|
||||
|
||||
```python
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
class MyTransformer(nn.Module):
|
||||
def load_weights(self, weights):
|
||||
params = dict(self.named_parameters())
|
||||
loaded = set()
|
||||
for name, tensor in weights:
|
||||
# Optional: remap names from diffusers to vllm-omni naming
|
||||
# e.g., "ff.net.0.proj" -> "ff.net_0.proj"
|
||||
|
||||
if name in params:
|
||||
param = params[name]
|
||||
if hasattr(param, "weight_loader"):
|
||||
param.weight_loader(param, tensor)
|
||||
else:
|
||||
default_weight_loader(param, tensor)
|
||||
loaded.add(name)
|
||||
return loaded
|
||||
```
|
||||
|
||||
#### QKV Fusion in load_weights
|
||||
|
||||
If you fused separate Q/K/V into a `QKVParallelLinear`, you need to map diffusers' separate weight names:
|
||||
|
||||
```python
|
||||
stacked_params_mapping = [
|
||||
("to_qkv", "to_q", "q"),
|
||||
("to_qkv", "to_k", "k"),
|
||||
("to_qkv", "to_v", "v"),
|
||||
]
|
||||
|
||||
def load_weights(self, weights):
|
||||
params = dict(self.named_parameters())
|
||||
loaded = set()
|
||||
for name, tensor in weights:
|
||||
for fused_name, orig_name, shard_id in stacked_params_mapping:
|
||||
if orig_name in name:
|
||||
name = name.replace(orig_name, fused_name)
|
||||
param = params[name]
|
||||
param.weight_loader(param, tensor, shard_id)
|
||||
loaded.add(name)
|
||||
break
|
||||
else:
|
||||
# Normal loading
|
||||
...
|
||||
return loaded
|
||||
```
|
||||
|
||||
### Class-Level Attributes for Features
|
||||
|
||||
```python
|
||||
class MyTransformer(nn.Module):
|
||||
# torch.compile: list block class names that repeat and can be compiled
|
||||
_repeated_blocks = ["MyTransformerBlock"]
|
||||
|
||||
# CPU offload: attribute name of the nn.ModuleList containing blocks
|
||||
_layerwise_offload_blocks_attr = "blocks"
|
||||
|
||||
# LoRA: mapping of fused param names to original param names
|
||||
packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
|
||||
|
||||
# Sequence parallelism plan (advanced — add after basic impl works)
|
||||
_sp_plan = {
|
||||
"blocks.0": SequenceParallelInput(split_dim=1),
|
||||
"proj_out": SequenceParallelOutput(gather_dim=1),
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,178 @@
|
||||
# Troubleshooting Reference
|
||||
|
||||
## Common Errors When Adding a Diffusion Model
|
||||
|
||||
### ImportError / ModuleNotFoundError
|
||||
|
||||
**Cause**: Missing or incorrect registration.
|
||||
|
||||
**Fix checklist**:
|
||||
1. Model registered in `vllm_omni/diffusion/registry.py` `_DIFFUSION_MODELS` dict
|
||||
2. `__init__.py` exports the pipeline class
|
||||
3. Pipeline file exists at the correct path: `vllm_omni/diffusion/models/{folder}/{file}.py`
|
||||
4. Class name in registry matches the actual class name in the file
|
||||
|
||||
### Shape Mismatch in Attention
|
||||
|
||||
**Symptom**: `RuntimeError: shape mismatch` or `expected 4D tensor`
|
||||
|
||||
**Cause**: QKV tensors not reshaped to `[batch, seq_len, num_heads, head_dim]`.
|
||||
|
||||
**Fix**: Before calling `self.attn(q, k, v, ...)`, ensure:
|
||||
```python
|
||||
q = q.view(batch, seq_len, self.num_heads, self.head_dim)
|
||||
k = k.view(batch, kv_seq_len, self.num_kv_heads, self.head_dim)
|
||||
v = v.view(batch, kv_seq_len, self.num_kv_heads, self.head_dim)
|
||||
```
|
||||
|
||||
After attention, reshape back:
|
||||
```python
|
||||
out = out.reshape(batch, seq_len, -1)
|
||||
```
|
||||
|
||||
### Weight Loading Failures
|
||||
|
||||
**Symptom**: `RuntimeError: size mismatch for parameter ...` or missing keys
|
||||
|
||||
**Debugging**:
|
||||
1. Print diffusers weight names: `safetensors.safe_open(path, "pt").keys()`
|
||||
2. Print model parameter names: `dict(model.named_parameters()).keys()`
|
||||
3. Compare and add name remappings in `load_weights()`
|
||||
|
||||
**Common remappings needed**:
|
||||
- `ff.net.0.proj` → `ff.net_0.proj` (PyTorch Sequential indexing)
|
||||
- `.to_out.0.` → `.to_out.` (Sequential unwrapping)
|
||||
- `scale_shift_table` → moved to a wrapper module
|
||||
|
||||
### Black/Blank/Noisy Output
|
||||
|
||||
**Possible causes**:
|
||||
1. **Wrong latent normalization**: Check VAE expects latents scaled by `vae.config.scaling_factor`
|
||||
2. **Wrong scheduler**: Using the wrong scheduler class or wrong `flow_shift`
|
||||
3. **Missing CFG**: Some models require `guidance_scale > 1.0` with negative prompt
|
||||
4. **Wrong timestep format**: Some schedulers expect float, others expect int/long
|
||||
5. **Missing post-processing**: Raw VAE output may need denormalization
|
||||
|
||||
**Quick test**: Run with diffusers directly using the same seed and compare latents at each step.
|
||||
|
||||
### OOM (Out of Memory)
|
||||
|
||||
**Solutions** (in order of preference):
|
||||
1. `--enforce-eager` to disable torch.compile (saves compile memory)
|
||||
2. `--enable-cpu-offload` for model-level offload
|
||||
3. `--enable-layerwise-offload` for block-level offload (better for large models)
|
||||
4. `--vae-use-slicing --vae-use-tiling` for VAE memory reduction
|
||||
5. Reduce resolution: `--height 480 --width 832`
|
||||
6. Use TP: `--tensor-parallel-size 2`
|
||||
|
||||
### Different Output vs Diffusers Reference
|
||||
|
||||
**Common causes**:
|
||||
1. **Attention backend difference**: FlashAttention vs SDPA may produce slightly different results. Set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` to match diffusers
|
||||
2. **Float precision**: vLLM-Omni may use bfloat16 where diffusers uses float32 for some operations
|
||||
3. **Missing normalization**: Check all LayerNorm/RMSNorm are preserved
|
||||
4. **Scheduler rounding**: Some schedulers have numerical sensitivity
|
||||
|
||||
### Tensor Parallel Errors
|
||||
|
||||
**Symptom**: `AssertionError: not divisible` or incorrect output with TP>1
|
||||
|
||||
**Fix**:
|
||||
1. Verify `num_heads % tp_size == 0` and `num_kv_heads % tp_size == 0`
|
||||
2. Ensure `ColumnParallelLinear` / `RowParallelLinear` are used correctly
|
||||
3. Check that norms between parallel layers use distributed norm if needed
|
||||
4. Verify `load_weights` handles TP sharding for norm weights
|
||||
5. Use `self.to_qkv.num_heads` (local heads per GPU) for QKV split sizes, not total heads
|
||||
|
||||
**Missing `input_is_parallel=True`**:
|
||||
|
||||
`RowParallelLinear` expects sharded input from `ColumnParallelLinear`:
|
||||
```python
|
||||
self.w1 = ColumnParallelLinear(dim, hidden_dim, return_bias=False)
|
||||
self.w2 = RowParallelLinear(hidden_dim, dim, input_is_parallel=True, return_bias=False)
|
||||
```
|
||||
|
||||
### Sequence Parallel Errors
|
||||
|
||||
**Symptom**: Incorrect output or crashes with `--ulysses-degree N` or `--usp N`
|
||||
|
||||
**Possible causes**:
|
||||
1. **Inline operations between shard/gather points**: `torch.cat()`, `pad_sequence()` etc. not at `nn.Module` boundaries. Fix: extract into submodule.
|
||||
2. **Wrong `split_dim`**: Check the tensor shape at the shard point. Sequence dimension is typically `dim=1` for `[B, S, D]` tensors.
|
||||
3. **RoPE not sharded**: If RoPE is computed separately, add it to `_sp_plan` with `split_output=True`.
|
||||
4. **Sequence not divisible by SP degree**: Use `auto_pad=True` in `SequenceParallelInput` or switch to `ulysses_mode="advanced_uaa"`.
|
||||
|
||||
**Debugging**: Add `expected_dims=N` to `SequenceParallelInput`/`Output` for shape validation at runtime.
|
||||
|
||||
### CFG Parallel Errors
|
||||
|
||||
**Symptom**: CFG parallel not activating, no speedup
|
||||
|
||||
**Fix checklist**:
|
||||
1. Pipeline inherits `CFGParallelMixin`
|
||||
2. `guidance_scale > 1.0`
|
||||
3. Negative prompt provided (even if empty string)
|
||||
4. `--cfg-parallel-size 2` specified
|
||||
5. `diffuse()` method calls `predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()`
|
||||
|
||||
**Symptom**: Different output with CFG parallel vs sequential
|
||||
|
||||
**Possible cause**: Non-deterministic scheduler. Fix: pass `generator=torch.Generator(device).manual_seed(seed)` to `scheduler_step_maybe_with_cfg()`.
|
||||
|
||||
### HSDP Errors
|
||||
|
||||
**Symptom**: HSDP not activating or errors during weight loading
|
||||
|
||||
**Fix checklist**:
|
||||
1. Transformer defines `_hsdp_shard_conditions` class attribute
|
||||
2. Shard condition functions return `True` for correct modules (test with `model.named_modules()`)
|
||||
3. Not combining with TP (HSDP and TP are incompatible)
|
||||
4. For standalone HSDP, `hsdp_shard_size` is specified explicitly
|
||||
|
||||
**Verify**: Check logs for "HSDP Inference: replicate_size=..., shard_size=..." and "Sharded N modules + root".
|
||||
|
||||
### Cache-DiT Not Applied
|
||||
|
||||
**Symptom**: No speedup, no cache-related log messages
|
||||
|
||||
**Fix checklist**:
|
||||
1. Model not in `_NO_CACHE_ACCELERATION` in `registry.py`
|
||||
2. Pipeline class name matches `CUSTOM_DIT_ENABLERS` key (if using custom enabler)
|
||||
3. `cache_backend="cache_dit"` specified
|
||||
4. Check logs for "Cache-dit enabled successfully on xxx"
|
||||
|
||||
**Verify pipeline name**: `print(pipeline.__class__.__name__)` — must match registry key.
|
||||
|
||||
### Cache-DiT Quality Degradation
|
||||
|
||||
**Symptom**: Artifacts or lower quality with cache-dit
|
||||
|
||||
**Fix**: Reduce aggressiveness:
|
||||
```python
|
||||
cache_config={
|
||||
"residual_diff_threshold": 0.12, # Lower from 0.24
|
||||
"max_warmup_steps": 6, # Increase from 4
|
||||
"max_continuous_cached_steps": 2, # Reduce if higher
|
||||
}
|
||||
```
|
||||
|
||||
If quality is still poor, the model may need a custom enabler with per-block-list `ParamsModifier` tuning.
|
||||
|
||||
### Model Not Detected / Wrong Pipeline Class
|
||||
|
||||
**Symptom**: `ValueError: Model class ... not found in diffusion model registry`
|
||||
|
||||
**Cause**: The model's `model_index.json` has a `_class_name` for the pipeline that doesn't match registry keys.
|
||||
|
||||
**Fix**: The registry key must match the diffusers pipeline class name from `model_index.json`. If using a different name, map it in the registry:
|
||||
```python
|
||||
"DiffusersPipelineClassName": ("your_folder", "your_file", "YourVllmClassName"),
|
||||
```
|
||||
|
||||
## Debugging Workflow
|
||||
|
||||
1. **Add verbose logging**: Use `logger.info()` to print tensor shapes at each stage
|
||||
2. **Compare step-by-step**: Run diffusers and vllm-omni side by side, comparing tensors after each major operation
|
||||
3. **Use small configs**: Reduce `num_inference_steps=2`, small resolution for fast iteration
|
||||
4. **Test transformer isolation**: Feed the same input to both diffusers and vllm-omni transformers, compare outputs
|
||||
5. **Binary search for bugs**: Comment out blocks/layers to isolate where divergence starts
|
||||
@@ -0,0 +1,284 @@
|
||||
---
|
||||
name: add-tts-model
|
||||
description: "Integrate a new text-to-speech model into vLLM-Omni from HuggingFace reference implementation through production-ready serving with streaming and CUDA graph acceleration. Use when adding a new TTS model, wiring stage separation for speech synthesis, enabling online voice generation serving, debugging TTS integration behavior, or building audio output pipelines."
|
||||
---
|
||||
|
||||
# TTS Model Integration Workflow
|
||||
|
||||
## Overview
|
||||
|
||||
```
|
||||
HF Reference -> Stage Separation -> Online Serving -> Async Chunk -> CUDA Graph
|
||||
(Phase 1) (Phase 2) (Phase 3) (Phase 4) (Phase 5)
|
||||
```
|
||||
|
||||
## Phase 1: HuggingFace Reference
|
||||
|
||||
**Goal**: Understand the reference implementation and verify it produces correct audio.
|
||||
|
||||
### Steps
|
||||
|
||||
1. **Run the reference model** end-to-end using the official HuggingFace / GitHub code
|
||||
2. **Document the architecture**:
|
||||
- What are the sub-models? (AR decoder, codec decoder, vocoder, etc.)
|
||||
- What is the token vocabulary? (semantic codes, RVQ codebooks, special tokens)
|
||||
- What is the output format? (sample rate, channels, codec type)
|
||||
3. **Capture reference outputs** for comparison during integration
|
||||
4. **Identify the config structure**: `config.json` fields, `model_type`, sub-model configs
|
||||
|
||||
### Key Questions
|
||||
|
||||
- How many codebooks? What are the codebook sizes?
|
||||
- What special tokens exist? (`<|voice|>`, `<|audio_start|>`, `<|im_end|>`, etc.)
|
||||
- What is the token-to-ID mapping for codec codes?
|
||||
- What is the hop length / frame rate of the codec?
|
||||
- Does the model support voice cloning? How? (reference audio encoding, speaker embeddings, etc.)
|
||||
|
||||
### Deliverables
|
||||
|
||||
- Working reference script that produces audio
|
||||
- Architecture diagram / notes
|
||||
- Token vocabulary mapping
|
||||
- Reference audio samples for regression testing
|
||||
|
||||
## Phase 2: Stage Separation (Offline Inference)
|
||||
|
||||
**Goal**: Split the model into vLLM-Omni stages and get offline inference working.
|
||||
|
||||
### Steps
|
||||
|
||||
1. **Register the model** in `vllm_omni/model_executor/models/registry.py`
|
||||
2. **Create config classes** (`configuration_<model>.py`) with `model_type` registration
|
||||
3. **Implement Stage 0** (AR model):
|
||||
- Subclass appropriate base (e.g., wrap Qwen3 decoder layers)
|
||||
- Implement `forward()` for autoregressive token generation
|
||||
- Handle special token logic (start/stop tokens, codec token mapping)
|
||||
- If dual-AR (like Fish Speech), implement Fast AR as a nested module
|
||||
4. **Implement Stage 1** (Decoder):
|
||||
- Load codec weights (may need lazy loading from separate checkpoint)
|
||||
- Implement `forward()`: codec codes -> audio waveform
|
||||
- Return `OmniOutput` with `multimodal_outputs`
|
||||
5. **Create stage config YAML** defining both stages, memory allocation, and model paths
|
||||
6. **Create stage input processor** for prompt building
|
||||
7. **Write end2end.py** test script
|
||||
|
||||
### Critical Parameters to Get Right
|
||||
|
||||
| Parameter | Impact if Wrong |
|
||||
|-----------|----------------|
|
||||
| Hop length | Audio duration wrong, streaming noise |
|
||||
| Token ID mapping | Garbage codes -> noise output |
|
||||
| Codebook count/size | Shape mismatch crashes |
|
||||
| Stop token | Generation never stops or stops too early |
|
||||
| dtype / autocast | Numerical issues, silent quality degradation |
|
||||
| Repetition penalty | Must match reference (often 1.0 for TTS) |
|
||||
|
||||
### Debugging Priority (from experience)
|
||||
|
||||
When audio output is wrong, check in this order:
|
||||
|
||||
1. **RoPE / attention**: Are position encodings correct? Is the attention mask right?
|
||||
2. **Normalization**: RMSNorm epsilon, layer norm placement (pre vs post)
|
||||
3. **Hop length**: Product of all upsample rates in the codec decoder
|
||||
4. **Token mapping**: Are codec IDs correctly offset from the vocabulary base?
|
||||
5. **Sampling parameters**: Temperature, top_k, top_p, repetition_penalty
|
||||
6. **Tensor layout**: Codebook-major vs frame-major ordering
|
||||
7. **dtype**: Float32 for codec decoders (autocast can corrupt audio)
|
||||
|
||||
### Deliverables
|
||||
|
||||
- Model files in `vllm_omni/model_executor/models/<model_name>/`
|
||||
- Stage config YAML
|
||||
- Working `end2end.py` with correct audio output
|
||||
- README.md in the example directory
|
||||
|
||||
## Phase 3: Online Serving
|
||||
|
||||
**Goal**: Expose the model via `/v1/audio/speech` API endpoint.
|
||||
|
||||
### Steps
|
||||
|
||||
1. **Register in `serving_speech.py`**:
|
||||
- Add model stage name to `_TTS_MODEL_STAGES` set
|
||||
- Add model detection flag (e.g., `_is_fish_speech`)
|
||||
- Implement prompt builder method (e.g., `_build_fish_speech_prompt()`)
|
||||
2. **Handle model-specific parameters**:
|
||||
- Voice cloning: `ref_audio` encoding and prompt injection
|
||||
- `max_new_tokens` override in sampling params
|
||||
- Model-specific default values
|
||||
3. **Create client scripts**: `speech_client.py`, `run_server.sh`
|
||||
4. **Test all response formats**: wav, mp3, flac, pcm
|
||||
5. **Add Gradio demo**: Interactive web UI with streaming support
|
||||
|
||||
### Voice Cloning Pattern
|
||||
|
||||
```python
|
||||
import base64
|
||||
from pathlib import Path
|
||||
|
||||
def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list:
|
||||
"""Build prompt with reference audio for voice cloning in serving_speech.py."""
|
||||
audio_bytes = Path(ref_audio_path).read_bytes()
|
||||
codes = codec.encode(audio_bytes) # Encode on CPU using model's codec (e.g., DAC)
|
||||
token_ids = [code + codec.vocab_offset for code in codes.flatten().tolist()]
|
||||
return [
|
||||
{"role": "system", "content": f"<|voice|>{''.join(chr(t) for t in token_ids)}"},
|
||||
{"role": "user", "content": text},
|
||||
]
|
||||
```
|
||||
|
||||
### Deliverables
|
||||
|
||||
- Updated `serving_speech.py` with model-specific prompt builder
|
||||
- Client scripts and server launcher
|
||||
- Gradio demo with streaming and voice cloning UI
|
||||
- Documentation (offline + online serving docs)
|
||||
|
||||
## Phase 4: Async Chunk (Streaming)
|
||||
|
||||
**Goal**: Enable inter-stage streaming so audio chunks are produced while AR generation continues.
|
||||
|
||||
### Steps
|
||||
|
||||
1. **Update stage config YAML**:
|
||||
```yaml
|
||||
async_chunk: true
|
||||
codec_chunk_frames: 25 # frames per chunk
|
||||
codec_left_context_frames: 25 # overlap for smooth boundaries
|
||||
```
|
||||
2. **Implement chunk handling in Stage 1**:
|
||||
- Accept partial input (chunk of codec codes)
|
||||
- Handle left context for smooth audio boundaries
|
||||
- Return partial audio in `OmniOutput`
|
||||
3. **Test streaming**:
|
||||
- Verify audio quality matches non-streaming output
|
||||
- Check for artifacts at chunk boundaries
|
||||
- Measure TTFA (time to first audio)
|
||||
4. **Update online serving** to support `stream=true` with PCM output
|
||||
|
||||
### Streaming Architecture
|
||||
|
||||
```
|
||||
Stage 0 (AR) Stage 1 (Decoder)
|
||||
| |
|
||||
|-- chunk 0 (25 frames) ------> decode -> audio chunk 0 -> client
|
||||
|-- chunk 1 (25 frames) ------> decode -> audio chunk 1 -> client
|
||||
|-- chunk 2 (25 frames) ------> decode -> audio chunk 2 -> client
|
||||
...
|
||||
```
|
||||
|
||||
### Key Considerations
|
||||
|
||||
- **Left context overlap**: Prevents audible artifacts at chunk boundaries
|
||||
- **Hop length matters**: `context_audio_samples = context_frames * hop_length`
|
||||
- **First chunk latency**: Can use larger initial chunk for better quality, then smaller chunks
|
||||
|
||||
### Deliverables
|
||||
|
||||
- Updated stage config with async_chunk enabled
|
||||
- Smooth streaming audio without boundary artifacts
|
||||
- TTFA metrics
|
||||
|
||||
## Phase 5: CUDA Graph Acceleration
|
||||
|
||||
**Goal**: Capture the AR loop as a CUDA graph for significant speedup.
|
||||
|
||||
### Steps
|
||||
|
||||
1. **Identify the hot loop**: The AR decoding loop that runs N steps per token
|
||||
2. **Create static buffers**:
|
||||
- KV caches with fixed max sequence length
|
||||
- Pre-built causal masks and position tensors per step
|
||||
- Static input/output tensors
|
||||
3. **Implement graph capture**:
|
||||
- Warm up with real data
|
||||
- Capture the forward pass
|
||||
- Replay with updated inputs
|
||||
4. **Handle constraints**:
|
||||
- Use `torch.argmax` instead of `torch.multinomial` (graph-safe)
|
||||
- Fixed batch size (fall back to eager for other sizes)
|
||||
- No dynamic control flow inside the graph
|
||||
|
||||
### Example: Code Predictor CUDA Graph (Qwen3-TTS)
|
||||
|
||||
```python
|
||||
import torch
|
||||
|
||||
class CodePredictorGraph:
|
||||
"""Captures the 16-step code predictor AR loop as a single CUDA graph."""
|
||||
|
||||
def setup_graph(self, device: torch.device, kv_heads: int = 4, head_dim: int = 64):
|
||||
self.num_steps = 16
|
||||
self.kv_cache = torch.zeros(1, kv_heads, self.num_steps, head_dim, device=device)
|
||||
self.positions = torch.arange(self.num_steps, device=device)
|
||||
self.causal_mask = torch.tril(torch.ones(self.num_steps, self.num_steps, device=device))
|
||||
self.input_buf = torch.zeros(1, 1, kv_heads * head_dim, device=device)
|
||||
self.output_buf = torch.zeros(1, self.num_steps, device=device, dtype=torch.long)
|
||||
# Warm up, then: self.graph = torch.cuda.CUDAGraph(); self.graph.capture(...)
|
||||
|
||||
def run_graph(self, initial_input: torch.Tensor) -> torch.Tensor:
|
||||
self.input_buf.copy_(initial_input)
|
||||
self.graph.replay()
|
||||
return self.output_buf.clone()
|
||||
```
|
||||
|
||||
### Performance Expectations
|
||||
|
||||
Based on Qwen3-TTS code predictor experience:
|
||||
- **3-5x speedup** for the graphed component
|
||||
- Only effective for fixed batch sizes (typically batch_size=1)
|
||||
- Falls back to eager mode for unsupported configurations
|
||||
|
||||
### Deliverables
|
||||
|
||||
- CUDA graph implementation for the AR hot loop
|
||||
- Benchmark script comparing eager vs graph performance
|
||||
- Documentation of constraints and fallback behavior
|
||||
|
||||
## Integration Checklist
|
||||
|
||||
Use this checklist when integrating a new TTS model:
|
||||
|
||||
### Phase 1: HF Reference
|
||||
- [ ] Reference model runs and produces correct audio
|
||||
- [ ] Architecture documented (stages, codebooks, tokens, sample rate)
|
||||
- [ ] Reference audio samples saved for comparison
|
||||
|
||||
### Phase 2: Stage Separation
|
||||
- [ ] Model registered in `registry.py`
|
||||
- [ ] Config classes created with `model_type` registration
|
||||
- [ ] Stage 0 (AR) implemented and generates correct tokens
|
||||
- [ ] Stage 1 (Decoder) produces correct audio from tokens
|
||||
- [ ] Stage config YAML created
|
||||
- [ ] `end2end.py` produces audio matching reference quality
|
||||
- [ ] README.md written
|
||||
|
||||
### Phase 3: Online Serving
|
||||
- [ ] Model added to `serving_speech.py`
|
||||
- [ ] Prompt builder handles text input correctly
|
||||
- [ ] Voice cloning works (if supported)
|
||||
- [ ] All response formats work (wav, mp3, flac, pcm)
|
||||
- [ ] Client scripts and server launcher created
|
||||
- [ ] Gradio demo working
|
||||
- [ ] Documentation added (offline + online docs, nav, supported models)
|
||||
|
||||
### Phase 4: Async Chunk
|
||||
- [ ] Stage config updated with `async_chunk: true`
|
||||
- [ ] Stage 1 handles partial chunks correctly
|
||||
- [ ] No audio artifacts at chunk boundaries
|
||||
- [ ] Streaming via API (`stream=true`) works
|
||||
- [ ] TTFA measured and acceptable
|
||||
|
||||
### Phase 5: CUDA Graph
|
||||
- [ ] Hot loop identified and profiled
|
||||
- [ ] Static buffers allocated
|
||||
- [ ] Graph captured and replays correctly
|
||||
- [ ] Benchmark shows meaningful speedup
|
||||
- [ ] Fallback to eager works for unsupported configs
|
||||
|
||||
## References
|
||||
|
||||
- [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) -- supported models and usage
|
||||
- [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) -- complete example of Phases 1-3
|
||||
- [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) -- complete example of all 5 phases
|
||||
- [Adding a TTS model (developer guide)](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/model/adding_tts_model.md)
|
||||
@@ -0,0 +1,34 @@
|
||||
# Claude Skills for vLLM-Omni
|
||||
|
||||
This directory contains Claude Code skills maintained for the `vllm-omni`
|
||||
repository. These skills capture repeatable workflows for common contributor
|
||||
tasks such as model integration, pull request review, and release note
|
||||
generation.
|
||||
|
||||
## Directory Structure
|
||||
|
||||
Each skill lives in its own directory under `.claude/skills/`. A skill may
|
||||
include:
|
||||
|
||||
- `SKILL.md`: the main workflow and operating instructions
|
||||
- `references/`: focused reference material used by the skill
|
||||
- `scripts/`: small helper scripts used by the skill
|
||||
|
||||
## Available Skills
|
||||
|
||||
- `add-diffusion-model`: guides integration of a new diffusion model into
|
||||
`vllm-omni`
|
||||
- `add-omni-model`: covers addition of new omni-modality model support
|
||||
- `add-tts-model`: covers integration of new TTS models and related serving
|
||||
workflows
|
||||
- `generate-release-note`: helps prepare release notes for repository changes
|
||||
- `review-pr`: provides a structured workflow for reviewing pull requests
|
||||
|
||||
## Maintenance Guidelines
|
||||
|
||||
- Keep skill names short and task-oriented.
|
||||
- Prefer repository-local paths, commands, and examples.
|
||||
- Avoid hardcoding fast-changing support matrices unless the skill is actively
|
||||
maintained alongside those changes.
|
||||
- Treat skills as contributor tooling: optimize for clarity, actionability, and
|
||||
low maintenance overhead.
|
||||
+13
-1
@@ -158,7 +158,19 @@ cython_debug/
|
||||
|
||||
# Claude
|
||||
CLAUDE.md
|
||||
.claude/
|
||||
/.claude/*
|
||||
!.claude/skills/
|
||||
!.claude/skills/readme.md
|
||||
!.claude/skills/add-diffusion-model/
|
||||
!.claude/skills/add-diffusion-model/SKILL.md
|
||||
!.claude/skills/add-diffusion-model/references/
|
||||
!.claude/skills/add-diffusion-model/references/*.md
|
||||
!.claude/skills/add-tts-model/
|
||||
!.claude/skills/add-tts-model/SKILL.md
|
||||
!.claude/skills/review-pr/
|
||||
!.claude/skills/review-pr/SKILL.md
|
||||
!.claude/skills/review-pr/references/
|
||||
!.claude/skills/review-pr/references/*.md
|
||||
|
||||
# Codex
|
||||
AGENTS.md
|
||||
|
||||
@@ -98,6 +98,7 @@ nav:
|
||||
- design/feature/ray_based_execution.md
|
||||
- design/feature/omni_connectors/
|
||||
- design/feature/cfg_parallel.md
|
||||
- design/feature/expert_parallel.md
|
||||
- design/feature/sequence_parallel.md
|
||||
- design/feature/tensor_parallel.md
|
||||
- design/feature/vae_parallel.md
|
||||
|
||||
@@ -0,0 +1,221 @@
|
||||
# Expert Parallel
|
||||
|
||||
This section describes how to add Expert Parallel (EP) to a diffusion transformer that uses Mixture-of-Experts (MoE) layers.
|
||||
We use **HunyuanImage3.0** as the reference implementation.
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Step-by-Step Implementation](#step-by-step-implementation)
|
||||
- [Testing](#testing)
|
||||
- [Reference Implementations](#reference-implementations)
|
||||
- [Summary](#summary)
|
||||
|
||||
---
|
||||
|
||||
## Overview
|
||||
|
||||
### What is Expert Parallel?
|
||||
|
||||
**Expert Parallel** is a parallelism strategy in Mixture-of-Experts (MoE) models that distributes different expert networks across distinct computational devices. Each device holds and computes only a subset of experts (local experts), with tokens dispatched to and gathered from remote devices via collective communication operations (e.g., All-to-All, All-Gather).
|
||||
|
||||
| Backend | Description |
|
||||
|---------|-------------|
|
||||
| `allgather_reducescatter` | Default backend based on allgather/reducescatter primitives, suitable for general EP+DP deployments.|
|
||||
|
||||
## Configuration
|
||||
|
||||
Enable EP by setting the `--enable-expert-parallel` flag. The EP size is automatically calculated as:
|
||||
|
||||
```text
|
||||
EP_SIZE = TP_SIZE × SP_SIZE × CFG_SIZE × DP_SIZE
|
||||
```
|
||||
|
||||
|
||||
Where:
|
||||
|
||||
- `TP_SIZE`: Tensor parallel size
|
||||
- `SP_SIZE`: Sequence parallel size
|
||||
- `CFG_SIZE`: Classifier-free guidance parallel size
|
||||
- `DP_SIZE`: Data parallel size
|
||||
- `EP_SIZE`: Expert parallel size (computed automatically)
|
||||
|
||||
Note:
|
||||
- Expert parallelism is only applicable to Mixture-of-Experts (MoE) models.
|
||||
- The EP group is created **per pipeline stage**, meaning it includes all ranks that participate in model parallelism except pipeline parallelism.
|
||||
- The underlying communication pattern for expert parallelism is **All-to-All** among the ranks in the EP group.
|
||||
|
||||
For example, consider a configuration with `TP=2`, `SP=1`, `CFG=2`, and `DP=4` (total 2×1×2×4 = 16 GPUs).
|
||||
|
||||
- Expert layers are handled by an EP group of size 16.
|
||||
|
||||
- Attention layers use tensor parallelism of size 2 within each of the 8 DP groups (because `DP×CFG×SP = 4×2×1 = 8` groups, each containing the 2 TP ranks). Inside each such group, the attention weights are sharded across the 2 GPUs.
|
||||
|
||||
|
||||
## Step-by-Step Implementation
|
||||
|
||||
### Step 1: Configure Expert Parallelism Settings
|
||||
|
||||
Calculate local experts per rank:
|
||||
|
||||
```
|
||||
ep_size = 8 # Expert Parallel size (typically equals TP size)
|
||||
num_experts = 64
|
||||
num_local_experts = num_experts // ep_size # 8 experts per card
|
||||
|
||||
# Check divisibility
|
||||
assert num_experts % ep_size == 0, "Experts must be divisible by EP size"
|
||||
```
|
||||
|
||||
### Step 2: Use Sparse MoE Block to enable EP routing.
|
||||
|
||||
Example:
|
||||
```
|
||||
from vllm.model_executor.layers.linear import ReplicatedLinear
|
||||
class HunYuanSparseMoeBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
config: PretrainedConfig,
|
||||
layer_id: int = -1,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
self.tp_size = get_tensor_model_parallel_world_size()
|
||||
self.n_routed_experts = config.num_experts # 64
|
||||
|
||||
# Calculate local experts per rank (key for EP)
|
||||
if self.tp_size > self.n_routed_experts:
|
||||
raise ValueError(f"TP size {self.tp_size} > experts {self.n_routed_experts}")
|
||||
|
||||
# Routing gate (replicated on all ranks, computes scores for all tokens to all experts)
|
||||
self.gate = ReplicatedLinear(
|
||||
config.hidden_size,
|
||||
config.num_experts,
|
||||
bias=False,
|
||||
quant_config=None,
|
||||
prefix=f"{prefix}.gate",
|
||||
)
|
||||
|
||||
# EP expert layer (factory loads platform-specific implementation)
|
||||
self.experts = HunyuanFusedMoE(...)
|
||||
```
|
||||
**Key Points:**
|
||||
- gate is **ReplicatedLinear** (replicated on all ranks)
|
||||
- experts is created via **HunyuanFusedMoE factory**, which automatically handles EP dispatch
|
||||
|
||||
### Step 3: Initialize EP Runtime
|
||||
|
||||
Initialize the EP communication context before model loading.
|
||||
```
|
||||
from vllm.utils.import_utils import resolve_obj_by_qualname
|
||||
# Call during __init__ or model loading
|
||||
op_name = "hunyuan_fused_moe"
|
||||
|
||||
# Prepare EP runtime: establish communication groups, assign local expert indices, init _expert_map
|
||||
current_omni_platform.prepare_diffusion_op_runtime(op_name)
|
||||
|
||||
# Factory automatically resolves platform implementation (GPU: FusedMoE / NPU: AscendFusedMoE)
|
||||
impl = resolve_obj_by_qualname(
|
||||
current_omni_platform.get_diffusion_model_impl_qualname(op_name)
|
||||
)
|
||||
```
|
||||
|
||||
### Step 4: Expert Weight Mapping & Loading
|
||||
|
||||
Each rank loads only the expert weights assigned to its local allocation.
|
||||
```
|
||||
# Get expert parameter mapping (different per rank)
|
||||
expert_mapping = HunyuanFusedMoE.make_expert_params_mapping(
|
||||
model=self,
|
||||
ckpt_gate_proj_name="gate_proj",
|
||||
ckpt_down_proj_name="down_proj",
|
||||
ckpt_up_proj_name="up_proj",
|
||||
num_experts=64,
|
||||
num_redundant_experts=0,
|
||||
)
|
||||
# Returns: [(param_name, weight_name, expert_id, shard_id), ...]
|
||||
# Note: Each rank only contains mappings for its local expert_ids
|
||||
|
||||
# Filter non-local experts during loading
|
||||
for name, loaded_weight in weights:
|
||||
if "mlp.experts" in name:
|
||||
# Parse expert_id from weight name (implementation needed)
|
||||
expert_id = parse_expert_id_from_name(name)
|
||||
local_expert_start = (ep_rank) * num_local_experts
|
||||
local_expert_end = (ep_rank + 1) * num_local_experts
|
||||
|
||||
if not (local_expert_start <= expert_id < local_expert_end):
|
||||
continue # Skip non-local expert weights
|
||||
```
|
||||
### Step 5: Forward Pass with EP
|
||||
|
||||
Example (MoE Forward):
|
||||
```
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
orig_shape = hidden_states.shape
|
||||
hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
|
||||
|
||||
# 1. Global routing computation (all tokens, all expert scores)
|
||||
# hidden_states: [num_tokens, hidden_dim] (full tensor)
|
||||
router_logits, _ = self.gate(hidden_states) # [num_tokens, num_experts]
|
||||
|
||||
# 2. EP dispatch and compute (HunyuanFusedMoE handles all_to_all internally)
|
||||
# - Dispatch: Send tokens to target ranks based on router_logits
|
||||
# - Local Compute: Each rank processes only its num_local_experts
|
||||
# - Combine: Results returned to original token positions
|
||||
final_hidden_states = self.experts(
|
||||
hidden_states=hidden_states,
|
||||
router_logits=router_logits,
|
||||
)
|
||||
|
||||
# 3. Add shared expert output (not EP, computed on all ranks)
|
||||
if self.shared_mlp is not None:
|
||||
shared_out = self.shared_mlp(hidden_states)
|
||||
final_hidden_states = final_hidden_states + shared_out
|
||||
|
||||
# 4. Tensor Parallel All-Reduce (synchronize across TP group)
|
||||
if self.tp_size > 1:
|
||||
final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(
|
||||
final_hidden_states
|
||||
)
|
||||
|
||||
return final_hidden_states.view(orig_shape)
|
||||
```
|
||||
|
||||
## Testing
|
||||
After adding Expert Parallel support, test via command line:
|
||||
```bash
|
||||
cd examples/offline_inference/text_to_image
|
||||
python text_to_image.py \
|
||||
--model Your-org/your-model \
|
||||
--prompt "a cup of coffee on the table" \
|
||||
--output "ep_enabled.png" \
|
||||
--num-inference-steps 50 \
|
||||
--guidance-scale 5.0 \
|
||||
--tensor-parallel-size 8 \
|
||||
--seed 1234 \
|
||||
--enable-expert-parallel
|
||||
```
|
||||
|
||||
vLLM‑Omni currently focuses on core diffusion model inference acceleration, so the Expert Parallel implementation includes only the basic multi‑GPU expert sharding functionality (enabled via --enable-expert-parallel). Advanced features such as communication backend selection (--all2all-backend), load balancing (--enable-eplb and its configuration), and multi‑node deployment belong to the extended capabilities of the main vLLM project and have not yet been integrated into Omni.
|
||||
|
||||
## Reference Implementations
|
||||
|
||||
Complete examples in the codebase:
|
||||
|
||||
| Model | Path | Pattern | Notes |
|
||||
|-------|------|---------|-------|
|
||||
| **HunyuanImage3.0** | `vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py` | Standard EP | Full implementation with validation |
|
||||
| **EP Tests** | `vllm-omni/tests/e2e/offline_inference/test_expert_parallel.py` | E2E testing | EP correctness and performance |
|
||||
| **Constraint Tests** | `vllm-omni/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py` | Unit testing | Validation logic |
|
||||
|
||||
---
|
||||
## Summary
|
||||
|
||||
Adding Expert Parallel support to diffusion model:
|
||||
|
||||
1. **Identify MoE layers** - Locate the router and expert networks in each transformer block.
|
||||
2. **Validate EP constraints** – Ensure num_experts is divisible by expert_parallel_size.
|
||||
3. **Test** - Run with enable-expert-parallel, check memory reduction, speedup, and output quality against single‑GPU baseline.
|
||||
@@ -33,8 +33,8 @@ runtime:
|
||||
zmq_port: 50051 # ZMQ base port (see "Port Offset Scheme" below)
|
||||
protocol: "rdma" # "rdma" or "tcp"
|
||||
device_name: "" # RDMA device (e.g., "mlx5_0"), empty for auto-detect
|
||||
memory_pool_size: 2147483648 # 2GB memory pool
|
||||
memory_pool_device: "cpu" # "cpu" for pinned memory, "cuda" for GPUDirect RDMA
|
||||
memory_pool_size: 4294967296 # 4 GB (CPU); use 2147483648 (2 GB) for GPU
|
||||
memory_pool_device: "cpu" # "cpu" for pinned memory (recommended), "cuda" for GPUDirect RDMA
|
||||
```
|
||||
|
||||
Wire stages to the connector:
|
||||
@@ -64,8 +64,8 @@ stage_args:
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|---|---|---|
|
||||
| `memory_pool_size` | 1 GB | Total size of the RDMA-registered memory pool in bytes. |
|
||||
| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity). |
|
||||
| `memory_pool_size` | 4 GB (CPU) / 2 GB (GPU) | Total size of the RDMA-registered memory pool in bytes. Recommended 4 GB for CPU pinned memory; 2 GB for GPU VRAM to conserve device memory. |
|
||||
| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended, works on all topologies). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity, PIX topology). |
|
||||
|
||||
### Networking
|
||||
|
||||
@@ -107,10 +107,10 @@ receiver_connect = remote_side_channel_port + tp_rank
|
||||
|
||||
## Memory Pool Modes
|
||||
|
||||
| Mode | Config | Data Flow | Best For |
|
||||
|---|---|---|---|
|
||||
| CPU Pinned | `memory_pool_device: "cpu"` | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
|
||||
| GPUDirect | `memory_pool_device: "cuda"` | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
|
||||
| Mode | Config | Recommended Pool Size | Data Flow | Best For |
|
||||
|---|---|---|---|---|
|
||||
| CPU Pinned | `memory_pool_device: "cpu"` | 4 GB | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
|
||||
| GPUDirect | `memory_pool_device: "cuda"` | 2 GB | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
|
||||
|
||||
> **Note**: GPUDirect RDMA requires the NIC and GPU to share a direct PCIe
|
||||
> switch (PIX topology). On systems where they are connected via PXB or NODE,
|
||||
|
||||
@@ -60,5 +60,6 @@ th {
|
||||
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
|
||||
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
|
||||
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
|
||||
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
|
||||
|
||||
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
|
||||
@@ -139,11 +139,15 @@ Factory function `get_offload_backend()` selects the appropriate backend based o
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attr (Layerwise specific) |
|
||||
|--------------|----------------|-----------|---------------------|-------------------|-------------|
|
||||
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
|
||||
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
|
||||
| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attrs (Layerwise specific) |
|
||||
|--------------|----------------|-----------|---------------------|-------------------|-----------------------------------|
|
||||
| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | `"transformer_blocks"`, `"single_transformer_blocks"` |
|
||||
| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | `"layers"` |
|
||||
| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | `"transformer"` |
|
||||
| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` |
|
||||
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` |
|
||||
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
|
||||
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
|
||||
|
||||
**Notes:**
|
||||
- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
|
||||
|
||||
@@ -12,7 +12,7 @@
|
||||
|
||||
vLLM-Omni supports various advanced features for diffusion models:
|
||||
|
||||
- Acceleration: **cache methods**, **parallelism methods**
|
||||
- Acceleration: **cache methods**, **parallelism methods**, **startup optimizations**
|
||||
- Memory optimization: **cpu offloading**, **quantization**
|
||||
- Extensions: **LoRA inference**
|
||||
- Execution modes: **step execution**
|
||||
@@ -44,6 +44,12 @@ Parallelism methods distribute computation across GPUs without quality loss (mat
|
||||
| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
|
||||
| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
|
||||
|
||||
#### Startup Optimization
|
||||
|
||||
| Method | Description | Best For |
|
||||
|--------|-------------|----------|
|
||||
| **[Multi-Thread Weight Loading](#multi-thread-weight-loading)** | Loads safetensors shards in parallel using a thread pool | All diffusion models; reduces startup from minutes to seconds |
|
||||
|
||||
**Note:** Some acceleration methods can be combined together for optimized performance. See [Feature Compatibility Table](#feature-compatibility) and [Feature Compatibility Tutorial](feature_compatibility.md) for detailed configuration examples.
|
||||
|
||||
### Memory Optimization
|
||||
@@ -107,19 +113,19 @@ The following tables show which models support each feature:
|
||||
| **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
|
||||
| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
|
||||
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
|
||||
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
|
||||
|
||||
> Notes:
|
||||
@@ -179,6 +185,59 @@ The following tables show which models support each feature:
|
||||
6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT) or LoRA.
|
||||
|
||||
|
||||
## Multi-Thread Weight Loading
|
||||
|
||||
Large diffusion models can take several minutes to load weights at startup (e.g., ~3 min for Qwen-Image, ~5 min for Wan2.2 I2V 14B). Multi-thread weight loading speeds up this process by loading safetensors shards in parallel using a thread pool instead of sequentially.
|
||||
|
||||
This optimization is **enabled by default** with 4 threads. No configuration is needed for the default behavior.
|
||||
|
||||
### Configuration
|
||||
|
||||
| Parameter | CLI Flag | Default | Description |
|
||||
|-----------|----------|---------|-------------|
|
||||
| `enable_multithread_weight_load` | `--disable-multithread-weight-load` | `True` (enabled) | Pass the flag to disable multi-thread loading |
|
||||
| `num_weight_load_threads` | `--num-weight-load-threads` | `4` | Number of threads for parallel weight loading |
|
||||
|
||||
!!! tip
|
||||
The default of 4 threads balances speed and disk I/O contention. On fast NVMe storage you may benefit from more threads (e.g., 8). On HDD or network storage, the default of 4 avoids saturating I/O bandwidth.
|
||||
|
||||
### Online Serving
|
||||
|
||||
```bash
|
||||
# Default (multi-thread enabled, 4 threads)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091
|
||||
|
||||
# Custom thread count
|
||||
vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --num-weight-load-threads 8
|
||||
|
||||
# Disable multi-thread loading
|
||||
vllm serve Qwen/Qwen-Image --omni --disable-multithread-weight-load
|
||||
```
|
||||
|
||||
### Offline Inference
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
# Default (multi-thread enabled, 4 threads)
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
|
||||
# Custom thread count
|
||||
omni = Omni(
|
||||
model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
num_weight_load_threads=8,
|
||||
)
|
||||
```
|
||||
|
||||
### Benchmarks
|
||||
|
||||
Measured on NVIDIA H800:
|
||||
|
||||
| Model | Before | After | Speedup |
|
||||
|-------|--------|-------|---------|
|
||||
| **Qwen/Qwen-Image** (53.7 GiB) | 168s | 27s | **6.2x** |
|
||||
| **Wan-AI/Wan2.2-I2V-A14B-Diffusers** (64.5 GiB) | 283s | 56s | **5.1x** |
|
||||
|
||||
## Learn More
|
||||
|
||||
**Cache Acceleration:**
|
||||
@@ -204,6 +263,10 @@ The following tables show which models support each feature:
|
||||
|
||||
- **[Step Execution Guide](diffusion/step_execution.md)** - Per-step denoise execution with mid-request abort support
|
||||
|
||||
**Startup Optimization:**
|
||||
|
||||
- **[Multi-Thread Weight Loading](#multi-thread-weight-loading)** - Speed up model startup by loading safetensors shards in parallel
|
||||
|
||||
**Advanced Topics:**
|
||||
|
||||
- **[Feature Compatibility](feature_compatibility.md)** - How to combine multiple features for maximum performance
|
||||
@@ -0,0 +1,110 @@
|
||||
# Dynin-Omni Offline End2End Example
|
||||
|
||||
This folder contains a unified offline inference entrypoint:
|
||||
|
||||
- `end2end.py`
|
||||
|
||||
## 1. Environment Setup
|
||||
|
||||
Run from repository root:
|
||||
|
||||
```bash
|
||||
cd <REPO_ROOT>
|
||||
```
|
||||
|
||||
If needed, install this repo in editable mode:
|
||||
|
||||
```bash
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
## 2. Extra Dependencies (EMOVA)
|
||||
|
||||
Install the following packages for EMOVA-related components:
|
||||
|
||||
```bash
|
||||
pip install \
|
||||
"phonemizer==3.3.0" \
|
||||
"Unidecode==1.4.0" \
|
||||
"hydra-core==1.3.2" \
|
||||
"pytorch-lightning==1.1.0" \
|
||||
"wget==3.2" \
|
||||
"wrapt==2.1.1" \
|
||||
"onnx==1.20.1" \
|
||||
"frozendict==2.4.7" \
|
||||
"inflect==7.5.0" \
|
||||
"braceexpand==0.1.7" \
|
||||
"webdataset==1.0.2" \
|
||||
"torch-stft==0.1.4" \
|
||||
"editdistance==0.8.1"
|
||||
```
|
||||
|
||||
## 3. Hardware and VRAM Requirements
|
||||
|
||||
This example uses a 3-stage pipeline on one GPU by default
|
||||
([`dynin_omni.yaml`](../../../vllm_omni/model_executor/stage_configs/dynin_omni.yaml)):
|
||||
|
||||
- Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
|
||||
- Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`
|
||||
- Stage-2 (`token2audio`): `gpu_memory_utilization: 0.1`
|
||||
|
||||
### Requested GPU Memory Budget from `gpu_memory_utilization`
|
||||
|
||||
| Stage | Utilization | A100 80GB | H200 141GB |
|
||||
| :-- | :-- | :-- | :-- |
|
||||
| Stage-0 (token2text) | 0.5 | ~40.0 GB | ~70.5 GB |
|
||||
| Stage-1 (token2image) | 0.1 | ~8.0 GB | ~14.1 GB |
|
||||
| Stage-2 (token2audio) | 0.1 | ~8.0 GB | ~14.1 GB |
|
||||
| Total requested budget | 0.7 | ~56.0 GB | ~98.7 GB |
|
||||
|
||||
### Observed Runtime Signal (from your log)
|
||||
|
||||
- Stage-0 reported: `Model loading took 15.12 GiB memory` (weights footprint signal).
|
||||
- Stages 1/2 can still add runtime memory depending on task path and backend allocations.
|
||||
- Keep extra headroom for CUDA/PyTorch overhead and temporary allocations.
|
||||
|
||||
### GPU Compatibility
|
||||
|
||||
- Confirmed target GPUs for this setup: **NVIDIA H200**, **NVIDIA A100**.
|
||||
- CI/e2e coverage in this repo also includes CUDA **L4** markers for Dynin tests.
|
||||
|
||||
## 4. End2End Run Examples
|
||||
|
||||
```bash
|
||||
# t2t
|
||||
python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
|
||||
--task t2t --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
|
||||
|
||||
# i2t
|
||||
python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
|
||||
--task i2t --model snu-aidas/Dynin-Omni --image <IMAGE_PATH> --text "Please describe this image in detail."
|
||||
|
||||
# s2t
|
||||
python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
|
||||
--task s2t --model snu-aidas/Dynin-Omni --audio <AUDIO_PATH> --text "Transcribe the given audio."
|
||||
|
||||
# t2i
|
||||
python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
|
||||
--task t2i --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
|
||||
|
||||
# v2t
|
||||
python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
|
||||
--task v2t --model snu-aidas/Dynin-Omni --video <VIDEO_PATH> --text "Describe this video in detail."
|
||||
|
||||
# i2i
|
||||
python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
|
||||
--task i2i --model snu-aidas/Dynin-Omni --image <IMAGE_PATH> --text <INSTRUCTION_TEXT>
|
||||
|
||||
# t2s
|
||||
python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
|
||||
--task t2s --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
|
||||
```
|
||||
|
||||
## 5. Notes
|
||||
|
||||
- Outputs are saved under task-specific directories in `/tmp` by default.
|
||||
- You can override output path with `--output-dir`.
|
||||
- If you want to force local config resolution, pass `--dynin-config-path <PATH_TO_DYNIN_OMNI_YAML>`.
|
||||
- If you see the warning
|
||||
`max_num_batched_tokens (32768) exceeds max_num_seqs * max_model_len (4096)`,
|
||||
reduce `max_num_batched_tokens` in stage config (for example, `4096` in CI config).
|
||||
File diff suppressed because it is too large.
Load diff
@@ -51,5 +51,6 @@ Key arguments:
|
||||
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
|
||||
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
|
||||
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
|
||||
- `--strength`: **Z-Image only** - controls the denoising start timestep for I2I (default: 0.6). Range: [0.0, 1.0]. Lower values preserve more of the original image; higher values allow more creative changes.
|
||||
|
||||
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
|
||||
@@ -146,7 +146,7 @@ def parse_args() -> argparse.Namespace:
|
||||
"--audio-sample-rate",
|
||||
type=int,
|
||||
default=24000,
|
||||
help="Sample rate for audio output when saved (default: 24000 for LTX2).",
|
||||
help="Sample rate for audio output when saved (default: 24000).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cache-backend",
|
||||
@@ -471,15 +471,9 @@ def main():
|
||||
|
||||
video_array = _ensure_frame_list(video_array)
|
||||
|
||||
use_ltx2_export = is_ltx2
|
||||
encode_video = None
|
||||
if use_ltx2_export:
|
||||
try:
|
||||
from diffusers.pipelines.ltx2.export_utils import encode_video
|
||||
except ImportError:
|
||||
encode_video = None
|
||||
if audio is not None:
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
|
||||
if use_ltx2_export and encode_video is not None:
|
||||
if isinstance(video_array, list):
|
||||
frames_np = np.stack(video_array, axis=0)
|
||||
elif isinstance(video_array, np.ndarray):
|
||||
@@ -490,25 +484,24 @@ def main():
|
||||
if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
|
||||
frames_np = frames_np[..., :3]
|
||||
|
||||
audio_out = None
|
||||
if audio is not None:
|
||||
if isinstance(audio, list):
|
||||
audio = audio[0] if audio else None
|
||||
if isinstance(audio, np.ndarray):
|
||||
audio = torch.from_numpy(audio)
|
||||
if isinstance(audio, torch.Tensor):
|
||||
audio_out = audio
|
||||
if audio_out.dim() > 1:
|
||||
audio_out = audio_out[0]
|
||||
audio_out = audio_out.float().cpu()
|
||||
frames_u8 = (np.clip(frames_np, 0.0, 1.0) * 255).round().clip(0, 255).astype("uint8")
|
||||
|
||||
encode_video(
|
||||
frames_np,
|
||||
fps=fps,
|
||||
audio=audio_out,
|
||||
audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
|
||||
output_path=str(output_path),
|
||||
audio_np = audio
|
||||
if isinstance(audio_np, list):
|
||||
audio_np = audio_np[0] if audio_np else None
|
||||
if isinstance(audio_np, torch.Tensor):
|
||||
audio_np = audio_np.detach().cpu().float().numpy()
|
||||
if isinstance(audio_np, np.ndarray):
|
||||
audio_np = np.squeeze(audio_np).astype(np.float32)
|
||||
|
||||
video_bytes = mux_video_audio_bytes(
|
||||
frames_u8,
|
||||
audio_np,
|
||||
fps=float(fps),
|
||||
audio_sample_rate=args.audio_sample_rate,
|
||||
)
|
||||
with open(str(output_path), "wb") as f:
|
||||
f.write(video_bytes)
|
||||
else:
|
||||
export_to_video(video_array, str(output_path), fps=fps)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
|
||||
@@ -94,16 +94,21 @@ def main():
|
||||
print(f"Video frames: shape={video_frames.shape}, dtype={video_frames.dtype}")
|
||||
|
||||
audio_waveform = None
|
||||
if hasattr(first, "multimodal_output") and first.multimodal_output:
|
||||
audio_waveform = first.multimodal_output.get("audio")
|
||||
mm = first.multimodal_output or {}
|
||||
if mm:
|
||||
audio_waveform = mm.get("audio")
|
||||
if audio_waveform is not None:
|
||||
print(f"Audio waveform: shape={audio_waveform.shape}, dtype={audio_waveform.dtype}")
|
||||
|
||||
output_fps = float(mm.get("fps", 25))
|
||||
output_sr = int(mm.get("audio_sample_rate", 24000))
|
||||
print(f"Using fps={output_fps}, audio_sample_rate={output_sr} from model output")
|
||||
|
||||
video_bytes = mux_video_audio_bytes(
|
||||
video_frames,
|
||||
audio_waveform,
|
||||
fps=25.0,
|
||||
audio_sample_rate=44100,
|
||||
fps=output_fps,
|
||||
audio_sample_rate=output_sr,
|
||||
)
|
||||
with open(args.output, "wb") as f:
|
||||
f.write(video_bytes)
|
||||
|
||||
@@ -376,7 +376,7 @@ def main():
|
||||
f"vae_patch_parallel_size={args.vae_patch_parallel_size}, "
|
||||
f"enable_expert_parallel={args.enable_expert_parallel}."
|
||||
)
|
||||
print(f" CPU offload: {args.enable_cpu_offload}")
|
||||
print(f" CPU offload: {args.enable_cpu_offload}; CPU Layerwise Offload: {args.enable_layerwise_offload}")
|
||||
print(f" Image size: {args.width}x{args.height}")
|
||||
if args.lora_path:
|
||||
print(f" LoRA: scale={args.lora_scale}")
|
||||
|
||||
@@ -160,7 +160,7 @@ def parse_args() -> argparse.Namespace:
|
||||
"--audio-sample-rate",
|
||||
type=int,
|
||||
default=24000,
|
||||
help="Sample rate for audio output when saved (default: 24000 for LTX2).",
|
||||
help="Sample rate for audio output when saved (default: 24000).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vae-patch-parallel-size",
|
||||
@@ -430,17 +430,8 @@ def main():
|
||||
|
||||
video_array = _ensure_frame_list(video_array)
|
||||
|
||||
use_ltx2_export = False
|
||||
if args.model and "ltx" in str(args.model).lower():
|
||||
use_ltx2_export = True
|
||||
if audio is not None:
|
||||
use_ltx2_export = True
|
||||
|
||||
if use_ltx2_export:
|
||||
try:
|
||||
from diffusers.pipelines.ltx2.export_utils import encode_video
|
||||
except ImportError:
|
||||
raise ImportError("diffusers is required for LTX2 encode_video.")
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
|
||||
if isinstance(video_array, list):
|
||||
frames_np = np.stack(video_array, axis=0)
|
||||
@@ -449,28 +440,24 @@ def main():
|
||||
else:
|
||||
frames_np = np.asarray(video_array)
|
||||
|
||||
frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
|
||||
video_tensor = torch.from_numpy(frames_u8)
|
||||
frames_u8 = (np.clip(frames_np, 0.0, 1.0) * 255).round().clip(0, 255).astype("uint8")
|
||||
|
||||
audio_out = None
|
||||
if audio is not None:
|
||||
if isinstance(audio, list):
|
||||
audio = audio[0] if audio else None
|
||||
if isinstance(audio, np.ndarray):
|
||||
audio = torch.from_numpy(audio)
|
||||
if isinstance(audio, torch.Tensor):
|
||||
audio_out = audio
|
||||
if audio_out.dim() > 1:
|
||||
audio_out = audio_out[0]
|
||||
audio_out = audio_out.float().cpu()
|
||||
audio_np = audio
|
||||
if isinstance(audio_np, list):
|
||||
audio_np = audio_np[0] if audio_np else None
|
||||
if isinstance(audio_np, torch.Tensor):
|
||||
audio_np = audio_np.detach().cpu().float().numpy()
|
||||
if isinstance(audio_np, np.ndarray):
|
||||
audio_np = np.squeeze(audio_np).astype(np.float32)
|
||||
|
||||
encode_video(
|
||||
video_tensor,
|
||||
fps=args.fps,
|
||||
audio=audio_out,
|
||||
audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
|
||||
output_path=str(output_path),
|
||||
video_bytes = mux_video_audio_bytes(
|
||||
frames_u8,
|
||||
audio_np,
|
||||
fps=float(args.fps),
|
||||
audio_sample_rate=args.audio_sample_rate,
|
||||
)
|
||||
with open(str(output_path), "wb") as f:
|
||||
f.write(video_bytes)
|
||||
else:
|
||||
export_to_video(video_array, str(output_path), fps=args.fps)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
# Dynin-Omni Online Serving Example
|
||||
|
||||
## Installation
|
||||
|
||||
Please refer to [README.md](../../../README.md).
|
||||
|
||||
## Launch the Server
|
||||
|
||||
First, find the `transformers_modules` path:
|
||||
|
||||
```bash
|
||||
python - <<'PY'
|
||||
from transformers.utils.hub import HF_MODULES_CACHE
|
||||
print(HF_MODULES_CACHE)
|
||||
PY
|
||||
```
|
||||
|
||||
Then export it for both `PYTHONPATH` and `HF_MODULES_CACHE`:
|
||||
|
||||
```bash
|
||||
export PYTHONPATH=<transformers_modules_path>:$PYTHONPATH
|
||||
export HF_MODULES_CACHE=<transformers_modules_path>
|
||||
```
|
||||
|
||||
Run from repository root:
|
||||
|
||||
```bash
|
||||
vllm-omni serve snu-aidas/Dynin-Omni \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
|
||||
```
|
||||
|
||||
If `vllm-omni` is not in PATH, run:
|
||||
|
||||
```bash
|
||||
PYTHONPATH="$(pwd)" python -m vllm_omni.entrypoints.cli.main serve snu-aidas/Dynin-Omni \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
|
||||
```
|
||||
|
||||
Wait until the server logs show both `All stages initialized successfully` and
|
||||
`Application startup complete.` before sending requests.
|
||||
|
||||
## Send Requests via Python Client
|
||||
|
||||
Move to the example directory:
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/dynin_omni
|
||||
```
|
||||
|
||||
### Text -> Image
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type t2i \
|
||||
--prompt "A realistic indoor living room with natural daylight."
|
||||
```
|
||||
|
||||
### Image -> Image
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type i2i \
|
||||
--image-path ../../offline_inference/dynin_omni/data/image/sofa_under_water.jpg \
|
||||
--prompt "Transform this surreal underwater setting into a realistic indoor living room while preserving the sofa layout."
|
||||
```
|
||||
|
||||
### Text -> Speech
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type t2s \
|
||||
--prompt "Hello. This is Dynin-omni."
|
||||
```
|
||||
|
||||
## CLI Arguments
|
||||
|
||||
- `--query-type` (`t2i|t2s|i2i`)
|
||||
- `--model` (default: `snu-aidas/Dynin-Omni`)
|
||||
- `--host` / `--port` (OpenAI-compatible vLLM endpoint)
|
||||
- `--prompt` (custom text)
|
||||
- `--image-path` (required for `i2i`)
|
||||
- `--modalities` (optional output modalities override)
|
||||
- `--output-dir` (default: `/tmp/dynin_online_outputs`)
|
||||
|
||||
## Notes
|
||||
|
||||
- This client currently supports only `t2i`, `t2s`, and `i2i`.
|
||||
- `t2t` is intentionally not exposed in this online example.
|
||||
- This example intentionally uses the OpenAI-compatible chat completion endpoint.
|
||||
- Task routing for non-text outputs relies on Dynin task trigger tokens (`<|t2i|>`, `<|i2i|>`, `<|t2s|>`) injected by the client.
|
||||
- Outputs are saved under `/tmp/dynin_online_outputs` by default.
|
||||
- Dynin stage-0 warmup can take a while on first startup; do not send requests before startup completes.
|
||||
- Dynin itself can execute text-returning tasks such as `t2t`, `s2t`, `i2t`, and `v2t`, but this online serving example currently runs stage-0 in `generation` mode. In that path, the generation worker does not surface the final text as `output.text`, so OpenAI chat responses for those text-output tasks may complete internally but still return empty text.
|
||||
+342
@@ -0,0 +1,342 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
DEFAULT_MODEL = "snu-aidas/Dynin-Omni"
|
||||
DEFAULT_OUTPUT_DIR = "/tmp/dynin_online_outputs"
|
||||
|
||||
QUERY_CHOICES = ("t2i", "t2s", "i2i")
|
||||
DEFAULT_PROMPT_BY_QUERY = {
|
||||
"t2i": "A high quality detailed living room interior photo.",
|
||||
"t2s": "Please read this sentence naturally: Hello from Dynin-Omni online serving.",
|
||||
"i2i": "Transform this image into a realistic indoor living room while preserving layout.",
|
||||
}
|
||||
DEFAULT_MODALITIES_BY_QUERY = {
|
||||
"t2i": ["image"],
|
||||
"t2s": ["audio"],
|
||||
"i2i": ["image"],
|
||||
}
|
||||
OFFLINE_PARITY_STAGE_COUNT = 3
|
||||
OFFLINE_PARITY_STAGE_SAMPLING = {
|
||||
"max_tokens": 1,
|
||||
"temperature": 0.0,
|
||||
"top_p": 1.0,
|
||||
"detokenize": False,
|
||||
}
|
||||
|
||||
|
||||
def _infer_mime_type(path: Path) -> str:
|
||||
mime_type, _ = mimetypes.guess_type(str(path))
|
||||
return mime_type or "application/octet-stream"
|
||||
|
||||
|
||||
def _encode_file_as_data_url(path: Path) -> str:
|
||||
mime_type = _infer_mime_type(path)
|
||||
raw = path.read_bytes()
|
||||
encoded = base64.b64encode(raw).decode("utf-8")
|
||||
return f"data:{mime_type};base64,{encoded}"
|
||||
|
||||
|
||||
def _to_image_url(path_or_url: str) -> str:
|
||||
value = str(path_or_url)
|
||||
if value.startswith(("http://", "https://", "data:image/")):
|
||||
return value
|
||||
path = Path(value).expanduser().resolve()
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Image file not found: {path}")
|
||||
return _encode_file_as_data_url(path)
|
||||
|
||||
|
||||
def _build_user_content(query_type: str, prompt: str, image_path: str | None) -> list[dict[str, Any]]:
|
||||
if query_type == "t2i":
|
||||
return [{"type": "text", "text": f"<|t2i|> {prompt}"}]
|
||||
|
||||
if query_type == "t2s":
|
||||
return [{"type": "text", "text": f"<|t2s|> {prompt}"}]
|
||||
|
||||
if query_type == "i2i":
|
||||
if not image_path:
|
||||
raise ValueError("--image-path is required for query type i2i")
|
||||
return [
|
||||
{"type": "text", "text": f"<|i2i|> {prompt}"},
|
||||
{"type": "image_url", "image_url": {"url": _to_image_url(image_path)}},
|
||||
]
|
||||
|
||||
raise ValueError(f"Unsupported query_type: {query_type}")
|
||||
|
||||
|
||||
def _collect_text_from_content(content: Any) -> list[str]:
|
||||
texts: list[str] = []
|
||||
if isinstance(content, str):
|
||||
stripped = content.strip()
|
||||
if stripped:
|
||||
texts.append(stripped)
|
||||
return texts
|
||||
|
||||
if isinstance(content, dict):
|
||||
for key in ("text", "content", "value", "output_text"):
|
||||
text_value = content.get(key)
|
||||
if isinstance(text_value, str) and text_value.strip():
|
||||
texts.append(text_value.strip())
|
||||
return texts
|
||||
|
||||
if isinstance(content, list):
|
||||
for item in content:
|
||||
texts.extend(_collect_text_from_content(item))
|
||||
return texts
|
||||
|
||||
content_text = getattr(content, "text", None)
|
||||
if isinstance(content_text, str) and content_text.strip():
|
||||
texts.append(content_text.strip())
|
||||
content_value = getattr(content, "content", None)
|
||||
if isinstance(content_value, str) and content_value.strip():
|
||||
texts.append(content_value.strip())
|
||||
output_text = getattr(content, "output_text", None)
|
||||
if isinstance(output_text, str) and output_text.strip():
|
||||
texts.append(output_text.strip())
|
||||
return texts
|
||||
|
||||
|
||||
def _extract_text_outputs(chat_completion: Any) -> list[str]:
|
||||
texts: list[str] = []
|
||||
for choice in getattr(chat_completion, "choices", []) or []:
|
||||
message = getattr(choice, "message", None)
|
||||
if message is None:
|
||||
continue
|
||||
content = getattr(message, "content", None)
|
||||
texts.extend(_collect_text_from_content(content))
|
||||
reasoning_content = getattr(message, "reasoning_content", None)
|
||||
if isinstance(reasoning_content, str) and reasoning_content.strip():
|
||||
texts.append(reasoning_content.strip())
|
||||
choice_text = getattr(choice, "text", None)
|
||||
if isinstance(choice_text, str) and choice_text.strip():
|
||||
texts.append(choice_text.strip())
|
||||
top_level_output_text = getattr(chat_completion, "output_text", None)
|
||||
if isinstance(top_level_output_text, str) and top_level_output_text.strip():
|
||||
texts.append(top_level_output_text.strip())
|
||||
return texts
|
||||
|
||||
|
||||
def _extract_image_data_urls(chat_completion: Any) -> list[str]:
|
||||
urls: list[str] = []
|
||||
for choice in getattr(chat_completion, "choices", []) or []:
|
||||
message = getattr(choice, "message", None)
|
||||
if message is None:
|
||||
continue
|
||||
content = getattr(message, "content", None)
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
for item in content:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
if item.get("type") != "image_url":
|
||||
continue
|
||||
image_url = (item.get("image_url") or {}).get("url")
|
||||
if isinstance(image_url, str) and image_url.startswith("data:image"):
|
||||
urls.append(image_url)
|
||||
return urls
|
||||
|
||||
|
||||
def _extract_audio_payloads(chat_completion: Any) -> list[bytes]:
|
||||
payloads: list[bytes] = []
|
||||
for choice in getattr(chat_completion, "choices", []) or []:
|
||||
message = getattr(choice, "message", None)
|
||||
if message is None:
|
||||
continue
|
||||
message_audio = getattr(message, "audio", None)
|
||||
if message_audio is None:
|
||||
continue
|
||||
data_b64 = getattr(message_audio, "data", None)
|
||||
if isinstance(data_b64, str) and data_b64:
|
||||
try:
|
||||
payloads.append(base64.b64decode(data_b64))
|
||||
except Exception:
|
||||
continue
|
||||
return payloads
|
||||
|
||||
|
||||
def _decode_data_url(data_url: str) -> tuple[bytes, str]:
|
||||
header, data = data_url.split(",", 1)
|
||||
mime_type = "image/png"
|
||||
if ";" in header and ":" in header:
|
||||
mime_type = header.split(":", 1)[1].split(";", 1)[0]
|
||||
return base64.b64decode(data), mime_type
|
||||
|
||||
|
||||
def _image_extension_from_mime(mime_type: str) -> str:
|
||||
if mime_type == "image/jpeg":
|
||||
return ".jpg"
|
||||
if mime_type == "image/webp":
|
||||
return ".webp"
|
||||
if mime_type == "image/gif":
|
||||
return ".gif"
|
||||
return ".png"
|
||||
|
||||
|
||||
def _save_outputs(
|
||||
*,
|
||||
query_type: str,
|
||||
chat_completion: Any,
|
||||
output_dir: Path,
|
||||
) -> None:
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
stamp = time.strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
text_outputs = _extract_text_outputs(chat_completion)
|
||||
image_data_urls = _extract_image_data_urls(chat_completion)
|
||||
audio_payloads = _extract_audio_payloads(chat_completion)
|
||||
|
||||
if text_outputs:
|
||||
text_path = output_dir / f"{query_type}_{stamp}.txt"
|
||||
text_path.write_text("\n\n".join(text_outputs) + "\n", encoding="utf-8")
|
||||
print(f"[dynin-online] text saved: {text_path}")
|
||||
print(text_outputs[0])
|
||||
|
||||
for idx, image_url in enumerate(image_data_urls):
|
||||
image_bytes, mime_type = _decode_data_url(image_url)
|
||||
ext = _image_extension_from_mime(mime_type)
|
||||
image_path = output_dir / f"{query_type}_{stamp}_{idx}{ext}"
|
||||
image_path.write_bytes(image_bytes)
|
||||
print(f"[dynin-online] image saved: {image_path}")
|
||||
|
||||
for idx, audio_bytes in enumerate(audio_payloads):
|
||||
audio_path = output_dir / f"{query_type}_{stamp}_{idx}.wav"
|
||||
audio_path.write_bytes(audio_bytes)
|
||||
print(f"[dynin-online] audio saved: {audio_path}")
|
||||
|
||||
if not text_outputs and not image_data_urls and not audio_payloads:
|
||||
print("[dynin-online] no output extracted from response")
|
||||
raw_path = output_dir / f"{query_type}_{stamp}_raw_response.json"
|
||||
try:
|
||||
if hasattr(chat_completion, "model_dump_json"):
|
||||
serialized = chat_completion.model_dump_json(indent=2)
|
||||
else:
|
||||
if hasattr(chat_completion, "model_dump"):
|
||||
raw_payload: Any = chat_completion.model_dump(mode="json")
|
||||
else:
|
||||
raw_payload = chat_completion
|
||||
try:
|
||||
serialized = json.dumps(raw_payload, ensure_ascii=False, indent=2)
|
||||
except Exception:
|
||||
serialized = json.dumps({"repr": repr(raw_payload)}, ensure_ascii=False, indent=2)
|
||||
raw_path.write_text(serialized + "\n", encoding="utf-8")
|
||||
print(f"[dynin-online] raw response saved: {raw_path}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _build_offline_parity_sampling_params_list() -> list[dict[str, Any]]:
|
||||
return [dict(OFFLINE_PARITY_STAGE_SAMPLING) for _ in range(OFFLINE_PARITY_STAGE_COUNT)]
|
||||
|
||||
|
||||
def run_request(args: argparse.Namespace) -> None:
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(
|
||||
api_key="EMPTY",
|
||||
base_url=f"http://{args.host}:{args.port}/v1",
|
||||
)
|
||||
prompt = args.prompt.strip() if args.prompt else DEFAULT_PROMPT_BY_QUERY[args.query_type]
|
||||
user_content = _build_user_content(
|
||||
query_type=args.query_type,
|
||||
prompt=prompt,
|
||||
image_path=args.image_path,
|
||||
)
|
||||
if args.modalities:
|
||||
modalities = [item.strip() for item in args.modalities.split(",") if item.strip()]
|
||||
else:
|
||||
modalities = DEFAULT_MODALITIES_BY_QUERY[args.query_type]
|
||||
|
||||
extra_body = {
|
||||
"sampling_params_list": _build_offline_parity_sampling_params_list(),
|
||||
}
|
||||
chat_completion = client.chat.completions.create(
|
||||
model=args.model,
|
||||
messages=[{"role": "user", "content": user_content}],
|
||||
modalities=modalities,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
_save_outputs(
|
||||
query_type=args.query_type,
|
||||
chat_completion=chat_completion,
|
||||
output_dir=Path(args.output_dir).expanduser(),
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Dynin-Omni online chat completion client")
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
type=str,
|
||||
default="t2i",
|
||||
choices=QUERY_CHOICES,
|
||||
help="Dynin query type",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
"-m",
|
||||
type=str,
|
||||
default=DEFAULT_MODEL,
|
||||
help="Model name/path",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--host",
|
||||
type=str,
|
||||
default="localhost",
|
||||
help="Host/IP of the vLLM Omni API server",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--port",
|
||||
type=int,
|
||||
default=8091,
|
||||
help="Port of the vLLM Omni API server",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
"-p",
|
||||
type=str,
|
||||
default="",
|
||||
help="Custom prompt text",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image-path",
|
||||
"-i",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Image path/URL for i2i",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--modalities",
|
||||
type=str,
|
||||
default="",
|
||||
help="Comma-separated output modalities override (e.g., text,image,audio)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
"-o",
|
||||
type=str,
|
||||
default=DEFAULT_OUTPUT_DIR,
|
||||
help="Directory to save outputs",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
|
||||
run_request(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -314,6 +314,7 @@ count, use `size` and `n` rather than `height`, `width`, or
|
||||
| `seed` | int | None | Random seed (reproducible) |
|
||||
| `negative_prompt` | str | None | Negative prompt |
|
||||
| `num_outputs_per_prompt` | int | 1 | Number of images to generate |
|
||||
| `strength` | float | 0.6 | **Z-Image only** - Denoising start timestep for I2I. Range: [0.0, 1.0]. Lower preserves more of original image. |
|
||||
| `layers` | int | 4 | Number of layers (Qwen-Image-Layered) |
|
||||
| `resolution` | int | 640 | Resolution, 640 or 1024 (Qwen-Image-Layered) |
|
||||
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Unit tests for metrics.py
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from vllm.benchmarks.serve import TaskType
|
||||
|
||||
from vllm_omni.benchmarks.metrics.metrics import calculate_metrics
|
||||
from vllm_omni.benchmarks.patch.patch import MixRequestFuncOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_output(prompt_len: int, output_tokens: int = 10) -> MixRequestFuncOutput:
|
||||
"""Build a minimal successful MixRequestFuncOutput for metrics aggregation."""
|
||||
output = MixRequestFuncOutput()
|
||||
output.success = True
|
||||
output.prompt_len = prompt_len
|
||||
output.output_tokens = output_tokens
|
||||
output.generated_text = "x" * output_tokens
|
||||
output.ttft = 0.1
|
||||
output.text_latency = 1.0
|
||||
output.latency = 1.0
|
||||
output.start_time = 0.0
|
||||
output.itl = [0.1] * max(output_tokens - 1, 0)
|
||||
output.audio_ttfp = 0.0
|
||||
output.audio_rtf = 0.0
|
||||
output.audio_duration = 0.0
|
||||
output.audio_frames = 0
|
||||
output.input_audio_duration = 0.0
|
||||
output.error = ""
|
||||
return output
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# total_input Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def test_total_input_aggregated_from_output_prompt_len():
|
||||
"""Test that total_input sums outputs[i].prompt_len, not input_requests[i].prompt_len."""
|
||||
outputs = [_make_output(4992), _make_output(3000)]
|
||||
|
||||
metrics, _ = calculate_metrics(
|
||||
input_requests=[],
|
||||
outputs=outputs,
|
||||
dur_s=10.0,
|
||||
tokenizer=None,
|
||||
selected_percentiles=[99.0],
|
||||
goodput_config_dict={},
|
||||
task_type=TaskType.GENERATION,
|
||||
selected_percentile_metrics=[],
|
||||
max_concurrency=None,
|
||||
request_rate=float("inf"),
|
||||
benchmark_duration=10.0,
|
||||
)
|
||||
|
||||
assert metrics.total_input == 7992, (
|
||||
"total_input should aggregate from outputs[i].prompt_len to reflect the true multimodal input token count"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "-s"])
|
||||
@@ -574,5 +574,59 @@ class TestTextLatencyAttribute:
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# prompt_len Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prompt_len_assigned_from_usage(mocker: MockerFixture):
|
||||
# Arrange: request claims prompt_len=100, but server reports 4992 (multimodal).
|
||||
request_input = RequestFuncInput(
|
||||
model="test-model",
|
||||
model_name="test-model",
|
||||
prompt="test prompt",
|
||||
api_url="http://test.com/v1/chat/completions",
|
||||
prompt_len=100,
|
||||
output_len=20,
|
||||
)
|
||||
|
||||
chunks = [
|
||||
create_sse_chunk(
|
||||
{
|
||||
"choices": [{"delta": {"content": "Hello"}}],
|
||||
"modality": "text",
|
||||
}
|
||||
),
|
||||
create_sse_chunk(
|
||||
{
|
||||
"choices": [{"delta": {"content": " world"}}],
|
||||
"modality": "text",
|
||||
}
|
||||
),
|
||||
# Final usage chunk emitted because stream_options.include_usage=True.
|
||||
create_sse_chunk(
|
||||
{
|
||||
"choices": [],
|
||||
"usage": {"prompt_tokens": 4992, "completion_tokens": 2, "total_tokens": 4994},
|
||||
}
|
||||
),
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
|
||||
mock_response = MockResponse(200, chunks)
|
||||
mock_session = mocker.AsyncMock()
|
||||
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
||||
|
||||
# Act
|
||||
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
||||
|
||||
# Assert
|
||||
assert output.success is True
|
||||
assert output.prompt_len == 4992, (
|
||||
"prompt_len should be overridden by usage.prompt_tokens to reflect the true multimodal input token count"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "-s"])
|
||||
+23
-5
@@ -16,6 +16,7 @@ if "VLLM_TARGET_DEVICE" not in os.environ:
|
||||
os.environ["VLLM_TARGET_DEVICE"] = "cpu"
|
||||
|
||||
import concurrent.futures
|
||||
import contextlib
|
||||
import gc
|
||||
import multiprocessing
|
||||
import socket
|
||||
@@ -52,6 +53,7 @@ from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
|
||||
PromptImageInput = list[Any] | Any | None
|
||||
PromptVideoInput = list[Any] | Any | None
|
||||
@@ -337,10 +339,10 @@ def log_test_name_before_test(request):
|
||||
|
||||
def _run_pre_test_cleanup(enable_force=False):
|
||||
if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force:
|
||||
print("GPU cleanup disabled")
|
||||
print("\nPre-test GPU cleanup skipped(Default off is typical when one worker/instance runs many tests.)\n")
|
||||
return
|
||||
|
||||
print("Pre-test GPU status:")
|
||||
print("\nPre-test GPU status:")
|
||||
|
||||
num_gpus = torch.cuda.device_count()
|
||||
if num_gpus > 0:
|
||||
@@ -1087,6 +1089,22 @@ def _merge_base64_audio_to_segment(base64_list: list[str]):
|
||||
return merged
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _serialize_whisper_small_model_download():
|
||||
"""Serialize Whisper ``small`` cache writes across processes (Linux; ``fcntl``)."""
|
||||
import fcntl
|
||||
|
||||
lock_path = Path.home() / ".cache" / "whisper" / ".small_model_download.lock"
|
||||
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
f = open(lock_path, "a+b")
|
||||
try:
|
||||
fcntl.flock(f.fileno(), fcntl.LOCK_EX)
|
||||
yield
|
||||
finally:
|
||||
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
|
||||
f.close()
|
||||
|
||||
|
||||
def _whisper_transcribe_in_current_process(output_path: str) -> str:
|
||||
import whisper
|
||||
|
||||
@@ -1107,7 +1125,8 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str:
|
||||
else:
|
||||
use_accelerator = False
|
||||
device = "cpu"
|
||||
model = whisper.load_model("small", device=device)
|
||||
with _serialize_whisper_small_model_download():
|
||||
model = whisper.load_model("small", device=device)
|
||||
try:
|
||||
text = model.transcribe(
|
||||
output_path,
|
||||
@@ -1126,8 +1145,7 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str:
|
||||
|
||||
|
||||
def convert_audio_file_to_text(output_path: str) -> str:
|
||||
"""Convert an audio file to text in an isolated subprocess."""
|
||||
# Import locally to avoid impacting test module import time.
|
||||
"""Convert an audio file to text in an isolated subprocess (spawn)."""
|
||||
ctx = multiprocessing.get_context("spawn")
|
||||
with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor:
|
||||
future = executor.submit(_whisper_transcribe_in_current_process, output_path)
|
||||
|
||||
@@ -9,6 +9,10 @@ those requests are permanently orphaned.
|
||||
import unittest
|
||||
from collections import deque
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeAdapter:
|
||||
"""Minimal mock of OmniChunkTransferAdapter tracking restore calls."""
|
||||
|
||||
@@ -0,0 +1,217 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_ltx2_baseline_eager",
|
||||
"description": "Single-device baseline with enforce-eager (no torch.compile)",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Lightricks/LTX-2",
|
||||
"serve_args": {
|
||||
"enforce-eager": true,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "256x256_145f_steps6",
|
||||
"dataset": "random",
|
||||
"task": "t2v",
|
||||
"backend": "v1/videos",
|
||||
"width": 256,
|
||||
"height": 256,
|
||||
"num-frames": 145,
|
||||
"fps": 24,
|
||||
"num-inference-steps": 6,
|
||||
"num-prompts": 3,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true
|
||||
},
|
||||
{
|
||||
"name": "480x768_41f_steps20",
|
||||
"dataset": "random",
|
||||
"task": "t2v",
|
||||
"backend": "v1/videos",
|
||||
"width": 768,
|
||||
"height": 480,
|
||||
"num-frames": 41,
|
||||
"fps": 24,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 3,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true
|
||||
}
|
||||
]
|
||||
},
|
||||
|
||||
{
|
||||
"test_name": "test_ltx2_torch_compile",
|
||||
"description": "Single-device with torch.compile (default, no enforce-eager)",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Lightricks/LTX-2",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "256x256_145f_steps6",
|
||||
"dataset": "random",
|
||||
"task": "t2v",
|
||||
"backend": "v1/videos",
|
||||
"width": 256,
|
||||
"height": 256,
|
||||
"num-frames": 145,
|
||||
"fps": 24,
|
||||
"num-inference-steps": 6,
|
||||
"num-prompts": 3,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true
|
||||
},
|
||||
{
|
||||
"name": "480x768_41f_steps20",
|
||||
"dataset": "random",
|
||||
"task": "t2v",
|
||||
"backend": "v1/videos",
|
||||
"width": 768,
|
||||
"height": 480,
|
||||
"num-frames": 41,
|
||||
"fps": 24,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 3,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true
|
||||
}
|
||||
]
|
||||
},
|
||||
|
||||
{
|
||||
"test_name": "test_ltx2_cfg2_eager",
|
||||
"description": "CFG-parallel=2 with enforce-eager",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Lightricks/LTX-2",
|
||||
"serve_args": {
|
||||
"cfg-parallel-size": 2,
|
||||
"enforce-eager": true,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "256x256_145f_steps6",
|
||||
"dataset": "random",
|
||||
"task": "t2v",
|
||||
"backend": "v1/videos",
|
||||
"width": 256,
|
||||
"height": 256,
|
||||
"num-frames": 145,
|
||||
"fps": 24,
|
||||
"num-inference-steps": 6,
|
||||
"num-prompts": 3,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true
|
||||
},
|
||||
{
|
||||
"name": "480x768_41f_steps20",
|
||||
"dataset": "random",
|
||||
"task": "t2v",
|
||||
"backend": "v1/videos",
|
||||
"width": 768,
|
||||
"height": 480,
|
||||
"num-frames": 41,
|
||||
"fps": 24,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 3,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true
|
||||
}
|
||||
]
|
||||
},
|
||||
|
||||
{
|
||||
"test_name": "test_ltx2_cfg2_compile",
|
||||
"description": "CFG-parallel=2 with torch.compile",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Lightricks/LTX-2",
|
||||
"serve_args": {
|
||||
"cfg-parallel-size": 2,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "256x256_145f_steps6",
|
||||
"dataset": "random",
|
||||
"task": "t2v",
|
||||
"backend": "v1/videos",
|
||||
"width": 256,
|
||||
"height": 256,
|
||||
"num-frames": 145,
|
||||
"fps": 24,
|
||||
"num-inference-steps": 6,
|
||||
"num-prompts": 3,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true
|
||||
},
|
||||
{
|
||||
"name": "480x768_41f_steps20",
|
||||
"dataset": "random",
|
||||
"task": "t2v",
|
||||
"backend": "v1/videos",
|
||||
"width": 768,
|
||||
"height": 480,
|
||||
"num-frames": 41,
|
||||
"fps": 24,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 3,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true
|
||||
}
|
||||
]
|
||||
},
|
||||
|
||||
{
|
||||
"test_name": "test_ltx2_cache_dit_eager",
|
||||
"description": "CacheDiT with enforce-eager",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Lightricks/LTX-2",
|
||||
"serve_args": {
|
||||
"cache-backend": "cache_dit",
|
||||
"enforce-eager": true,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "256x256_145f_steps6",
|
||||
"dataset": "random",
|
||||
"task": "t2v",
|
||||
"backend": "v1/videos",
|
||||
"width": 256,
|
||||
"height": 256,
|
||||
"num-frames": 145,
|
||||
"fps": 24,
|
||||
"num-inference-steps": 6,
|
||||
"num-prompts": 3,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true
|
||||
},
|
||||
{
|
||||
"name": "480x768_41f_steps20",
|
||||
"dataset": "random",
|
||||
"task": "t2v",
|
||||
"backend": "v1/videos",
|
||||
"width": 768,
|
||||
"height": 480,
|
||||
"num-frames": 41,
|
||||
"fps": 24,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 3,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
+7
-1
@@ -21,12 +21,13 @@ from unittest.mock import MagicMock, Mock, patch
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.utils import hardware_test
|
||||
from vllm_omni.diffusion.cache.teacache.extractors import extract_flux2_klein_context
|
||||
from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
|
||||
Flux2Transformer2DModel,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
pytestmark = [pytest.mark.core_model]
|
||||
|
||||
|
||||
@pytest.fixture(scope="function", autouse=True)
|
||||
@@ -113,6 +114,7 @@ class TestFlux2KleinExtractor(BaseExtractorTest):
|
||||
def get_sample_inputs(self, sample_inputs):
|
||||
return sample_inputs
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_modulated_input_shape(self, flux2_klein_module, sample_inputs):
|
||||
"""Test that modulated_input has correct shape matching the model's inner_dim.
|
||||
|
||||
@@ -126,16 +128,19 @@ class TestFlux2KleinExtractor(BaseExtractorTest):
|
||||
inner_dim = flux2_klein_module.inner_dim
|
||||
assert context.modulated_input.shape == (batch_size, img_seq_len, inner_dim)
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_run_transformer_blocks_callable(self, flux2_klein_module, sample_inputs):
|
||||
"""Test that run_transformer_blocks is callable."""
|
||||
context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
|
||||
assert callable(context.run_transformer_blocks)
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_postprocess_callable(self, flux2_klein_module, sample_inputs):
|
||||
"""Test that postprocess is callable."""
|
||||
context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
|
||||
assert callable(context.postprocess)
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_extra_states_contains_full_transformer(self, flux2_klein_module, sample_inputs):
|
||||
"""Test that extra_states contains run_flux2_full_transformer_with_single."""
|
||||
context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
|
||||
@@ -154,6 +159,7 @@ class TestFlux2KleinExtractor(BaseExtractorTest):
|
||||
assert context is not None
|
||||
assert context.temb is not None
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_invalid_module_raises_error(self):
|
||||
"""Test that invalid module without transformer_blocks raises ValueError."""
|
||||
invalid_module = Mock()
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_wan as wan_vae_module
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _DummyOmniAutoencoderKLWan(OmniAutoencoderKLWan):
|
||||
def __init__(self, *, dtype: torch.dtype):
|
||||
torch.nn.Module.__init__(self)
|
||||
self.register_parameter("dummy_weight", torch.nn.Parameter(torch.ones(1, dtype=dtype)))
|
||||
|
||||
|
||||
def test_wan_vae_execution_context_handles_fp32():
|
||||
model = _DummyOmniAutoencoderKLWan(dtype=torch.float32)
|
||||
with model._execution_context():
|
||||
output = model.dummy_weight + 1
|
||||
assert output.dtype == torch.float32
|
||||
|
||||
|
||||
def test_wan_vae_execution_context_handles_bf16():
|
||||
model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
|
||||
with model._execution_context():
|
||||
output = model.dummy_weight + 1
|
||||
assert output.dtype == torch.bfloat16
|
||||
|
||||
|
||||
def test_wan_vae_execution_context_uses_platform_autocast(mocker):
|
||||
sentinel = object()
|
||||
platform = mocker.Mock()
|
||||
platform.create_autocast_context.return_value = sentinel
|
||||
mocker.patch.object(wan_vae_module, "current_omni_platform", platform)
|
||||
|
||||
model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
|
||||
|
||||
assert model._execution_context() is sentinel
|
||||
platform.create_autocast_context.assert_called_once_with(
|
||||
device_type=model.dummy_weight.device.type,
|
||||
dtype=torch.bfloat16,
|
||||
enabled=True,
|
||||
)
|
||||
@@ -11,6 +11,8 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
|
||||
TileTask,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class E2EOperator:
|
||||
"""tiles with (2, 3) -- (H,W)"""
|
||||
|
||||
@@ -17,6 +17,7 @@ import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from tests.utils import hardware_test
|
||||
from vllm_omni.diffusion.attention.parallel.ulysses import (
|
||||
_all_gather_int,
|
||||
_ulysses_all_to_all_any_o,
|
||||
@@ -69,6 +70,8 @@ PERF_CASES: list[_PerfCase] = [
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", PERF_CASES)
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=4)
|
||||
def test_ulysses_advanced_uaa_comm_overhead(case: _PerfCase) -> None:
|
||||
available_gpus = current_omni_platform.get_device_count()
|
||||
if available_gpus < case.world_size:
|
||||
|
||||
@@ -0,0 +1,236 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for BAGEL trajectory recording in the denoising loop."""
|
||||
|
||||
import types
|
||||
from dataclasses import dataclass
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.bagel.bagel_transformer import (
|
||||
Bagel,
|
||||
NaiveCache,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
NUM_TOKENS = 8
|
||||
HIDDEN_DIM = 16
|
||||
NUM_TIMESTEPS = 5
|
||||
# generate_image uses timesteps[:-1], so actual steps = NUM_TIMESTEPS - 1
|
||||
EXPECTED_STEPS = NUM_TIMESTEPS - 1
|
||||
|
||||
|
||||
def _make_mock_bagel():
|
||||
"""Create a mock Bagel with forward returning constant velocity."""
|
||||
mock = MagicMock(spec=Bagel)
|
||||
mock._sp_size = 1
|
||||
|
||||
# forward returns a small constant velocity so x_t changes each step
|
||||
def fake_forward(self, x_t, **kwargs):
|
||||
return torch.ones_like(x_t) * 0.1
|
||||
|
||||
mock.forward = types.MethodType(fake_forward, mock)
|
||||
# _merge_naive_caches is called in the batched CFG path
|
||||
mock._merge_naive_caches = types.MethodType(lambda self, caches: NaiveCache(1), mock)
|
||||
|
||||
# Bind the real generate_image to our mock
|
||||
mock.generate_image = types.MethodType(Bagel.generate_image, mock)
|
||||
return mock
|
||||
|
||||
|
||||
def _make_generate_args(num_tokens=NUM_TOKENS, hidden_dim=HIDDEN_DIM, cfg=False):
|
||||
"""Tensor arguments for generate_image.
|
||||
|
||||
Args:
|
||||
cfg: If True, enable batched CFG path (cfg_text_scale > 1.0).
|
||||
"""
|
||||
seq_len = num_tokens + 2 # packed_seqlens includes 2 extra tokens
|
||||
base = dict(
|
||||
packed_text_ids=torch.zeros(2, dtype=torch.long),
|
||||
packed_text_indexes=torch.tensor([0, 1], dtype=torch.long),
|
||||
packed_init_noises=torch.randn(num_tokens, hidden_dim),
|
||||
packed_vae_position_ids=torch.arange(num_tokens, dtype=torch.long),
|
||||
packed_vae_token_indexes=torch.arange(2, seq_len, dtype=torch.long),
|
||||
packed_seqlens=torch.tensor([seq_len], dtype=torch.int),
|
||||
packed_position_ids=torch.arange(seq_len, dtype=torch.long),
|
||||
packed_indexes=torch.arange(seq_len, dtype=torch.long),
|
||||
past_key_values=NaiveCache(1),
|
||||
key_values_lens=torch.tensor([0], dtype=torch.int),
|
||||
packed_key_value_indexes=torch.zeros(0, dtype=torch.long),
|
||||
num_timesteps=NUM_TIMESTEPS,
|
||||
timestep_shift=1.0,
|
||||
cfg_text_scale=1.0,
|
||||
cfg_img_scale=1.0,
|
||||
)
|
||||
if cfg:
|
||||
base |= dict(
|
||||
cfg_text_scale=4.0,
|
||||
cfg_text_packed_query_indexes=torch.arange(seq_len, dtype=torch.long),
|
||||
cfg_text_packed_position_ids=torch.arange(seq_len, dtype=torch.long),
|
||||
cfg_text_past_key_values=NaiveCache(1),
|
||||
cfg_text_key_values_lens=torch.tensor([0], dtype=torch.int),
|
||||
cfg_text_packed_key_value_indexes=torch.zeros(0, dtype=torch.long),
|
||||
)
|
||||
return base
|
||||
|
||||
|
||||
@pytest.fixture(params=[False, True], ids=["no_cfg", "batched_cfg"])
|
||||
def bagel_and_args(request):
|
||||
"""Mock Bagel instance and generate_image arguments.
|
||||
|
||||
Parametrized over CFG mode so every test runs on both the no-CFG
|
||||
and batched-CFG code paths.
|
||||
"""
|
||||
cfg = request.param
|
||||
with patch(
|
||||
"vllm_omni.diffusion.models.bagel.bagel_transformer.get_classifier_free_guidance_world_size",
|
||||
return_value=1,
|
||||
):
|
||||
yield _make_mock_bagel(), _make_generate_args(cfg=cfg)
|
||||
|
||||
|
||||
class TestTrajectoryRecording:
|
||||
"""Tests for trajectory latent/timestep recording in generate_image."""
|
||||
|
||||
def test_trajectory_disabled_returns_none(self, bagel_and_args):
|
||||
bagel, args = bagel_and_args
|
||||
|
||||
unpacked, trajectory_latents, trajectory_timesteps, trajectory_log_probs = bagel.generate_image(
|
||||
**args, return_trajectory_latents=False
|
||||
)
|
||||
|
||||
assert isinstance(unpacked, (list, tuple))
|
||||
assert len(unpacked) == 1 # one sequence
|
||||
assert trajectory_latents is None
|
||||
assert trajectory_timesteps is None
|
||||
assert trajectory_log_probs is None
|
||||
|
||||
def test_trajectory_enabled_returns_correct_count(self, bagel_and_args):
|
||||
bagel, args = bagel_and_args
|
||||
|
||||
_, trajectory_latents, trajectory_timesteps, trajectory_log_probs = bagel.generate_image(
|
||||
**args, return_trajectory_latents=True
|
||||
)
|
||||
|
||||
assert trajectory_latents is not None
|
||||
assert trajectory_timesteps is not None
|
||||
assert len(trajectory_latents) == EXPECTED_STEPS
|
||||
assert len(trajectory_timesteps) == EXPECTED_STEPS
|
||||
# log_probs is None without a scheduler (default ODE path)
|
||||
assert trajectory_log_probs is None
|
||||
|
||||
def test_trajectory_latents_shape_matches_input(self, bagel_and_args):
|
||||
bagel, args = bagel_and_args
|
||||
expected_shape = args["packed_init_noises"].shape
|
||||
|
||||
_, trajectory_latents, *_ = bagel.generate_image(**args, return_trajectory_latents=True)
|
||||
|
||||
for i, lat in enumerate(trajectory_latents):
|
||||
assert lat.shape == expected_shape, f"Step {i}: expected {expected_shape}, got {lat.shape}"
|
||||
|
||||
def test_trajectory_latents_are_distinct(self, bagel_and_args):
|
||||
bagel, args = bagel_and_args
|
||||
|
||||
_, trajectory_latents, *_ = bagel.generate_image(**args, return_trajectory_latents=True)
|
||||
|
||||
for i in range(1, len(trajectory_latents)):
|
||||
assert not torch.equal(trajectory_latents[i], trajectory_latents[i - 1]), (
|
||||
f"Steps {i - 1} and {i} should differ"
|
||||
)
|
||||
|
||||
def test_trajectory_timesteps_are_decreasing(self, bagel_and_args):
|
||||
bagel, args = bagel_and_args
|
||||
|
||||
_, _, trajectory_timesteps, _ = bagel.generate_image(**args, return_trajectory_latents=True)
|
||||
|
||||
for i in range(1, len(trajectory_timesteps)):
|
||||
assert trajectory_timesteps[i] < trajectory_timesteps[i - 1], (
|
||||
f"Timestep {i} ({trajectory_timesteps[i]:.4f}) should be less than "
|
||||
f"timestep {i - 1} ({trajectory_timesteps[i - 1]:.4f})"
|
||||
)
|
||||
|
||||
def test_trajectory_final_latent_matches_output(self, bagel_and_args):
|
||||
bagel, args = bagel_and_args
|
||||
|
||||
unpacked, trajectory_latents, *_ = bagel.generate_image(**args, return_trajectory_latents=True)
|
||||
|
||||
# Reconstruct the full final latent from unpacked pieces
|
||||
final_latent = torch.cat(unpacked, dim=0)
|
||||
assert torch.allclose(trajectory_latents[-1], final_latent, atol=1e-6), (
|
||||
"Last trajectory latent should match the final output"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Mock scheduler for log-prob tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class _MockStepOutput:
|
||||
prev_sample: torch.Tensor
|
||||
log_prob: torch.Tensor
|
||||
|
||||
|
||||
class _MockScheduler:
|
||||
"""Minimal scheduler: Euler step + constant log-prob per step."""
|
||||
|
||||
def step(self, model_output, sigma, sample, dt, **kwargs):
|
||||
prev_sample = sample - model_output * dt
|
||||
log_prob = torch.tensor(-1.0)
|
||||
return _MockStepOutput(prev_sample=prev_sample, log_prob=log_prob)
|
||||
|
||||
|
||||
class TestTrajectoryLogProbs:
|
||||
"""Tests for log-prob recording when a scheduler is provided."""
|
||||
|
||||
@pytest.fixture()
|
||||
def bagel_scheduler_args(self):
|
||||
with patch(
|
||||
"vllm_omni.diffusion.models.bagel.bagel_transformer.get_classifier_free_guidance_world_size",
|
||||
return_value=1,
|
||||
):
|
||||
yield _make_mock_bagel(), _make_generate_args(), _MockScheduler()
|
||||
|
||||
def test_log_probs_recorded_with_scheduler(self, bagel_scheduler_args):
|
||||
bagel, args, scheduler = bagel_scheduler_args
|
||||
|
||||
_, _, _, trajectory_log_probs = bagel.generate_image(
|
||||
**args, return_trajectory_latents=True, scheduler=scheduler
|
||||
)
|
||||
|
||||
assert trajectory_log_probs is not None
|
||||
assert len(trajectory_log_probs) == EXPECTED_STEPS
|
||||
|
||||
def test_log_probs_are_finite(self, bagel_scheduler_args):
|
||||
bagel, args, scheduler = bagel_scheduler_args
|
||||
|
||||
_, _, _, trajectory_log_probs = bagel.generate_image(
|
||||
**args, return_trajectory_latents=True, scheduler=scheduler
|
||||
)
|
||||
|
||||
for i, lp in enumerate(trajectory_log_probs):
|
||||
assert torch.isfinite(lp).all(), f"Step {i}: log_prob is not finite"
|
||||
|
||||
def test_log_probs_none_without_scheduler(self, bagel_scheduler_args):
|
||||
bagel, args, _ = bagel_scheduler_args
|
||||
|
||||
_, _, _, trajectory_log_probs = bagel.generate_image(**args, return_trajectory_latents=True, scheduler=None)
|
||||
|
||||
assert trajectory_log_probs is None
|
||||
|
||||
def test_scheduler_updates_latents(self, bagel_scheduler_args):
|
||||
"""Verify the scheduler's prev_sample is used (not the raw Euler step)."""
|
||||
bagel, args, scheduler = bagel_scheduler_args
|
||||
|
||||
_, traj_with_sched, *_ = bagel.generate_image(**args, return_trajectory_latents=True, scheduler=scheduler)
|
||||
_, traj_without, *_ = bagel.generate_image(**args, return_trajectory_latents=True, scheduler=None)
|
||||
|
||||
# Mock scheduler does the same Euler step, so latents should match
|
||||
for i in range(len(traj_with_sched)):
|
||||
assert torch.allclose(traj_with_sched[i], traj_without[i], atol=1e-5), (
|
||||
f"Step {i}: scheduler and ODE paths should produce same latents"
|
||||
)
|
||||
@@ -3,6 +3,7 @@ from unittest.mock import MagicMock, patch
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.utils import hardware_test
|
||||
from vllm_omni.diffusion.models.flux2.flux2_transformer import (
|
||||
Flux2PosEmbed,
|
||||
Flux2Transformer2DModel,
|
||||
@@ -24,6 +25,8 @@ def setup_tp_group():
|
||||
class TestFlux2TransformerWeightLoading:
|
||||
"""Test Flux2Transformer weight loading functionality"""
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_weight_loading_tp2(self, setup_tp_group):
|
||||
"""Verify weights load correctly with TP=2"""
|
||||
# Prepare test data
|
||||
@@ -78,6 +81,8 @@ class TestFlux2TransformerWeightLoading:
|
||||
class TestFlux2RopePositionEmbedding:
|
||||
"""Test Flux2 RoPE position embedding functionality"""
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_rope_position_embedding(self):
|
||||
"""Verify RoPE produces correct embeddings for 4D coordinates"""
|
||||
# Prepare test data - use model default configuration
|
||||
@@ -132,6 +137,8 @@ class TestFlux2RopePositionEmbedding:
|
||||
class TestFlux2PackedModuleMapping:
|
||||
"""Test Flux2 packed module mapping functionality"""
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_packed_module_mapping(self, setup_tp_group):
|
||||
"""Verify to_qkv packing matches HF checkpoint"""
|
||||
model = Flux2Transformer2DModel(
|
||||
@@ -208,6 +215,8 @@ class TestFlux2PackedModuleMapping:
|
||||
f"add_kv_proj weight dimension should be {expected_add_kv_shape}, got {attn_block.add_kv_proj.weight.shape}"
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_packed_mapping_edge_cases(self, setup_tp_group):
|
||||
"""Test edge cases for packed mapping"""
|
||||
model = Flux2Transformer2DModel(
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_pipeline(sequence_parallel_size: int = 1) -> LTX2Pipeline:
|
||||
pipeline = object.__new__(LTX2Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.audio_vae_temporal_compression_ratio = 4
|
||||
pipeline.audio_vae_mel_compression_ratio = 4
|
||||
pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size))
|
||||
# Mock audio_vae with identity normalization (mean=0, std=1) so
|
||||
# _normalize_audio_latents is a no-op and test values are preserved.
|
||||
pipeline.audio_vae = SimpleNamespace(
|
||||
latents_mean=torch.tensor(0.0),
|
||||
latents_std=torch.tensor(1.0),
|
||||
)
|
||||
return pipeline
|
||||
|
||||
|
||||
def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
|
||||
pipeline = _make_pipeline(sequence_parallel_size=4)
|
||||
latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
|
||||
|
||||
padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
|
||||
batch_size=1,
|
||||
num_channels_latents=2,
|
||||
num_mel_bins=8,
|
||||
audio_latent_length=10,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
latents=latents,
|
||||
)
|
||||
|
||||
assert original_num_frames == 10
|
||||
assert padded_num_frames == 12
|
||||
assert padded.shape == (1, 12, 4)
|
||||
torch.testing.assert_close(padded[:, :10], latents)
|
||||
torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4))
|
||||
|
||||
|
||||
def test_unpad_audio_latents_restores_original_frames_before_unpack():
|
||||
pipeline = _make_pipeline()
|
||||
original = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
|
||||
padded = torch.cat([original, torch.full((1, 2, 4), 999.0)], dim=1)
|
||||
|
||||
unpadded = pipeline._unpad_audio_latents(padded, 10)
|
||||
unpacked = pipeline._unpack_audio_latents(unpadded, latent_length=10, num_mel_bins=2)
|
||||
expected = pipeline._unpack_audio_latents(original, latent_length=10, num_mel_bins=2)
|
||||
|
||||
assert unpacked.shape == (1, 2, 10, 2)
|
||||
assert not (unpacked == 999.0).any()
|
||||
torch.testing.assert_close(unpacked, expected)
|
||||
@@ -5,7 +5,7 @@
|
||||
import pytest
|
||||
from torch import nn
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_build_quant_config_fp8():
|
||||
|
||||
@@ -8,9 +8,10 @@ import pytest
|
||||
import torch
|
||||
|
||||
import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
|
||||
from tests.utils import hardware_test
|
||||
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
pytestmark = [pytest.mark.diffusion]
|
||||
|
||||
|
||||
@contextmanager
|
||||
@@ -64,6 +65,8 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
|
||||
return runner
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypatch):
|
||||
"""Guard cache diagnostics with runtime backend state to avoid stale-config crashes."""
|
||||
runner = _make_runner(cache_backend=None, cache_backend_name="cache_dit")
|
||||
@@ -84,6 +87,8 @@ def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypa
|
||||
assert cache_summary_calls == []
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeypatch):
|
||||
class _EnabledCacheBackend:
|
||||
def is_enabled(self):
|
||||
@@ -107,6 +112,8 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp
|
||||
assert cache_summary_calls == [(runner.pipeline, True)]
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
|
||||
class _DummyLoader:
|
||||
def __init__(self, load_config, od_config=None):
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import asdict
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
|
||||
async def run_test():
|
||||
captured = {}
|
||||
|
||||
def step(request):
|
||||
captured["request"] = request
|
||||
return [
|
||||
SimpleNamespace(
|
||||
images=["img-1"],
|
||||
_multimodal_output={},
|
||||
_custom_output={},
|
||||
metrics={},
|
||||
stage_durations={},
|
||||
peak_memory_mb=0.0,
|
||||
latents=None,
|
||||
trajectory_latents=None,
|
||||
trajectory_timesteps=None,
|
||||
trajectory_log_probs=None,
|
||||
trajectory_decoded=None,
|
||||
final_output_type="image",
|
||||
),
|
||||
SimpleNamespace(
|
||||
images=["img-2"],
|
||||
_multimodal_output={},
|
||||
_custom_output={},
|
||||
metrics={},
|
||||
stage_durations={},
|
||||
peak_memory_mb=0.0,
|
||||
latents=None,
|
||||
trajectory_latents=None,
|
||||
trajectory_timesteps=None,
|
||||
trajectory_log_probs=None,
|
||||
trajectory_decoded=None,
|
||||
final_output_type="image",
|
||||
),
|
||||
]
|
||||
|
||||
proc = object.__new__(StageDiffusionProc)
|
||||
proc._engine = SimpleNamespace(step=step)
|
||||
proc._executor = ThreadPoolExecutor(max_workers=1)
|
||||
|
||||
try:
|
||||
result = await proc._process_batch_request(
|
||||
request_id="req-parent",
|
||||
prompts=["hello", "world"],
|
||||
sampling_params_dict=asdict(OmniDiffusionSamplingParams()),
|
||||
kv_sender_info={0: {"host": "10.0.0.2", "zmq_port": 50151}},
|
||||
)
|
||||
finally:
|
||||
proc._executor.shutdown(wait=True)
|
||||
|
||||
request = captured["request"]
|
||||
assert request.request_id == "req-parent"
|
||||
assert request.request_ids == ["req-parent-0", "req-parent-1"]
|
||||
assert request.kv_sender_info == {0: {"host": "10.0.0.2", "zmq_port": 50151}}
|
||||
assert result.request_id == "req-parent"
|
||||
assert result.images == ["img-1", "img-2"]
|
||||
|
||||
asyncio.run(run_test())
|
||||
@@ -9,7 +9,7 @@ from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
|
||||
from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
|
||||
from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
|
||||
|
||||
# pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_basic_serialization():
|
||||
@@ -120,3 +120,61 @@ def test_get_invalid_metadata(shm_connector):
|
||||
|
||||
result = shm_connector.get("stage_0", "stage_1", "req_3", {"unknown": "format"})
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_mooncake_connector_defaults_missing_host_to_detected_ip(monkeypatch: pytest.MonkeyPatch):
|
||||
import vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector as mooncake_module
|
||||
|
||||
class _FakePool:
|
||||
is_cuda = False
|
||||
|
||||
def pin_memory(self):
|
||||
return self
|
||||
|
||||
def data_ptr(self):
|
||||
return 1234
|
||||
|
||||
class _FakeTransferEngine:
|
||||
def initialize(self, host, mode, protocol, device_name):
|
||||
self.host = host
|
||||
self.mode = mode
|
||||
self.protocol = protocol
|
||||
self.device_name = device_name
|
||||
return 0
|
||||
|
||||
def get_rpc_port(self):
|
||||
return 23456
|
||||
|
||||
def register_memory(self, base_ptr, pool_size):
|
||||
del base_ptr, pool_size
|
||||
return 0
|
||||
|
||||
def unregister_memory(self, base_ptr):
|
||||
del base_ptr
|
||||
return 0
|
||||
|
||||
monkeypatch.setattr(mooncake_module, "TransferEngine", _FakeTransferEngine)
|
||||
monkeypatch.setattr(mooncake_module.torch, "empty", lambda *args, **kwargs: _FakePool())
|
||||
monkeypatch.setattr(
|
||||
mooncake_module.MooncakeTransferEngineConnector,
|
||||
"_get_local_ip",
|
||||
lambda self: "10.20.30.40",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
mooncake_module.MooncakeTransferEngineConnector,
|
||||
"_zmq_listener_loop",
|
||||
lambda self: self._listener_ready.set(),
|
||||
)
|
||||
|
||||
connector = mooncake_module.MooncakeTransferEngineConnector(
|
||||
{
|
||||
"zmq_port": 50051,
|
||||
"memory_pool_size": 4096,
|
||||
}
|
||||
)
|
||||
try:
|
||||
assert connector.host == "10.20.30.40"
|
||||
assert connector.engine.host == "10.20.30.40"
|
||||
assert connector.get_connection_info()["host"] == "10.20.30.40"
|
||||
finally:
|
||||
connector.close()
|
||||
@@ -1,8 +1,14 @@
|
||||
import json
|
||||
import struct
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm_omni.distributed.omni_connectors.kv_transfer_manager as kv_transfer_manager_module
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
|
||||
KVCacheTransferData,
|
||||
OmniKVCacheConfig,
|
||||
OmniKVTransferManager,
|
||||
)
|
||||
@@ -60,6 +66,35 @@ def common_constants():
|
||||
}
|
||||
|
||||
|
||||
def _decode_stored_payload(data):
|
||||
if isinstance(data, torch.Tensor) and data.dtype == torch.uint8 and data.dim() == 1:
|
||||
return KVCacheTransferData.from_bytes(data.cpu().numpy().tobytes())
|
||||
|
||||
if isinstance(data, (bytes, bytearray, memoryview)):
|
||||
return KVCacheTransferData.from_bytes(data)
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def _make_serialized_payload() -> tuple[bytes, torch.Tensor]:
|
||||
key_tensor = torch.arange(12, dtype=torch.float32).reshape(3, 4)
|
||||
payload = KVCacheTransferData(
|
||||
request_id="req-payload",
|
||||
layer_blocks={"key_cache": [key_tensor], "value_cache": [None]},
|
||||
block_ids=[1],
|
||||
metadata={"seq_len": 3},
|
||||
).to_bytes()
|
||||
return payload, key_tensor
|
||||
|
||||
|
||||
def _rewrite_serialized_header(payload: bytes, mutate_header) -> bytes:
|
||||
header_len = struct.unpack(">I", payload[:4])[0]
|
||||
header = json.loads(payload[4 : 4 + header_len])
|
||||
mutate_header(header)
|
||||
new_header = json.dumps(header, separators=(",", ":")).encode("utf-8")
|
||||
return struct.pack(">I", len(new_header)) + new_header + payload[4 + header_len :]
|
||||
|
||||
|
||||
def test_manager_extraction(kv_config, mock_connector, common_constants):
|
||||
"""Test extraction and sending logic in OmniKVTransferManager."""
|
||||
num_layers = common_constants["num_layers"]
|
||||
@@ -95,7 +130,7 @@ def test_manager_extraction(kv_config, mock_connector, common_constants):
|
||||
expected_key = f"stage1->stage2:{full_request_id}"
|
||||
assert expected_key in mock_connector.store
|
||||
|
||||
data = mock_connector.store[expected_key]
|
||||
data = _decode_stored_payload(mock_connector.store[expected_key])
|
||||
assert data["request_id"] == req_id
|
||||
assert "layer_blocks" in data
|
||||
assert len(data["layer_blocks"]["key_cache"]) == num_layers
|
||||
@@ -106,6 +141,116 @@ def test_manager_extraction(kv_config, mock_connector, common_constants):
|
||||
assert data["layer_blocks"]["key_cache"][0].shape == expected_shape
|
||||
|
||||
|
||||
def test_from_bytes_rejects_out_of_bounds_header_len():
|
||||
payload, _ = _make_serialized_payload()
|
||||
bad_payload = struct.pack(">I", len(payload)) + payload[4:]
|
||||
|
||||
with pytest.raises(ValueError, match="header_len"):
|
||||
KVCacheTransferData.from_bytes(bad_payload)
|
||||
|
||||
with pytest.raises(ValueError, match="header_len"):
|
||||
KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8))
|
||||
|
||||
|
||||
def test_from_bytes_rejects_out_of_bounds_tensor_span():
|
||||
payload, _ = _make_serialized_payload()
|
||||
bad_payload = _rewrite_serialized_header(payload, lambda header: header["td"][0].update({"o": 4096}))
|
||||
|
||||
with pytest.raises(ValueError, match="tensor span"):
|
||||
KVCacheTransferData.from_bytes(bad_payload)
|
||||
|
||||
with pytest.raises(ValueError, match="tensor span"):
|
||||
KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8))
|
||||
|
||||
|
||||
def test_from_bytes_rejects_unsupported_dtype():
|
||||
payload, _ = _make_serialized_payload()
|
||||
bad_payload = _rewrite_serialized_header(payload, lambda header: header["td"][0].update({"d": "cuda"}))
|
||||
|
||||
with pytest.raises(ValueError, match="Unsupported dtype"):
|
||||
KVCacheTransferData.from_bytes(bad_payload)
|
||||
|
||||
with pytest.raises(ValueError, match="Unsupported dtype"):
|
||||
KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8))
|
||||
|
||||
|
||||
def test_from_bytes_uses_explicit_layer_index_descriptor():
|
||||
payload, key_tensor = _make_serialized_payload()
|
||||
payload_with_explicit_index = _rewrite_serialized_header(
|
||||
payload,
|
||||
lambda header: header["td"][0].update({"n": "key_cache_extra_suffix", "i": 0}),
|
||||
)
|
||||
|
||||
data = KVCacheTransferData.from_bytes(payload_with_explicit_index)
|
||||
|
||||
assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor)
|
||||
|
||||
|
||||
def test_update_sender_info_uses_configured_source_stage():
|
||||
config = OmniKVCacheConfig(
|
||||
connector_config={"type": "mock"},
|
||||
stage_id=2,
|
||||
engine_input_source=[1],
|
||||
need_recv_cache=True,
|
||||
)
|
||||
manager = OmniKVTransferManager(config)
|
||||
|
||||
manager.update_sender_info(
|
||||
{
|
||||
0: {"host": "10.0.0.1", "zmq_port": 50151},
|
||||
1: {"host": "10.0.0.2", "zmq_port": 50152},
|
||||
}
|
||||
)
|
||||
|
||||
assert manager.config.connector_config["sender_host"] == "10.0.0.2"
|
||||
assert manager.config.connector_config["sender_zmq_port"] == 50152
|
||||
|
||||
|
||||
def test_clone_received_payload_tensors_breaks_buffer_alias():
|
||||
payload, key_tensor = _make_serialized_payload()
|
||||
raw = np.frombuffer(bytearray(payload), dtype=np.uint8)
|
||||
data = KVCacheTransferData.from_bytes(memoryview(raw))
|
||||
|
||||
OmniKVTransferManager._clone_received_payload_tensors(data)
|
||||
raw[:] = 0
|
||||
|
||||
assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor)
|
||||
|
||||
|
||||
def test_receive_kv_cache_uses_exponential_backoff(monkeypatch):
|
||||
config = OmniKVCacheConfig(
|
||||
connector_config={"type": "mock"},
|
||||
from_stage="sender",
|
||||
stage_id="receiver",
|
||||
need_recv_cache=True,
|
||||
recv_timeout=0.3,
|
||||
)
|
||||
manager = OmniKVTransferManager(config)
|
||||
|
||||
class _NeverReadyConnector:
|
||||
def get(self, **kwargs):
|
||||
del kwargs
|
||||
return None
|
||||
|
||||
manager._connector = _NeverReadyConnector()
|
||||
|
||||
now = {"value": 0.0}
|
||||
sleep_intervals = []
|
||||
|
||||
monkeypatch.setattr(kv_transfer_manager_module.time, "time", lambda: now["value"])
|
||||
|
||||
def _fake_sleep(interval: float) -> None:
|
||||
sleep_intervals.append(interval)
|
||||
now["value"] += interval
|
||||
|
||||
monkeypatch.setattr(kv_transfer_manager_module.time, "sleep", _fake_sleep)
|
||||
|
||||
data, size = manager.receive_kv_cache_for_request("req-backoff")
|
||||
|
||||
assert (data, size) == (None, 0)
|
||||
assert sleep_intervals == pytest.approx([0.01, 0.02, 0.04, 0.08, 0.16])
|
||||
|
||||
|
||||
def test_manager_extraction_tuple_layout(kv_config, mock_connector, common_constants):
|
||||
"""Test extraction with tuple layout."""
|
||||
num_layers = common_constants["num_layers"]
|
||||
@@ -135,7 +280,7 @@ def test_manager_extraction_tuple_layout(kv_config, mock_connector, common_const
|
||||
expected_key = f"stage1->stage2:{full_request_id}"
|
||||
assert expected_key in mock_connector.store
|
||||
|
||||
data = mock_connector.store[expected_key]
|
||||
data = _decode_stored_payload(mock_connector.store[expected_key])
|
||||
expected_shape = (seq_len, num_heads, head_dim)
|
||||
for idx in range(len(kv_caches)):
|
||||
assert data["layer_blocks"]["key_cache"][idx].shape == expected_shape
|
||||
@@ -165,7 +310,7 @@ def test_manager_extraction_mismatched_kv_block_counts(kv_config, mock_connector
|
||||
expected_key = f"stage1->stage2:{full_request_id}"
|
||||
assert expected_key in mock_connector.store
|
||||
|
||||
data = mock_connector.store[expected_key]
|
||||
data = _decode_stored_payload(mock_connector.store[expected_key])
|
||||
expected_shape = (2 * block_size, num_heads, head_dim)
|
||||
assert data["layer_blocks"]["key_cache"][0].shape == expected_shape
|
||||
assert data["layer_blocks"]["value_cache"][0].shape == expected_shape
|
||||
@@ -254,6 +399,82 @@ def test_manager_reception(kv_config, mock_connector, common_constants):
|
||||
assert req.kv_metadata["seq_len"] == seq_len
|
||||
|
||||
|
||||
def test_manager_reception_prefers_parent_request_id_for_batched_request(kv_config, mock_connector, common_constants):
|
||||
"""Batched diffusion requests must fetch KV using the parent/global request ID."""
|
||||
num_layers = common_constants["num_layers"]
|
||||
num_heads = common_constants["num_heads"]
|
||||
head_dim = common_constants["head_dim"]
|
||||
seq_len = common_constants["seq_len"]
|
||||
parent_req_id = common_constants["req_id"]
|
||||
|
||||
expected_shape = (seq_len, num_heads, head_dim)
|
||||
key_cache = [torch.randn(expected_shape) for _ in range(num_layers)]
|
||||
value_cache = [torch.randn(expected_shape) for _ in range(num_layers)]
|
||||
|
||||
data_to_receive = {
|
||||
"request_id": parent_req_id,
|
||||
"layer_blocks": {"key_cache": key_cache, "value_cache": value_cache},
|
||||
"metadata": {"seq_len": seq_len},
|
||||
"block_ids": [],
|
||||
}
|
||||
|
||||
manager = OmniKVTransferManager(kv_config)
|
||||
manager._connector = mock_connector
|
||||
|
||||
full_request_id = f"omni_stage1_to_stage2_kv_cache_{parent_req_id}"
|
||||
store_key = f"stage1->stage2:{full_request_id}"
|
||||
mock_connector.store[store_key] = data_to_receive
|
||||
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=["prompt-a", "prompt-b"],
|
||||
sampling_params=OmniDiffusionSamplingParams(),
|
||||
request_ids=[f"{parent_req_id}-0", f"{parent_req_id}-1"],
|
||||
request_id=parent_req_id,
|
||||
)
|
||||
|
||||
success = manager.receive_kv_cache(req, target_device=torch.device("cpu"))
|
||||
|
||||
assert success
|
||||
assert req.kv_metadata["seq_len"] == seq_len
|
||||
assert torch.allclose(req.past_key_values.key_cache[0], key_cache[0])
|
||||
|
||||
|
||||
def test_receive_multi_kv_cache_uses_parent_request_id_for_cfg_collection(kv_config):
|
||||
manager = OmniKVTransferManager(kv_config)
|
||||
|
||||
seen = {}
|
||||
|
||||
def collect_cfg(request_id, cfg_request_ids, kv_transfer_manager, target_device):
|
||||
seen["request_id"] = request_id
|
||||
seen["cfg_request_ids"] = cfg_request_ids
|
||||
seen["kv_transfer_manager"] = kv_transfer_manager
|
||||
seen["target_device"] = target_device
|
||||
return {"cfg_text_kv_metadata": {"ok": True}}
|
||||
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=["prompt-a", "prompt-b"],
|
||||
sampling_params=OmniDiffusionSamplingParams(),
|
||||
request_ids=["req-parent-0", "req-parent-1"],
|
||||
request_id="req-parent",
|
||||
)
|
||||
req.sampling_params.cfg_kv_request_ids = {"cfg_text": "req-parent__cfg_text"}
|
||||
|
||||
manager.receive_kv_cache = lambda request, target_device=None: request is req
|
||||
|
||||
success = manager.receive_multi_kv_cache(
|
||||
req,
|
||||
cfg_kv_collect_func=collect_cfg,
|
||||
target_device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
assert success
|
||||
assert seen["request_id"] == "req-parent"
|
||||
assert seen["cfg_request_ids"] == {"cfg_text": "req-parent__cfg_text"}
|
||||
assert seen["kv_transfer_manager"] is manager
|
||||
assert seen["target_device"] == torch.device("cpu")
|
||||
assert req.sampling_params.cfg_text_kv_metadata == {"ok": True}
|
||||
|
||||
|
||||
def test_integration_flow(common_constants):
|
||||
"""Simulate extraction -> connector -> reception."""
|
||||
num_layers = common_constants["num_layers"]
|
||||
|
||||
@@ -3,12 +3,16 @@
|
||||
|
||||
from time import time
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.distributed.omni_coordinator import (
|
||||
InstanceInfo,
|
||||
RandomBalancer,
|
||||
StageStatus,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_load_balancer_select_returns_valid_index():
|
||||
"""Verify RandomBalancer.select() returns a valid index for instances."""
|
||||
|
||||
@@ -12,6 +12,8 @@ from vllm_omni.distributed.omni_coordinator import (
|
||||
OmniCoordClientForHub,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _bind_pub() -> tuple[zmq.Context, zmq.Socket, str]:
|
||||
ctx = zmq.Context.instance()
|
||||
|
||||
@@ -17,6 +17,8 @@ from vllm_omni.distributed.omni_coordinator import (
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _bind_router() -> tuple[zmq.Context, zmq.Socket, str]:
|
||||
ctx = zmq.Context.instance()
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
import json
|
||||
import time
|
||||
|
||||
import pytest
|
||||
import zmq
|
||||
from vllm.v1.utils import get_engine_client_zmq_addr
|
||||
|
||||
@@ -13,6 +14,8 @@ from vllm_omni.distributed.omni_coordinator import (
|
||||
StageStatus,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
|
||||
"""Receive InstanceList JSON from SUB socket. Returns None on timeout."""
|
||||
|
||||
@@ -6,7 +6,8 @@
|
||||
This pipeline follows the structure of the user's reference implementation:
|
||||
- supports pre-tokenized prompt IDs via OmniCustomPrompt-style dict input
|
||||
- uses an SDE scheduler that can return step logprobs
|
||||
- returns rich custom_output fields for testing
|
||||
- returns structured trajectory_* fields (latents, timesteps, log_probs)
|
||||
consistent with the BAGEL trajectory recording design
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -393,10 +394,10 @@ class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
|
||||
|
||||
return DiffusionOutput(
|
||||
output=_maybe_to_cpu(image),
|
||||
trajectory_latents=_maybe_to_cpu(all_latents),
|
||||
trajectory_log_probs=_maybe_to_cpu(all_log_probs),
|
||||
trajectory_timesteps=_maybe_to_cpu(all_timesteps),
|
||||
custom_output={
|
||||
"all_latents": _maybe_to_cpu(all_latents),
|
||||
"all_log_probs": _maybe_to_cpu(all_log_probs),
|
||||
"all_timesteps": _maybe_to_cpu(all_timesteps),
|
||||
"prompt_embeds": _maybe_to_cpu(prompt_embeds),
|
||||
"prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask),
|
||||
"negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds),
|
||||
|
||||
@@ -1,7 +1,12 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""E2E tests for AsyncOmni Qwen-Image generation flow (no Ray, no HTTP server)."""
|
||||
"""E2E tests for AsyncOmni Qwen-Image generation with trajectory_* fields.
|
||||
|
||||
Validates that the custom Qwen-Image pipeline returns structured trajectory
|
||||
outputs (latents, timesteps, log_probs) via OmniRequestOutput's trajectory_*
|
||||
fields instead of the legacy custom_output dict.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -191,10 +196,17 @@ async def test_async_omni_generate_with_logprobs():
|
||||
|
||||
_assert_valid_image_output(output)
|
||||
|
||||
all_log_probs = output.custom_output.get("all_log_probs")
|
||||
assert all_log_probs is not None, "all_log_probs should be present when logprobs=True"
|
||||
assert hasattr(all_log_probs, "shape")
|
||||
assert all_log_probs.numel() > 0
|
||||
assert output.trajectory_latents is not None, "trajectory_latents should be present"
|
||||
assert hasattr(output.trajectory_latents, "shape")
|
||||
assert output.trajectory_latents.numel() > 0
|
||||
|
||||
assert output.trajectory_timesteps is not None, "trajectory_timesteps should be present"
|
||||
assert hasattr(output.trajectory_timesteps, "shape")
|
||||
assert output.trajectory_timesteps.numel() > 0
|
||||
|
||||
assert output.trajectory_log_probs is not None, "trajectory_log_probs should be present when logprobs=True"
|
||||
assert hasattr(output.trajectory_log_probs, "shape")
|
||||
assert output.trajectory_log_probs.numel() > 0
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
|
||||
@@ -0,0 +1,198 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
End-to-end test for BAGEL LoRA support (Stage 1 / DiT).
|
||||
|
||||
Validates that LoRA adapters are correctly loaded, applied with controllable
|
||||
scale, and cleanly deactivated. Uses a synthetic rank-1 adapter targeting the
|
||||
first decoder layer's QKV projection.
|
||||
|
||||
Assertions:
|
||||
(a) LoRA at scale=1.0 visibly changes the output (diff > 0.5)
|
||||
(b) scale=2.0 produces a larger delta than scale=1.0 (linearity)
|
||||
(c) The delta is bounded (diff < 80, not corrupted)
|
||||
(d) Deactivating LoRA exactly restores the baseline (diff == 0)
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
from vllm_omni.inputs.data import OmniSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
from safetensors.torch import save_file
|
||||
|
||||
from tests.conftest import modify_stage_config
|
||||
from tests.utils import hardware_test
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.lora.request import LoRARequest
|
||||
from vllm_omni.lora.utils import stable_lora_int_id
|
||||
|
||||
MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
|
||||
BAGEL_STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
|
||||
DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers (reused from test_bagel_text2img.py patterns)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _resolve_stage_config(config_path: str, run_level: str) -> str:
|
||||
if run_level == "advanced_model":
|
||||
return modify_stage_config(
|
||||
config_path,
|
||||
deletes={
|
||||
"stage_args": {
|
||||
0: ["engine_args.load_format"],
|
||||
1: ["engine_args.load_format"],
|
||||
}
|
||||
},
|
||||
)
|
||||
return config_path
|
||||
|
||||
|
||||
def _configure_sampling_params(omni: Omni, num_inference_steps: int = 10) -> list[OmniSamplingParams]:
|
||||
params_list = omni.default_sampling_params_list
|
||||
if len(params_list) > 1:
|
||||
params_list[1].num_inference_steps = num_inference_steps
|
||||
params_list[1].extra_args = {
|
||||
"cfg_text_scale": 4.0,
|
||||
"cfg_img_scale": 1.5,
|
||||
}
|
||||
return params_list
|
||||
|
||||
|
||||
def _extract_generated_image(omni_outputs: list[OmniRequestOutput]) -> Image.Image | None:
|
||||
for req_output in omni_outputs:
|
||||
if req_output.images:
|
||||
return req_output.images[0]
|
||||
return None
|
||||
|
||||
|
||||
def _generate_bagel_image(omni: Omni) -> Image.Image:
|
||||
params_list = _configure_sampling_params(omni)
|
||||
params_list[1].lora_request = None
|
||||
outputs = list(
|
||||
omni.generate(
|
||||
prompts=[{"prompt": DEFAULT_PROMPT, "modalities": ["image"]}],
|
||||
sampling_params_list=params_list,
|
||||
)
|
||||
)
|
||||
img = _extract_generated_image(outputs)
|
||||
assert img is not None, "No image generated"
|
||||
return img
|
||||
|
||||
|
||||
def _generate_bagel_image_with_lora(
|
||||
omni: Omni,
|
||||
lora_request: LoRARequest,
|
||||
lora_scale: float = 1.0,
|
||||
) -> Image.Image:
|
||||
params_list = _configure_sampling_params(omni)
|
||||
params_list[1].lora_request = lora_request
|
||||
params_list[1].lora_scale = lora_scale
|
||||
outputs = list(
|
||||
omni.generate(
|
||||
prompts=[{"prompt": DEFAULT_PROMPT, "modalities": ["image"]}],
|
||||
sampling_params_list=params_list,
|
||||
)
|
||||
)
|
||||
img = _extract_generated_image(outputs)
|
||||
assert img is not None, "No image generated with LoRA"
|
||||
return img
|
||||
|
||||
|
||||
# BAGEL uses GQA: hidden_size=3584, 28 Q heads, 4 KV heads, head_dim=128
|
||||
# QKV packed dim = 28*128 + 4*128 + 4*128 = 3584 + 512 + 512 = 4608
|
||||
_LORA_DIM = 3584
|
||||
_LORA_QKV_DIM = 4608
|
||||
_LORA_MODULE = "bagel.language_model.model.layers.0.self_attn.qkv_proj"
|
||||
_LORA_RANK = 4
|
||||
|
||||
|
||||
def _make_file_lora_request(adapter_dir: Path) -> LoRARequest:
|
||||
"""Write synthetic adapter to disk and return a file-backed LoRARequest."""
|
||||
adapter_dir.mkdir(parents=True, exist_ok=True)
|
||||
gen = torch.Generator().manual_seed(42)
|
||||
lora_a = torch.randn((_LORA_RANK, _LORA_DIM), dtype=torch.float32, generator=gen) * 0.1
|
||||
lora_b = torch.randn((_LORA_QKV_DIM, _LORA_RANK), dtype=torch.float32, generator=gen) * 0.5
|
||||
save_file(
|
||||
{
|
||||
f"base_model.model.{_LORA_MODULE}.lora_A.weight": lora_a,
|
||||
f"base_model.model.{_LORA_MODULE}.lora_B.weight": lora_b,
|
||||
},
|
||||
str(adapter_dir / "adapter_model.safetensors"),
|
||||
)
|
||||
(adapter_dir / "adapter_config.json").write_text(
|
||||
json.dumps({"r": _LORA_RANK, "lora_alpha": _LORA_RANK, "target_modules": [_LORA_MODULE]}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
lora_dir = str(adapter_dir)
|
||||
return LoRARequest(lora_name="test_file", lora_int_id=stable_lora_int_id(lora_dir), lora_path=lora_dir)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
|
||||
def test_bagel_lora_scale_and_deactivation(run_level, tmp_path):
|
||||
"""Validate LoRA effect, bounded perturbation, and clean deactivation."""
|
||||
config_path = _resolve_stage_config(BAGEL_STAGE_CONFIG, run_level)
|
||||
omni = Omni(model=MODEL, stage_configs_path=config_path, stage_init_timeout=300)
|
||||
try:
|
||||
lora_request = _make_file_lora_request(tmp_path / "bagel_lora")
|
||||
|
||||
# 1) Baseline (no LoRA)
|
||||
baseline = _generate_bagel_image(omni)
|
||||
|
||||
# 2) LoRA with scale=1.0
|
||||
img_1x = _generate_bagel_image_with_lora(omni, lora_request, lora_scale=1.0)
|
||||
|
||||
# 3) LoRA with scale=2.0
|
||||
img_2x = _generate_bagel_image_with_lora(omni, lora_request, lora_scale=2.0)
|
||||
|
||||
# 4) No LoRA again (deactivation)
|
||||
restored = _generate_bagel_image(omni)
|
||||
|
||||
baseline_arr = np.array(baseline, dtype=np.int16)
|
||||
img_1x_arr = np.array(img_1x, dtype=np.int16)
|
||||
img_2x_arr = np.array(img_2x, dtype=np.int16)
|
||||
restored_arr = np.array(restored, dtype=np.int16)
|
||||
|
||||
diff_1x = np.abs(baseline_arr - img_1x_arr).mean()
|
||||
diff_2x = np.abs(baseline_arr - img_2x_arr).mean()
|
||||
diff_restored = np.abs(baseline_arr - restored_arr).mean()
|
||||
|
||||
# (a) Adapter has visible effect at both scales
|
||||
assert diff_1x > 0.5, f"LoRA scale=1.0 had no visible effect: diff={diff_1x}"
|
||||
assert diff_2x > 0.5, f"LoRA scale=2.0 had no visible effect: diff={diff_2x}"
|
||||
|
||||
# (b) Different scales produce different outputs
|
||||
assert not np.isclose(diff_1x, diff_2x, atol=1.0), (
|
||||
f"LoRA scale has no effect: diff_1x={diff_1x:.2f}, diff_2x={diff_2x:.2f}"
|
||||
)
|
||||
|
||||
# (c) Output is not corrupted
|
||||
assert diff_1x < 80, f"LoRA output looks corrupted: diff_1x={diff_1x}"
|
||||
assert diff_2x < 80, f"LoRA output looks corrupted: diff_2x={diff_2x}"
|
||||
|
||||
# (d) Deactivation fully restores base model
|
||||
assert diff_restored == 0.0, f"Base model not restored after LoRA deactivation: diff={diff_restored}"
|
||||
finally:
|
||||
omni.close()
|
||||
@@ -0,0 +1,419 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
E2E offline smoke tests for Dynin-Omni.
|
||||
|
||||
- model: "snu-aidas/Dynin-Omni"
|
||||
- stage config: tests/e2e/stage_configs/dynin_omni_ci.yaml
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from tests.conftest import OmniRunner
|
||||
from tests.utils import hardware_test
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
_DEFAULT_DYNIN_CONFIG_PATH: Path | None = None
|
||||
_DEFAULT_STAGE_CONFIG_PATH = _REPO_ROOT / "tests" / "e2e" / "stage_configs" / "dynin_omni_ci.yaml"
|
||||
|
||||
models = ["snu-aidas/Dynin-Omni"]
|
||||
stage_configs = [str(_DEFAULT_STAGE_CONFIG_PATH)]
|
||||
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
|
||||
|
||||
DYNIN_CONFIG_PATH = str(_DEFAULT_DYNIN_CONFIG_PATH) if _DEFAULT_DYNIN_CONFIG_PATH is not None else None
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.omni,
|
||||
]
|
||||
|
||||
|
||||
# prompting util
|
||||
def _build_mmu_prompt(tokenizer: Any, question: str, dynin_config_path: str | None) -> dict[str, Any]:
|
||||
encoded = tokenizer(question, return_tensors="pt", add_special_tokens=True)
|
||||
token_ids = [int(v) for v in encoded["input_ids"][0].tolist()]
|
||||
attention_mask = [int(v) for v in encoded["attention_mask"][0].tolist()]
|
||||
additional_information: dict[str, Any] = {
|
||||
"task": ["mmu"],
|
||||
"detok_id": [0],
|
||||
"prompt_length": [len(token_ids)],
|
||||
"attention_mask": [attention_mask],
|
||||
"max_new_tokens": [64],
|
||||
"steps": [64],
|
||||
"block_length": [16],
|
||||
"temperature": [0.0],
|
||||
}
|
||||
if dynin_config_path:
|
||||
additional_information["dynin_config_path"] = [str(dynin_config_path)]
|
||||
return {
|
||||
"prompt_token_ids": token_ids,
|
||||
"additional_information": additional_information,
|
||||
"modalities": ["text"],
|
||||
}
|
||||
|
||||
|
||||
def _build_mmu_multimodal_prompt(
|
||||
tokenizer: Any,
|
||||
question: str,
|
||||
dynin_config_path: str | None,
|
||||
*,
|
||||
image: Any | None = None,
|
||||
audio: tuple[np.ndarray, int] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if image is None and audio is None:
|
||||
raise ValueError("At least one multimodal input (image or audio) must be provided.")
|
||||
|
||||
prefix_chunks: list[str] = []
|
||||
mm_data: dict[str, Any] = {}
|
||||
if image is not None:
|
||||
prefix_chunks.append("<|soi|><|image|><|eoi|>")
|
||||
mm_data["image"] = image
|
||||
if audio is not None:
|
||||
prefix_chunks.append("<|soa|><|audio|><|eoa|>")
|
||||
mm_data["audio"] = audio
|
||||
|
||||
prefixed_question = " ".join(prefix_chunks + [question]).strip()
|
||||
prompt = _build_mmu_prompt(
|
||||
tokenizer=tokenizer,
|
||||
question=prefixed_question,
|
||||
dynin_config_path=dynin_config_path,
|
||||
)
|
||||
prompt["multi_modal_data"] = mm_data
|
||||
prompt["modalities"] = ["text"]
|
||||
return prompt
|
||||
|
||||
|
||||
def _generate_synthetic_image(width: int = 224, height: int = 224) -> np.ndarray:
|
||||
x = np.linspace(0, 255, width, dtype=np.uint8)
|
||||
y = np.linspace(0, 255, height, dtype=np.uint8)[:, None]
|
||||
red = np.tile(x, (height, 1))
|
||||
green = np.tile(y, (1, width))
|
||||
blue = ((red.astype(np.uint16) + green.astype(np.uint16)) // 2).astype(np.uint8)
|
||||
return np.stack([red, green, blue], axis=-1)
|
||||
|
||||
|
||||
def _generate_synthetic_audio(duration_s: int = 5, sample_rate: int = 48_000) -> tuple[np.ndarray, int]:
|
||||
t = np.linspace(0, duration_s, int(sample_rate * duration_s), endpoint=False, dtype=np.float32)
|
||||
waveform = 0.1 * np.sin(2.0 * np.pi * 440.0 * t)
|
||||
return waveform.astype(np.float32), sample_rate
|
||||
|
||||
|
||||
# prompting util
|
||||
def _build_t2s_decode_prompt(dynin_config_path: str | None) -> dict[str, Any]:
|
||||
# Bypass stage-0 generation and directly validate token->audio decode path.
|
||||
generated_audio_token_ids = [int(v) for v in ([10, 11, 12, 13, 14] * 32)]
|
||||
additional_information: dict[str, Any] = {
|
||||
"task": ["t2s"],
|
||||
"detok_id": [1],
|
||||
"generated_token_ids": [generated_audio_token_ids],
|
||||
"audio_codebook_size": [4096],
|
||||
}
|
||||
if dynin_config_path:
|
||||
additional_information["dynin_config_path"] = [str(dynin_config_path)]
|
||||
return {
|
||||
"prompt_token_ids": [0],
|
||||
"additional_information": additional_information,
|
||||
"modalities": ["audio"],
|
||||
}
|
||||
|
||||
|
||||
# prompting util
|
||||
def _build_t2i_decode_prompt(dynin_config_path: str | None) -> dict[str, Any]:
|
||||
# Bypass stage-0 generation and directly validate token->image decode path.
|
||||
# MAGVIT decode path expects a square token grid; 1024 tokens -> 32x32.
|
||||
generated_image_token_ids = [int(v) for v in ([10, 11, 12, 13, 14, 15, 16, 17] * 128)]
|
||||
additional_information: dict[str, Any] = {
|
||||
"task": ["t2i"],
|
||||
"detok_id": [2],
|
||||
"generated_token_ids": [generated_image_token_ids],
|
||||
"codebook_size": [8192],
|
||||
}
|
||||
if dynin_config_path:
|
||||
additional_information["dynin_config_path"] = [str(dynin_config_path)]
|
||||
return {
|
||||
"prompt_token_ids": [0],
|
||||
"additional_information": additional_information,
|
||||
"modalities": ["image"],
|
||||
}
|
||||
|
||||
|
||||
def _configure_dynin_config_env() -> None:
|
||||
if DYNIN_CONFIG_PATH:
|
||||
os.environ["DYNIN_CONFIG_PATH"] = str(DYNIN_CONFIG_PATH)
|
||||
else:
|
||||
os.environ.pop("DYNIN_CONFIG_PATH", None)
|
||||
|
||||
|
||||
def _is_finished_request_output(request_output: Any) -> bool:
|
||||
if request_output is None:
|
||||
return False
|
||||
req_list = request_output if isinstance(request_output, list) else [request_output]
|
||||
for req in req_list:
|
||||
if req is not None and bool(getattr(req, "finished", False)):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _find_stage_output(outputs: list[Any], output_type: str) -> Any | None:
|
||||
matched = [
|
||||
stage_output for stage_output in outputs if getattr(stage_output, "final_output_type", None) == output_type
|
||||
]
|
||||
if not matched:
|
||||
return None
|
||||
|
||||
# Prefer the latest finished chunk to avoid picking an intermediate stream output.
|
||||
for stage_output in reversed(matched):
|
||||
if _is_finished_request_output(getattr(stage_output, "request_output", None)):
|
||||
return stage_output
|
||||
return matched[-1]
|
||||
|
||||
|
||||
def _to_token_list(value: Any) -> list[int]:
|
||||
if value is None:
|
||||
return []
|
||||
if hasattr(value, "detach"):
|
||||
value = value.detach()
|
||||
if hasattr(value, "cpu"):
|
||||
value = value.cpu()
|
||||
if hasattr(value, "flatten"):
|
||||
value = value.flatten().tolist()
|
||||
if isinstance(value, tuple):
|
||||
value = list(value)
|
||||
if not isinstance(value, list):
|
||||
return []
|
||||
out: list[int] = []
|
||||
for token in value:
|
||||
if isinstance(token, bool):
|
||||
continue
|
||||
try:
|
||||
out.append(int(token))
|
||||
except Exception:
|
||||
continue
|
||||
return out
|
||||
|
||||
|
||||
def _extract_text(stage_output: Any, tokenizer: Any | None = None) -> str:
|
||||
request_output = getattr(stage_output, "request_output", None)
|
||||
if request_output is None:
|
||||
return ""
|
||||
req_list = request_output if isinstance(request_output, list) else [request_output]
|
||||
for req in req_list:
|
||||
completions = getattr(req, "outputs", None) or []
|
||||
if not completions:
|
||||
continue
|
||||
completion = completions[0]
|
||||
mm_out = (
|
||||
getattr(completion, "multimodal_output", None)
|
||||
or getattr(req, "multimodal_output", None)
|
||||
or getattr(stage_output, "multimodal_output", None)
|
||||
or {}
|
||||
)
|
||||
text = mm_out.get("text")
|
||||
if isinstance(text, list) and text:
|
||||
text = text[-1]
|
||||
if isinstance(text, str) and text.strip():
|
||||
return text.strip()
|
||||
if tokenizer is not None:
|
||||
for key in ("text_tokens", "token_ids"):
|
||||
token_ids = _to_token_list(mm_out.get(key))
|
||||
if not token_ids:
|
||||
continue
|
||||
decoded = tokenizer.decode(token_ids, skip_special_tokens=True)
|
||||
if isinstance(decoded, str) and decoded.strip():
|
||||
return decoded.strip()
|
||||
fallback = getattr(completion, "text", None)
|
||||
if isinstance(fallback, str) and fallback.strip():
|
||||
return fallback.strip()
|
||||
return ""
|
||||
|
||||
|
||||
def _extract_audio(stage_output: Any) -> Any | None:
|
||||
request_output = getattr(stage_output, "request_output", None)
|
||||
if request_output is None:
|
||||
return None
|
||||
req_list = request_output if isinstance(request_output, list) else [request_output]
|
||||
for req in req_list:
|
||||
completions = getattr(req, "outputs", None) or []
|
||||
if not completions:
|
||||
continue
|
||||
completion = completions[0]
|
||||
mm_out = getattr(completion, "multimodal_output", None) or {}
|
||||
if "audio" in mm_out:
|
||||
return mm_out["audio"]
|
||||
return None
|
||||
|
||||
|
||||
def _extract_image(stage_output: Any) -> Any | None:
|
||||
request_output = getattr(stage_output, "request_output", None)
|
||||
if request_output is None:
|
||||
return None
|
||||
req_list = request_output if isinstance(request_output, list) else [request_output]
|
||||
for req in req_list:
|
||||
completions = getattr(req, "outputs", None) or []
|
||||
if not completions:
|
||||
continue
|
||||
completion = completions[0]
|
||||
mm_out = getattr(completion, "multimodal_output", None) or {}
|
||||
if "image" in mm_out:
|
||||
return mm_out["image"]
|
||||
return None
|
||||
|
||||
|
||||
def _numel(value: Any) -> int:
|
||||
if value is None:
|
||||
return 0
|
||||
if isinstance(value, torch.Tensor):
|
||||
return int(value.numel())
|
||||
shape = getattr(value, "shape", None)
|
||||
if shape is not None:
|
||||
try:
|
||||
total = 1
|
||||
for dim in shape:
|
||||
total *= int(dim)
|
||||
return int(total)
|
||||
except Exception:
|
||||
pass
|
||||
if isinstance(value, (list, tuple)):
|
||||
return len(value)
|
||||
return 0
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
|
||||
@pytest.mark.parametrize("test_config", test_params)
|
||||
def test_dynin_t2i_decode_to_image(test_config: tuple[str, str]) -> None:
|
||||
model, stage_config_path = test_config
|
||||
_configure_dynin_config_env()
|
||||
prompt = _build_t2i_decode_prompt(dynin_config_path=DYNIN_CONFIG_PATH)
|
||||
|
||||
with OmniRunner(
|
||||
model,
|
||||
seed=42,
|
||||
stage_configs_path=stage_config_path,
|
||||
stage_init_timeout=600,
|
||||
init_timeout=600,
|
||||
) as runner:
|
||||
outputs = runner.generate([prompt])
|
||||
|
||||
image_output = _find_stage_output(outputs, "image")
|
||||
assert image_output is not None
|
||||
image_value = _extract_image(image_output)
|
||||
assert image_value is not None
|
||||
assert _numel(image_value) > 0
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
|
||||
@pytest.mark.parametrize("test_config", test_params)
|
||||
def test_dynin_mmu_to_text(test_config: tuple[str, str]) -> None:
|
||||
model, stage_config_path = test_config
|
||||
_configure_dynin_config_env()
|
||||
tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
|
||||
prompt = _build_mmu_prompt(
|
||||
tokenizer=tokenizer,
|
||||
question="What is 2 + 2? Answer in one short sentence.",
|
||||
dynin_config_path=DYNIN_CONFIG_PATH,
|
||||
)
|
||||
|
||||
with OmniRunner(
|
||||
model,
|
||||
seed=42,
|
||||
stage_configs_path=stage_config_path,
|
||||
stage_init_timeout=600,
|
||||
init_timeout=600,
|
||||
) as runner:
|
||||
outputs = runner.generate([prompt])
|
||||
|
||||
text_output = _find_stage_output(outputs, "text")
|
||||
assert text_output is not None
|
||||
text_content = _extract_text(text_output, tokenizer=tokenizer)
|
||||
assert text_content
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
|
||||
@pytest.mark.parametrize("test_config", test_params)
|
||||
def test_dynin_image_to_text(test_config: tuple[str, str]) -> None:
|
||||
model, stage_config_path = test_config
|
||||
_configure_dynin_config_env()
|
||||
tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
|
||||
prompt = _build_mmu_multimodal_prompt(
|
||||
tokenizer=tokenizer,
|
||||
question="Describe the image briefly in one sentence.",
|
||||
dynin_config_path=DYNIN_CONFIG_PATH,
|
||||
image=_generate_synthetic_image(),
|
||||
)
|
||||
|
||||
with OmniRunner(
|
||||
model,
|
||||
seed=42,
|
||||
stage_configs_path=stage_config_path,
|
||||
stage_init_timeout=600,
|
||||
init_timeout=600,
|
||||
) as runner:
|
||||
outputs = runner.generate([prompt])
|
||||
|
||||
text_output = _find_stage_output(outputs, "text")
|
||||
assert text_output is not None
|
||||
text_content = _extract_text(text_output, tokenizer=tokenizer)
|
||||
assert text_content
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
|
||||
@pytest.mark.parametrize("test_config", test_params)
|
||||
def test_dynin_speech_to_text(test_config: tuple[str, str]) -> None:
|
||||
model, stage_config_path = test_config
|
||||
_configure_dynin_config_env()
|
||||
tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
|
||||
prompt = _build_mmu_multimodal_prompt(
|
||||
tokenizer=tokenizer,
|
||||
question="Transcribe the audio briefly in one sentence.",
|
||||
dynin_config_path=DYNIN_CONFIG_PATH,
|
||||
audio=_generate_synthetic_audio(),
|
||||
)
|
||||
|
||||
with OmniRunner(
|
||||
model,
|
||||
seed=42,
|
||||
stage_configs_path=stage_config_path,
|
||||
stage_init_timeout=600,
|
||||
init_timeout=600,
|
||||
) as runner:
|
||||
outputs = runner.generate([prompt])
|
||||
|
||||
text_output = _find_stage_output(outputs, "text")
|
||||
assert text_output is not None
|
||||
text_content = _extract_text(text_output, tokenizer=tokenizer)
|
||||
assert text_content
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
|
||||
@pytest.mark.parametrize("test_config", test_params)
|
||||
def test_dynin_t2s_decode_to_audio(test_config: tuple[str, str]) -> None:
|
||||
model, stage_config_path = test_config
|
||||
_configure_dynin_config_env()
|
||||
prompt = _build_t2s_decode_prompt(dynin_config_path=DYNIN_CONFIG_PATH)
|
||||
|
||||
with OmniRunner(
|
||||
model,
|
||||
seed=42,
|
||||
stage_configs_path=stage_config_path,
|
||||
stage_init_timeout=600,
|
||||
init_timeout=600,
|
||||
) as runner:
|
||||
outputs = runner.generate([prompt])
|
||||
|
||||
audio_output = _find_stage_output(outputs, "audio")
|
||||
assert audio_output is not None
|
||||
audio_value = _extract_audio(audio_output)
|
||||
assert audio_value is not None
|
||||
assert _numel(audio_value) > 0
|
||||
@@ -0,0 +1,243 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from tests.utils import hardware_test
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
T2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "text_to_video" / "text_to_video.py"
|
||||
I2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "image_to_video" / "image_to_video.py"
|
||||
|
||||
T2V_PROMPT = (
|
||||
"At sunrise, a glowing paper lantern boat drifts through a narrow canal between mossy stone walls, "
|
||||
"soft fog above the water, the camera slowly gliding forward as golden reflections shimmer across "
|
||||
"the ripples, cinematic, realistic, highly detailed."
|
||||
)
|
||||
T2V_NEGATIVE_PROMPT = "worst quality, blurry, jittery motion, distorted, oversaturated, artifacts"
|
||||
I2V_PROMPT = "A cinematic dolly shot of a boat drifting on calm water at sunset"
|
||||
I2V_NEGATIVE_PROMPT = "worst quality, blurry, jittery motion"
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
|
||||
def _get_ltx2_model() -> str:
|
||||
return os.environ.get("VLLM_TEST_LTX2_MODEL", "Lightricks/LTX-2")
|
||||
|
||||
|
||||
def _md5(path: Path) -> str:
|
||||
digest = hashlib.md5(usedforsecurity=False)
|
||||
with path.open("rb") as f:
|
||||
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _make_deterministic_test_image(path: Path) -> None:
|
||||
"""Create a deterministic 256x256 test image for I2V tests."""
|
||||
rng = np.random.RandomState(42)
|
||||
img = Image.fromarray(rng.randint(0, 255, (256, 256, 3), dtype=np.uint8))
|
||||
img.save(path)
|
||||
|
||||
|
||||
def _run_and_check(cmd: list[str], env: dict, output_path: Path, expected_md5: str) -> None:
|
||||
result = subprocess.run(cmd, cwd=REPO_ROOT, env=env, capture_output=True, text=True, check=False)
|
||||
assert result.returncode == 0, (
|
||||
f"Command failed (exit {result.returncode}).\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
|
||||
)
|
||||
generated_md5 = _md5(output_path)
|
||||
assert generated_md5 == expected_md5, (
|
||||
f"Unexpected output md5: {generated_md5} != {expected_md5}.\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
|
||||
)
|
||||
|
||||
|
||||
# ── T2V tests ──
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parallel
|
||||
@pytest.mark.slow
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=2)
|
||||
def test_ltx2_t2v_cfg_parallel(tmp_path: Path):
|
||||
"""T2V with CFG=4.0, cfg-parallel-size=2."""
|
||||
output = tmp_path / "t2v_cfg4.mp4"
|
||||
env = os.environ.copy()
|
||||
env.setdefault("CUDA_VISIBLE_DEVICES", "0,1")
|
||||
cmd = [
|
||||
sys.executable,
|
||||
str(T2V_EXAMPLE),
|
||||
"--model",
|
||||
_get_ltx2_model(),
|
||||
"--prompt",
|
||||
T2V_PROMPT,
|
||||
"--negative-prompt",
|
||||
T2V_NEGATIVE_PROMPT,
|
||||
"--height",
|
||||
"256",
|
||||
"--width",
|
||||
"256",
|
||||
"--num-frames",
|
||||
"145",
|
||||
"--num-inference-steps",
|
||||
"6",
|
||||
"--guidance-scale",
|
||||
"4.0",
|
||||
"--frame-rate",
|
||||
"24",
|
||||
"--fps",
|
||||
"24",
|
||||
"--seed",
|
||||
"42",
|
||||
"--cfg-parallel-size",
|
||||
"2",
|
||||
"--enforce-eager",
|
||||
"--output",
|
||||
str(output),
|
||||
]
|
||||
_run_and_check(cmd, env, output, expected_md5="08e606b9c522fee4b6f30cee8b77db40")
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.slow
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_ltx2_t2v_no_cfg(tmp_path: Path):
|
||||
"""T2V with CFG=1.0 (no classifier-free guidance)."""
|
||||
output = tmp_path / "t2v_nocfg.mp4"
|
||||
env = os.environ.copy()
|
||||
env.setdefault("CUDA_VISIBLE_DEVICES", "0")
|
||||
cmd = [
|
||||
sys.executable,
|
||||
str(T2V_EXAMPLE),
|
||||
"--model",
|
||||
_get_ltx2_model(),
|
||||
"--prompt",
|
||||
T2V_PROMPT,
|
||||
"--height",
|
||||
"256",
|
||||
"--width",
|
||||
"256",
|
||||
"--num-frames",
|
||||
"145",
|
||||
"--num-inference-steps",
|
||||
"6",
|
||||
"--guidance-scale",
|
||||
"1.0",
|
||||
"--frame-rate",
|
||||
"24",
|
||||
"--fps",
|
||||
"24",
|
||||
"--seed",
|
||||
"42",
|
||||
"--enforce-eager",
|
||||
"--output",
|
||||
str(output),
|
||||
]
|
||||
_run_and_check(cmd, env, output, expected_md5="a83994b94b6e67c54a524e0383c45ce8")
|
||||
|
||||
|
||||
# ── I2V tests ──
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parallel
|
||||
@pytest.mark.slow
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=2)
|
||||
def test_ltx2_i2v_cfg_parallel(tmp_path: Path):
|
||||
"""I2V with CFG=4.0, cfg-parallel-size=2."""
|
||||
test_image = tmp_path / "test_input.png"
|
||||
_make_deterministic_test_image(test_image)
|
||||
output = tmp_path / "i2v_cfg4.mp4"
|
||||
env = os.environ.copy()
|
||||
env.setdefault("CUDA_VISIBLE_DEVICES", "0,1")
|
||||
cmd = [
|
||||
sys.executable,
|
||||
str(I2V_EXAMPLE),
|
||||
"--model",
|
||||
_get_ltx2_model(),
|
||||
"--model-class-name",
|
||||
"LTX2ImageToVideoPipeline",
|
||||
"--image",
|
||||
str(test_image),
|
||||
"--prompt",
|
||||
I2V_PROMPT,
|
||||
"--negative-prompt",
|
||||
I2V_NEGATIVE_PROMPT,
|
||||
"--height",
|
||||
"256",
|
||||
"--width",
|
||||
"256",
|
||||
"--num-frames",
|
||||
"73",
|
||||
"--num-inference-steps",
|
||||
"6",
|
||||
"--guidance-scale",
|
||||
"4.0",
|
||||
"--frame-rate",
|
||||
"24",
|
||||
"--fps",
|
||||
"24",
|
||||
"--seed",
|
||||
"42",
|
||||
"--cfg-parallel-size",
|
||||
"2",
|
||||
"--enforce-eager",
|
||||
"--output",
|
||||
str(output),
|
||||
]
|
||||
_run_and_check(cmd, env, output, expected_md5="aed7e56084b36373244d8f839b16d115")
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.slow
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_ltx2_i2v_no_cfg(tmp_path: Path):
|
||||
"""I2V with CFG=1.0 (no classifier-free guidance)."""
|
||||
test_image = tmp_path / "test_input.png"
|
||||
_make_deterministic_test_image(test_image)
|
||||
output = tmp_path / "i2v_nocfg.mp4"
|
||||
env = os.environ.copy()
|
||||
env.setdefault("CUDA_VISIBLE_DEVICES", "0")
|
||||
cmd = [
|
||||
sys.executable,
|
||||
str(I2V_EXAMPLE),
|
||||
"--model",
|
||||
_get_ltx2_model(),
|
||||
"--model-class-name",
|
||||
"LTX2ImageToVideoPipeline",
|
||||
"--image",
|
||||
str(test_image),
|
||||
"--prompt",
|
||||
I2V_PROMPT,
|
||||
"--height",
|
||||
"256",
|
||||
"--width",
|
||||
"256",
|
||||
"--num-frames",
|
||||
"73",
|
||||
"--num-inference-steps",
|
||||
"6",
|
||||
"--guidance-scale",
|
||||
"1.0",
|
||||
"--frame-rate",
|
||||
"24",
|
||||
"--fps",
|
||||
"24",
|
||||
"--seed",
|
||||
"42",
|
||||
"--enforce-eager",
|
||||
"--output",
|
||||
str(output),
|
||||
]
|
||||
_run_and_check(cmd, env, output, expected_md5="81b21ede12753e9e14a357a6c548b666")
|
||||
@@ -110,16 +110,31 @@ def test_magi_human_e2e(run_level):
|
||||
assert isinstance(video_frames, np.ndarray), f"Expected numpy array, got {type(video_frames)}"
|
||||
assert video_frames.ndim == 4, f"Expected 4D array (T,H,W,3), got shape {video_frames.shape}"
|
||||
|
||||
audio_waveform = None
|
||||
if hasattr(first, "multimodal_output") and first.multimodal_output:
|
||||
audio_waveform = first.multimodal_output.get("audio")
|
||||
mm = first.multimodal_output
|
||||
assert mm, "multimodal_output is empty or missing"
|
||||
|
||||
audio_waveform = mm.get("audio")
|
||||
assert audio_waveform is not None, "No audio waveform in multimodal_output"
|
||||
|
||||
audio_sample_rate = mm.get("audio_sample_rate")
|
||||
assert audio_sample_rate is not None, (
|
||||
"audio_sample_rate not found in multimodal_output; model post-process must propagate it"
|
||||
)
|
||||
assert isinstance(audio_sample_rate, (int, float)), (
|
||||
f"audio_sample_rate should be numeric, got {type(audio_sample_rate)}"
|
||||
)
|
||||
assert int(audio_sample_rate) > 0, f"audio_sample_rate must be positive, got {audio_sample_rate}"
|
||||
|
||||
fps = mm.get("fps")
|
||||
assert fps is not None, "fps not found in multimodal_output; model post-process must propagate it"
|
||||
assert isinstance(fps, (int, float)), f"fps should be numeric, got {type(fps)}"
|
||||
assert int(fps) > 0, f"fps must be positive, got {fps}"
|
||||
|
||||
video_bytes = mux_video_audio_bytes(
|
||||
video_frames,
|
||||
audio_waveform,
|
||||
fps=25.0,
|
||||
audio_sample_rate=44100,
|
||||
fps=float(fps),
|
||||
audio_sample_rate=int(audio_sample_rate),
|
||||
)
|
||||
assert isinstance(video_bytes, bytes), f"Expected MP4 bytes, got {type(video_bytes)}"
|
||||
assert len(video_bytes) > 1000, f"MP4 too small ({len(video_bytes)} bytes)"
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Example online tests for Dynin-Omni model.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import gc
|
||||
import os
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import soundfile as sf
|
||||
from vllm.assets.image import ImageAsset
|
||||
|
||||
from tests import conftest as tests_conftest
|
||||
from tests.conftest import OmniServerParams
|
||||
from tests.utils import hardware_test
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
|
||||
|
||||
MODEL = "snu-aidas/Dynin-Omni"
|
||||
STAGE_CONFIG = str(Path(__file__).parent.parent / "stage_configs" / "dynin_omni_ci.yaml")
|
||||
_WHISPER_SAMPLE_RATE_HZ = 16_000
|
||||
|
||||
T2I_PROMPT = "A high quality detailed living room interior photo."
|
||||
T2S_PROMPT = "Please read this sentence naturally: Hello from Dynin-Omni online serving."
|
||||
I2I_PROMPT = "Transform this outdoor nature boardwalk scene into a painting style with vivid colors."
|
||||
|
||||
TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG)]
|
||||
_STAGE_COUNT = 3
|
||||
_I2I_STAGE_SAMPLING = {"max_tokens": 1, "temperature": 0.0, "top_p": 1.0, "detokenize": False}
|
||||
|
||||
|
||||
def _prepare_audio_waveform_for_whisper(audio_data: np.ndarray, samplerate: int) -> np.ndarray:
|
||||
"""Normalize decoded audio into a mono 16 kHz float32 waveform for Whisper."""
|
||||
if samplerate <= 0:
|
||||
raise ValueError(f"Invalid audio sample rate: {samplerate}")
|
||||
|
||||
waveform = np.asarray(audio_data, dtype=np.float32)
|
||||
if waveform.ndim == 0:
|
||||
raise ValueError("Audio waveform must have at least one dimension")
|
||||
if waveform.ndim > 1:
|
||||
waveform = np.mean(waveform, axis=1)
|
||||
if waveform.size == 0:
|
||||
raise ValueError("Empty audio waveform")
|
||||
|
||||
if samplerate != _WHISPER_SAMPLE_RATE_HZ:
|
||||
target_num_samples = max(int(round(waveform.shape[0] * _WHISPER_SAMPLE_RATE_HZ / samplerate)), 1)
|
||||
source_positions = np.arange(waveform.shape[0], dtype=np.float64)
|
||||
target_positions = np.linspace(
|
||||
0.0,
|
||||
max(waveform.shape[0] - 1, 0),
|
||||
num=target_num_samples,
|
||||
dtype=np.float64,
|
||||
)
|
||||
waveform = np.interp(target_positions, source_positions, waveform).astype(np.float32)
|
||||
|
||||
return np.ascontiguousarray(np.clip(waveform, -1.0, 1.0), dtype=np.float32)
|
||||
|
||||
|
||||
def _convert_audio_bytes_to_text_without_ffmpeg(raw_bytes: bytes) -> str:
|
||||
"""Dynin t2s keeps Whisper transcription local to this test module and avoids ffmpeg."""
|
||||
import whisper
|
||||
|
||||
data, samplerate = sf.read(BytesIO(raw_bytes), dtype="float32", always_2d=True)
|
||||
audio_waveform = _prepare_audio_waveform_for_whisper(data, samplerate)
|
||||
|
||||
model = whisper.load_model("small", device="cpu")
|
||||
try:
|
||||
transcript = model.transcribe(
|
||||
audio_waveform,
|
||||
temperature=0.0,
|
||||
word_timestamps=True,
|
||||
condition_on_previous_text=False,
|
||||
)["text"]
|
||||
finally:
|
||||
del model
|
||||
gc.collect()
|
||||
|
||||
return transcript or ""
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dynin_t2s_openai_client(openai_client, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
tests_conftest,
|
||||
"convert_audio_bytes_to_text",
|
||||
_convert_audio_bytes_to_text_without_ffmpeg,
|
||||
)
|
||||
return openai_client
|
||||
|
||||
|
||||
def _build_t2i_messages(prompt: str) -> list[dict]:
|
||||
return [{"role": "user", "content": [{"type": "text", "text": f"<|t2i|> {prompt}"}]}]
|
||||
|
||||
|
||||
def _build_t2s_messages(prompt: str) -> list[dict]:
|
||||
return [{"role": "user", "content": [{"type": "text", "text": f"<|t2s|> {prompt}"}]}]
|
||||
|
||||
|
||||
def _build_i2i_messages(prompt: str) -> list[dict]:
|
||||
input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
|
||||
buffer = BytesIO()
|
||||
input_image.save(buffer, format="JPEG")
|
||||
image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": f"<|i2i|> {prompt}"},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
def test_send_i2i_request_001(omni_server, openai_client) -> None:
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": _build_i2i_messages(I2I_PROMPT),
|
||||
"modalities": ["image"],
|
||||
"extra_body": {
|
||||
"sampling_params_list": [dict(_I2I_STAGE_SAMPLING) for _ in range(_STAGE_COUNT)],
|
||||
},
|
||||
}
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
def test_send_t2i_request_001(omni_server, openai_client) -> None:
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": _build_t2i_messages(T2I_PROMPT),
|
||||
"modalities": ["image"],
|
||||
}
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
def test_send_t2s_request_001(omni_server, dynin_t2s_openai_client) -> None:
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": _build_t2s_messages(T2S_PROMPT),
|
||||
"modalities": ["audio"],
|
||||
}
|
||||
dynin_t2s_openai_client.send_omni_request(request_config)
|
||||
@@ -0,0 +1,84 @@
|
||||
# stage config for running dynin_omni with a 3-stage architecture.
|
||||
# this config is intended for e2e smoke tests.
|
||||
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
runtime:
|
||||
process: true
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: token2text
|
||||
model_arch: DyninOmniForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
engine_output_type: latent
|
||||
trust_remote_code: true
|
||||
gpu_memory_utilization: 0.5
|
||||
enforce_eager: true
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: false
|
||||
max_num_batched_tokens: 4096
|
||||
is_comprehension: true
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
process: true
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: token2image
|
||||
model_arch: DyninOmniForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
engine_output_type: latent
|
||||
trust_remote_code: true
|
||||
gpu_memory_utilization: 0.2
|
||||
enforce_eager: true
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: false
|
||||
max_num_batched_tokens: 4096
|
||||
engine_input_source: [0]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image
|
||||
final_output: true
|
||||
final_output_type: image
|
||||
|
||||
- stage_id: 2
|
||||
stage_type: llm
|
||||
runtime:
|
||||
process: true
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: token2audio
|
||||
model_arch: DyninOmniForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
engine_output_type: latent
|
||||
trust_remote_code: true
|
||||
gpu_memory_utilization: 0.2
|
||||
enforce_eager: true
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: false
|
||||
max_num_batched_tokens: 4096
|
||||
engine_input_source: [1]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 1
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
- from: 1
|
||||
to: 2
|
||||
window_size: -1
|
||||
@@ -0,0 +1,44 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for cross-stage LoRA routing in the orchestrator."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from vllm_omni.engine.orchestrator import build_engine_core_request_from_tokens
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class TestBuildEngineCoreRequestLoRA:
|
||||
"""Verify build_engine_core_request_from_tokens passes LoRA from params."""
|
||||
|
||||
def test_lora_extracted_from_diffusion_params(self):
|
||||
lr = LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/fake")
|
||||
params = OmniDiffusionSamplingParams(lora_request=lr)
|
||||
|
||||
# OmniDiffusionSamplingParams is not a SamplingParams, so
|
||||
# build_engine_core_request_from_tokens takes the pooling path.
|
||||
# We only care that lora_request is extracted via getattr.
|
||||
request = build_engine_core_request_from_tokens(
|
||||
request_id="req-1",
|
||||
prompt={"prompt_token_ids": [1, 2, 3]},
|
||||
params=params,
|
||||
model_config=None,
|
||||
)
|
||||
assert request.lora_request is lr
|
||||
|
||||
def test_no_lora_on_sampling_params(self):
|
||||
params = SamplingParams(max_tokens=10)
|
||||
|
||||
request = build_engine_core_request_from_tokens(
|
||||
request_id="req-2",
|
||||
prompt={"prompt_token_ids": [1, 2, 3]},
|
||||
params=params,
|
||||
model_config=None,
|
||||
)
|
||||
assert request.lora_request is None
|
||||
@@ -0,0 +1,510 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import janus
|
||||
import pytest
|
||||
from vllm.outputs import CompletionOutput, RequestOutput
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from vllm_omni.engine.orchestrator import Orchestrator
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@dataclass
|
||||
class OrchestratorFixture:
|
||||
orchestrator: Orchestrator
|
||||
request_sync_q: Any
|
||||
output_sync_q: Any
|
||||
queues: tuple[janus.Queue, ...]
|
||||
thread: threading.Thread
|
||||
result_future: concurrent.futures.Future[None]
|
||||
|
||||
|
||||
class FakeStageClient:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
stage_type: str = "llm",
|
||||
final_output: bool = False,
|
||||
final_output_type: str = "text",
|
||||
next_inputs: list[dict] | None = None,
|
||||
) -> None:
|
||||
self.stage_type = stage_type
|
||||
self.final_output = final_output
|
||||
self.final_output_type = final_output_type
|
||||
self.next_inputs = list(next_inputs or [])
|
||||
self.custom_process_input_func = None
|
||||
self.add_request_calls: list[tuple] = []
|
||||
self.abort_calls: list[list[str]] = []
|
||||
self.shutdown_calls = 0
|
||||
self._engine_core_outputs = queue.Queue()
|
||||
self._diffusion_outputs = queue.Queue()
|
||||
|
||||
# Orchestrator-facing interface.
|
||||
async def add_request_async(self, *args, **_kwargs) -> None:
|
||||
self.add_request_calls.append(args)
|
||||
|
||||
async def get_output_async(self):
|
||||
try:
|
||||
return self._engine_core_outputs.get_nowait()
|
||||
except queue.Empty:
|
||||
return SimpleNamespace(outputs=[])
|
||||
|
||||
def get_diffusion_output_nowait(self):
|
||||
try:
|
||||
return self._diffusion_outputs.get_nowait()
|
||||
except queue.Empty:
|
||||
return None
|
||||
|
||||
def set_engine_outputs(self, outputs) -> None:
|
||||
return None
|
||||
|
||||
def process_engine_inputs(self, stage_list, prompt=None):
|
||||
return list(self.next_inputs)
|
||||
|
||||
async def abort_requests_async(self, request_ids: list[str]) -> None:
|
||||
self.abort_calls.append(list(request_ids))
|
||||
|
||||
def shutdown(self) -> None:
|
||||
self.shutdown_calls += 1
|
||||
|
||||
# Test helpers for seeding fake stage outputs.
|
||||
def push_engine_core_outputs(self, outputs) -> None:
|
||||
self._engine_core_outputs.put_nowait(outputs)
|
||||
|
||||
def push_diffusion_output(self, output) -> None:
|
||||
self._diffusion_outputs.put_nowait(output)
|
||||
|
||||
|
||||
class FakeOutputProcessor:
|
||||
def __init__(self, *, request_outputs: list[object] | None = None) -> None:
|
||||
self.request_outputs = list(request_outputs or [])
|
||||
|
||||
def add_request(self, *_args, **_kwargs) -> None:
|
||||
return None
|
||||
|
||||
def process_outputs(self, *_args, **_kwargs):
|
||||
return SimpleNamespace(
|
||||
request_outputs=list(self.request_outputs),
|
||||
reqs_to_abort=[],
|
||||
)
|
||||
|
||||
def update_scheduler_stats(self, _scheduler_stats) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def _sampling_params(max_tokens: int = 4) -> SamplingParams:
|
||||
return SamplingParams(max_tokens=max_tokens)
|
||||
|
||||
|
||||
def _engine_core_outputs(tag: str, timestamp: float) -> SimpleNamespace:
|
||||
return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None)
|
||||
|
||||
|
||||
def _build_request_output(
|
||||
request_id: str,
|
||||
*,
|
||||
token_ids: list[int] | None = None,
|
||||
prompt_token_ids: list[int] | None = None,
|
||||
finished: bool = True,
|
||||
text: str = "test",
|
||||
) -> RequestOutput:
|
||||
completion = CompletionOutput(
|
||||
index=0,
|
||||
text=text,
|
||||
token_ids=list(token_ids or [1, 2]),
|
||||
cumulative_logprob=0.0,
|
||||
logprobs=None,
|
||||
finish_reason="stop" if finished else None,
|
||||
stop_reason=None,
|
||||
)
|
||||
return RequestOutput(
|
||||
request_id=request_id,
|
||||
prompt="prompt",
|
||||
prompt_token_ids=list(prompt_token_ids or [10, 11]),
|
||||
prompt_logprobs=None,
|
||||
outputs=[completion],
|
||||
finished=finished,
|
||||
metrics=None,
|
||||
lora_request=None,
|
||||
)
|
||||
|
||||
|
||||
def _build_harness(
|
||||
stage_clients: list[object],
|
||||
*,
|
||||
output_processors: list[object] | None = None,
|
||||
stage_vllm_configs: list[object] | None = None,
|
||||
async_chunk: bool = False,
|
||||
) -> OrchestratorFixture:
|
||||
if output_processors is None:
|
||||
output_processors = [FakeOutputProcessor() for _ in stage_clients]
|
||||
if stage_vllm_configs is None:
|
||||
stage_vllm_configs = [SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) for _ in stage_clients]
|
||||
|
||||
ready_future: concurrent.futures.Future[tuple[Orchestrator, janus.Queue, janus.Queue, janus.Queue]] = (
|
||||
concurrent.futures.Future()
|
||||
)
|
||||
result_future: concurrent.futures.Future[None] = concurrent.futures.Future()
|
||||
|
||||
def _runner() -> None:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
async def _run() -> None:
|
||||
request_queue = janus.Queue()
|
||||
output_queue = janus.Queue()
|
||||
rpc_queue = janus.Queue()
|
||||
orchestrator = Orchestrator(
|
||||
request_async_queue=request_queue.async_q,
|
||||
output_async_queue=output_queue.async_q,
|
||||
rpc_async_queue=rpc_queue.async_q,
|
||||
stage_clients=stage_clients,
|
||||
output_processors=output_processors,
|
||||
stage_vllm_configs=stage_vllm_configs,
|
||||
async_chunk=async_chunk,
|
||||
)
|
||||
ready_future.set_result((orchestrator, request_queue, output_queue, rpc_queue))
|
||||
await orchestrator.run()
|
||||
|
||||
try:
|
||||
loop.run_until_complete(_run())
|
||||
result_future.set_result(None)
|
||||
except Exception as exc:
|
||||
result_future.set_exception(exc)
|
||||
finally:
|
||||
try:
|
||||
pending = [task for task in asyncio.all_tasks(loop) if not task.done()]
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
if pending:
|
||||
loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
|
||||
loop.run_until_complete(loop.shutdown_asyncgens())
|
||||
finally:
|
||||
asyncio.set_event_loop(None)
|
||||
loop.close()
|
||||
|
||||
thread = threading.Thread(target=_runner, daemon=True, name="test-orchestrator")
|
||||
thread.start()
|
||||
|
||||
orchestrator, request_queue, output_queue, rpc_queue = ready_future.result(timeout=5)
|
||||
return OrchestratorFixture(
|
||||
orchestrator=orchestrator,
|
||||
request_sync_q=request_queue.sync_q,
|
||||
output_sync_q=output_queue.sync_q,
|
||||
queues=(request_queue, output_queue, rpc_queue),
|
||||
thread=thread,
|
||||
result_future=result_future,
|
||||
)
|
||||
|
||||
|
||||
async def _shutdown_orchestrator(orchestrator_fixture: OrchestratorFixture) -> None:
|
||||
orchestrator_fixture.request_sync_q.put_nowait({"type": "shutdown"})
|
||||
await asyncio.to_thread(orchestrator_fixture.thread.join, 5)
|
||||
if orchestrator_fixture.thread.is_alive():
|
||||
raise AssertionError("Timed out waiting for orchestrator thread shutdown")
|
||||
orchestrator_fixture.result_future.result(timeout=0)
|
||||
|
||||
|
||||
async def _wait_for(predicate, *, timeout: float = 2.0) -> None:
|
||||
deadline = time.monotonic() + timeout
|
||||
while not predicate():
|
||||
if time.monotonic() >= deadline:
|
||||
raise AssertionError("Timed out waiting for predicate")
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
|
||||
async def _get_output_message(orchestrator_fixture: OrchestratorFixture, *, timeout: float = 2.0) -> dict:
|
||||
deadline = time.monotonic() + timeout
|
||||
while True:
|
||||
if time.monotonic() >= deadline:
|
||||
raise AssertionError("Timed out waiting for orchestrator output")
|
||||
try:
|
||||
msg = orchestrator_fixture.output_sync_q.get_nowait()
|
||||
except queue.Empty:
|
||||
await asyncio.sleep(0.01)
|
||||
continue
|
||||
if msg.get("type") == "output":
|
||||
return msg
|
||||
|
||||
|
||||
async def _enqueue_add_request(
|
||||
orchestrator_fixture: OrchestratorFixture,
|
||||
*,
|
||||
request_id: str,
|
||||
prompt,
|
||||
original_prompt,
|
||||
sampling_params_list,
|
||||
final_stage_id: int,
|
||||
) -> None:
|
||||
orchestrator_fixture.request_sync_q.put_nowait(
|
||||
{
|
||||
"type": "add_request",
|
||||
"request_id": request_id,
|
||||
"prompt": prompt,
|
||||
"original_prompt": original_prompt,
|
||||
"sampling_params_list": sampling_params_list,
|
||||
"final_stage_id": final_stage_id,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _enqueue_abort_request(orchestrator_fixture: OrchestratorFixture, request_ids: list[str]) -> None:
|
||||
orchestrator_fixture.request_sync_q.put_nowait(
|
||||
{
|
||||
"type": "abort",
|
||||
"request_ids": request_ids,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def orchestrator_factory():
|
||||
fixtures: list[OrchestratorFixture] = []
|
||||
|
||||
def _factory(*args, **kwargs) -> OrchestratorFixture:
|
||||
fixture = _build_harness(*args, **kwargs)
|
||||
fixtures.append(fixture)
|
||||
return fixture
|
||||
|
||||
yield _factory
|
||||
|
||||
for fixture in fixtures:
|
||||
if fixture.thread.is_alive():
|
||||
fixture.request_sync_q.put_nowait({"type": "shutdown"})
|
||||
fixture.thread.join(timeout=5)
|
||||
for q in fixture.queues:
|
||||
q.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_two_stage_llm(orchestrator_factory) -> None:
|
||||
stage0 = FakeStageClient(stage_type="llm", final_output=False)
|
||||
stage1 = FakeStageClient(
|
||||
stage_type="llm",
|
||||
final_output=True,
|
||||
next_inputs=[{"prompt_token_ids": [7, 8, 9]}],
|
||||
)
|
||||
processors = [
|
||||
FakeOutputProcessor(request_outputs=[_build_request_output("req-llm", token_ids=[3, 4], finished=True)]),
|
||||
FakeOutputProcessor(request_outputs=[_build_request_output("req-llm", token_ids=[10, 11], finished=True)]),
|
||||
]
|
||||
orchestrator_fixture = orchestrator_factory([stage0, stage1], output_processors=processors)
|
||||
request = SimpleNamespace(request_id="req-llm", prompt_token_ids=[1, 2, 3])
|
||||
|
||||
try:
|
||||
await _enqueue_add_request(
|
||||
orchestrator_fixture,
|
||||
request_id="req-llm",
|
||||
prompt=request,
|
||||
original_prompt={"prompt": "hello"},
|
||||
sampling_params_list=[_sampling_params(), _sampling_params()],
|
||||
final_stage_id=1,
|
||||
)
|
||||
|
||||
await _wait_for(lambda: len(stage0.add_request_calls) == 1)
|
||||
stage0.push_engine_core_outputs(_engine_core_outputs("stage0-raw", 1.0))
|
||||
|
||||
await _wait_for(lambda: len(stage1.add_request_calls) == 1)
|
||||
stage1_request = stage1.add_request_calls[0][0]
|
||||
assert stage1_request.request_id == "req-llm"
|
||||
assert stage1_request.prompt_token_ids == [7, 8, 9]
|
||||
|
||||
stage1.push_engine_core_outputs(_engine_core_outputs("stage1-raw", 2.0))
|
||||
|
||||
output_msg = await _get_output_message(orchestrator_fixture)
|
||||
|
||||
assert output_msg["request_id"] == "req-llm"
|
||||
assert output_msg["stage_id"] == 1
|
||||
assert output_msg["finished"] is True
|
||||
assert output_msg["engine_outputs"].request_id == "req-llm"
|
||||
assert "req-llm" not in orchestrator_fixture.orchestrator.request_states
|
||||
finally:
|
||||
await _shutdown_orchestrator(orchestrator_fixture)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_single_stage_diffusion(orchestrator_factory) -> None:
|
||||
stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
|
||||
orchestrator_fixture = orchestrator_factory([stage0])
|
||||
params = OmniDiffusionSamplingParams()
|
||||
|
||||
try:
|
||||
await _enqueue_add_request(
|
||||
orchestrator_fixture,
|
||||
request_id="req-diff",
|
||||
prompt={"prompt": "draw a cat"},
|
||||
original_prompt={"prompt": "draw a cat"},
|
||||
sampling_params_list=[params],
|
||||
final_stage_id=0,
|
||||
)
|
||||
|
||||
await _wait_for(lambda: len(stage0.add_request_calls) == 1)
|
||||
stage0.push_diffusion_output(
|
||||
OmniRequestOutput.from_diffusion(
|
||||
request_id="req-diff",
|
||||
images=[],
|
||||
final_output_type="image",
|
||||
)
|
||||
)
|
||||
|
||||
output_msg = await _get_output_message(orchestrator_fixture)
|
||||
|
||||
assert output_msg["request_id"] == "req-diff"
|
||||
assert output_msg["stage_id"] == 0
|
||||
assert output_msg["finished"] is True
|
||||
assert output_msg["engine_outputs"].request_id == "req-diff"
|
||||
assert "req-diff" not in orchestrator_fixture.orchestrator.request_states
|
||||
finally:
|
||||
await _shutdown_orchestrator(orchestrator_fixture)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_llm_to_diffusion(orchestrator_factory) -> None:
|
||||
stage0 = FakeStageClient(stage_type="llm", final_output=False)
|
||||
stage1 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
|
||||
processors = [
|
||||
FakeOutputProcessor(request_outputs=[_build_request_output("req-img", token_ids=[3, 4], finished=True)]),
|
||||
FakeOutputProcessor(),
|
||||
]
|
||||
orchestrator_fixture = orchestrator_factory([stage0, stage1], output_processors=processors)
|
||||
request = SimpleNamespace(request_id="req-img", prompt_token_ids=[1, 2, 3])
|
||||
params = OmniDiffusionSamplingParams()
|
||||
original_prompt = {"prompt": "draw a fox"}
|
||||
|
||||
try:
|
||||
await _enqueue_add_request(
|
||||
orchestrator_fixture,
|
||||
request_id="req-img",
|
||||
prompt=request,
|
||||
original_prompt=original_prompt,
|
||||
sampling_params_list=[_sampling_params(), params],
|
||||
final_stage_id=1,
|
||||
)
|
||||
|
||||
await _wait_for(lambda: len(stage0.add_request_calls) == 1)
|
||||
stage0.push_engine_core_outputs(_engine_core_outputs("stage0-raw", 1.0))
|
||||
|
||||
await _wait_for(lambda: len(stage1.add_request_calls) == 1)
|
||||
assert stage1.add_request_calls[0] == ("req-img", original_prompt, params)
|
||||
|
||||
stage1.push_diffusion_output(
|
||||
OmniRequestOutput.from_diffusion(
|
||||
request_id="req-img",
|
||||
images=[],
|
||||
final_output_type="image",
|
||||
)
|
||||
)
|
||||
|
||||
output_msg = await _get_output_message(orchestrator_fixture)
|
||||
|
||||
assert output_msg["request_id"] == "req-img"
|
||||
assert output_msg["stage_id"] == 1
|
||||
assert output_msg["finished"] is True
|
||||
assert output_msg["engine_outputs"].request_id == "req-img"
|
||||
assert "req-img" not in orchestrator_fixture.orchestrator.request_states
|
||||
finally:
|
||||
await _shutdown_orchestrator(orchestrator_fixture)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_chunk(orchestrator_factory) -> None:
|
||||
stage0 = FakeStageClient(stage_type="llm", final_output=False)
|
||||
stage1 = FakeStageClient(stage_type="llm", final_output=True)
|
||||
processors = [
|
||||
FakeOutputProcessor(request_outputs=[_build_request_output("req-async", token_ids=[1], finished=True)]),
|
||||
FakeOutputProcessor(request_outputs=[_build_request_output("req-async", token_ids=[20, 21], finished=True)]),
|
||||
]
|
||||
orchestrator_fixture = orchestrator_factory(
|
||||
[stage0, stage1],
|
||||
output_processors=processors,
|
||||
async_chunk=True,
|
||||
)
|
||||
request = SimpleNamespace(request_id="req-async", prompt_token_ids=[1, 2, 3, 4])
|
||||
|
||||
try:
|
||||
await _enqueue_add_request(
|
||||
orchestrator_fixture,
|
||||
request_id="req-async",
|
||||
prompt=request,
|
||||
original_prompt={"prompt": "hello async"},
|
||||
sampling_params_list=[_sampling_params(), _sampling_params()],
|
||||
final_stage_id=1,
|
||||
)
|
||||
|
||||
await _wait_for(lambda: len(stage1.add_request_calls) == 1)
|
||||
prewarmed_request = stage1.add_request_calls[0][0]
|
||||
assert prewarmed_request.request_id == "req-async"
|
||||
assert prewarmed_request.prompt_token_ids
|
||||
assert all(token_id == 0 for token_id in prewarmed_request.prompt_token_ids)
|
||||
|
||||
stage1.push_engine_core_outputs(_engine_core_outputs("stage1-final", 3.0))
|
||||
|
||||
output_msg = await _get_output_message(orchestrator_fixture)
|
||||
|
||||
assert output_msg["request_id"] == "req-async"
|
||||
assert output_msg["stage_id"] == 1
|
||||
assert output_msg["finished"] is True
|
||||
assert "req-async" not in orchestrator_fixture.orchestrator.request_states
|
||||
finally:
|
||||
await _shutdown_orchestrator(orchestrator_fixture)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_shutdown(orchestrator_factory) -> None:
|
||||
stages = [
|
||||
FakeStageClient(stage_type="llm", final_output=False),
|
||||
FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image"),
|
||||
]
|
||||
orchestrator_fixture = orchestrator_factory(stages)
|
||||
|
||||
await _shutdown_orchestrator(orchestrator_fixture)
|
||||
|
||||
assert not orchestrator_fixture.thread.is_alive()
|
||||
for stage in stages:
|
||||
assert stage.shutdown_calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_abort(orchestrator_factory) -> None:
|
||||
stages = [
|
||||
FakeStageClient(stage_type="llm", final_output=False),
|
||||
FakeStageClient(stage_type="llm", final_output=True),
|
||||
]
|
||||
processors = [
|
||||
FakeOutputProcessor(request_outputs=[_build_request_output("req-abort", token_ids=[1], finished=True)]),
|
||||
FakeOutputProcessor(request_outputs=[_build_request_output("req-abort", token_ids=[2], finished=True)]),
|
||||
]
|
||||
orchestrator_fixture = orchestrator_factory(stages, output_processors=processors)
|
||||
request = SimpleNamespace(request_id="req-abort", prompt_token_ids=[1, 2, 3])
|
||||
|
||||
try:
|
||||
await _enqueue_add_request(
|
||||
orchestrator_fixture,
|
||||
request_id="req-abort",
|
||||
prompt=request,
|
||||
original_prompt={"prompt": "cancel me"},
|
||||
sampling_params_list=[_sampling_params(), _sampling_params()],
|
||||
final_stage_id=1,
|
||||
)
|
||||
await _wait_for(lambda: len(stages[0].add_request_calls) == 1)
|
||||
|
||||
await _enqueue_abort_request(orchestrator_fixture, ["req-abort"])
|
||||
await _wait_for(lambda: all(stage.abort_calls for stage in stages))
|
||||
|
||||
for stage in stages:
|
||||
assert stage.abort_calls == [["req-abort"]]
|
||||
assert "req-abort" not in orchestrator_fixture.orchestrator.request_states
|
||||
finally:
|
||||
await _shutdown_orchestrator(orchestrator_fixture)
|
||||
@@ -0,0 +1,207 @@
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from vllm import SamplingParams
|
||||
|
||||
from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
|
||||
from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _DummySenderStage:
|
||||
def __init__(self, sender_info):
|
||||
self._sender_info = sender_info
|
||||
self.engine_outputs = None
|
||||
|
||||
def set_engine_outputs(self, outputs):
|
||||
self.engine_outputs = outputs
|
||||
|
||||
def get_kv_sender_info(self):
|
||||
return self._sender_info
|
||||
|
||||
|
||||
class _DummyDiffusionStage:
|
||||
stage_type = "diffusion"
|
||||
custom_process_input_func = None
|
||||
|
||||
def __init__(self, engine_input_source=None):
|
||||
self.engine_input_source = engine_input_source or [0]
|
||||
self.calls = []
|
||||
|
||||
async def add_request_async(self, request_id, prompt, sampling_params, kv_sender_info=None):
|
||||
self.calls.append(
|
||||
{
|
||||
"request_id": request_id,
|
||||
"prompt": prompt,
|
||||
"sampling_params": sampling_params,
|
||||
"kv_sender_info": kv_sender_info,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_stage_engine_core_client_builds_kv_sender_info_from_tcp_address():
|
||||
client = object.__new__(StageEngineCoreClient)
|
||||
client.stage_id = 0
|
||||
client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
|
||||
client._omni_kv_config = None
|
||||
client._kv_sender_info = None
|
||||
client._kv_sender_initialized = False
|
||||
client._kv_sender_host = client._resolve_contact_host()
|
||||
client._initialize_kv_sender_endpoint()
|
||||
|
||||
assert client.get_kv_sender_info() == {
|
||||
"host": "10.20.30.40",
|
||||
"zmq_port": 50151,
|
||||
}
|
||||
|
||||
|
||||
def test_stage_engine_core_client_falls_back_to_detected_ip_for_loopback(monkeypatch):
|
||||
client = object.__new__(StageEngineCoreClient)
|
||||
client.stage_id = 1
|
||||
client.client_addresses = {"input_address": "tcp://127.0.0.1:1234"}
|
||||
client._omni_kv_config = None
|
||||
client._kv_sender_info = None
|
||||
client._kv_sender_initialized = False
|
||||
monkeypatch.setattr(client, "_detect_local_ip", lambda: "192.168.0.12")
|
||||
client._kv_sender_host = client._resolve_contact_host()
|
||||
client._initialize_kv_sender_endpoint()
|
||||
|
||||
assert client.get_kv_sender_info() == {
|
||||
"host": "192.168.0.12",
|
||||
"zmq_port": 50152,
|
||||
}
|
||||
|
||||
|
||||
def test_stage_engine_core_client_uses_connector_config_for_sender_port():
|
||||
client = object.__new__(StageEngineCoreClient)
|
||||
client.stage_id = 3
|
||||
client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
|
||||
client._kv_sender_info = None
|
||||
client._kv_sender_initialized = False
|
||||
client._omni_kv_config = {
|
||||
"omni_from_stage": "3",
|
||||
"connector_config": {
|
||||
"type": "MooncakeTransferEngineConnector",
|
||||
"role": "sender",
|
||||
"host": "10.20.30.99",
|
||||
"zmq_port": 51000,
|
||||
},
|
||||
}
|
||||
client._kv_sender_host = client._resolve_contact_host()
|
||||
client._initialize_kv_sender_endpoint()
|
||||
|
||||
assert client.get_kv_sender_info() == {
|
||||
"host": "10.20.30.99",
|
||||
"zmq_port": 51103,
|
||||
}
|
||||
|
||||
|
||||
def test_stage_engine_core_client_preserves_explicit_loopback_sender_host():
|
||||
client = object.__new__(StageEngineCoreClient)
|
||||
client.stage_id = 2
|
||||
client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
|
||||
client._kv_sender_info = None
|
||||
client._kv_sender_initialized = False
|
||||
client._omni_kv_config = {
|
||||
"omni_from_stage": "2",
|
||||
"connector_config": {
|
||||
"type": "MooncakeTransferEngineConnector",
|
||||
"role": "sender",
|
||||
"host": "127.0.0.1",
|
||||
"zmq_port": 51000,
|
||||
},
|
||||
}
|
||||
client._kv_sender_host = client._resolve_contact_host()
|
||||
client._initialize_kv_sender_endpoint()
|
||||
|
||||
assert client.get_kv_sender_info() == {
|
||||
"host": "127.0.0.1",
|
||||
"zmq_port": 51102,
|
||||
}
|
||||
|
||||
|
||||
def test_forward_to_diffusion_attaches_kv_sender_info():
|
||||
orchestrator = object.__new__(Orchestrator)
|
||||
sender_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151})
|
||||
diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
|
||||
|
||||
orchestrator.num_stages = 2
|
||||
orchestrator.stage_clients = [sender_stage, diffusion_stage]
|
||||
orchestrator._companion_map = {}
|
||||
orchestrator.stage_vllm_configs = [None, None]
|
||||
orchestrator.output_processors = [None, None]
|
||||
|
||||
params = OmniDiffusionSamplingParams()
|
||||
req_state = OrchestratorRequestState(
|
||||
request_id="req-1",
|
||||
prompt={"prompt": "hello"},
|
||||
sampling_params_list=[SamplingParams(max_tokens=4), params],
|
||||
final_stage_id=1,
|
||||
)
|
||||
|
||||
output = SimpleNamespace(request_id="req-1", finished=True)
|
||||
asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-1", 0, output, req_state))
|
||||
|
||||
assert sender_stage.engine_outputs == [output]
|
||||
assert diffusion_stage.calls[0]["request_id"] == "req-1"
|
||||
assert diffusion_stage.calls[0]["kv_sender_info"] == {
|
||||
0: {"host": "10.0.0.2", "zmq_port": 50151},
|
||||
}
|
||||
assert req_state.stage_submit_ts[1] > 0
|
||||
|
||||
|
||||
def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info():
|
||||
orchestrator = object.__new__(Orchestrator)
|
||||
source_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151})
|
||||
previous_stage = _DummySenderStage({"host": "10.0.0.9", "zmq_port": 59999})
|
||||
diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
|
||||
|
||||
orchestrator.num_stages = 3
|
||||
orchestrator.stage_clients = [source_stage, previous_stage, diffusion_stage]
|
||||
orchestrator._companion_map = {}
|
||||
orchestrator.stage_vllm_configs = [None, None, None]
|
||||
orchestrator.output_processors = [None, None, None]
|
||||
|
||||
params = OmniDiffusionSamplingParams()
|
||||
req_state = OrchestratorRequestState(
|
||||
request_id="req-3",
|
||||
prompt={"prompt": "hello"},
|
||||
sampling_params_list=[SamplingParams(max_tokens=4), SamplingParams(max_tokens=4), params],
|
||||
final_stage_id=2,
|
||||
)
|
||||
|
||||
output = SimpleNamespace(request_id="req-3", finished=True)
|
||||
asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-3", 1, output, req_state))
|
||||
|
||||
assert previous_stage.engine_outputs == [output]
|
||||
assert diffusion_stage.calls[0]["kv_sender_info"] == {
|
||||
0: {"host": "10.0.0.2", "zmq_port": 50151},
|
||||
}
|
||||
|
||||
|
||||
def test_prewarm_diffusion_attaches_kv_sender_info():
|
||||
orchestrator = object.__new__(Orchestrator)
|
||||
sender_stage = _DummySenderStage({"host": "10.0.0.3", "zmq_port": 50151})
|
||||
diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
|
||||
|
||||
orchestrator.stage_clients = [sender_stage, diffusion_stage]
|
||||
orchestrator.num_stages = 2
|
||||
|
||||
req_state = OrchestratorRequestState(
|
||||
request_id="req-2",
|
||||
prompt={"prompt": "hello"},
|
||||
sampling_params_list=[SamplingParams(max_tokens=4), OmniDiffusionSamplingParams()],
|
||||
final_stage_id=1,
|
||||
)
|
||||
|
||||
stage0_request = SimpleNamespace(prompt_token_ids=[1, 2, 3])
|
||||
asyncio.run(Orchestrator._prewarm_async_chunk_stages(orchestrator, "req-2", stage0_request, req_state))
|
||||
|
||||
assert diffusion_stage.calls[0]["request_id"] == "req-2"
|
||||
assert diffusion_stage.calls[0]["kv_sender_info"] == {
|
||||
0: {"host": "10.0.0.3", "zmq_port": 50151},
|
||||
}
|
||||
assert req_state.stage_submit_ts[1] > 0
|
||||
@@ -12,6 +12,7 @@ import pytest
|
||||
import torch
|
||||
|
||||
# ── Load modules without triggering vllm_omni.__init__ ─────────────
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_ENGINE_DIR = Path(__file__).resolve().parents[2] / "vllm_omni" / "engine"
|
||||
|
||||
|
||||
@@ -178,7 +178,7 @@ def test_client(mock_async_diffusion):
|
||||
)
|
||||
app.state.args = Namespace(
|
||||
default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}',
|
||||
max_generated_image_size=4096, # 64*64
|
||||
max_generated_image_size=1024 * 1792,
|
||||
)
|
||||
|
||||
return TestClient(app)
|
||||
@@ -245,7 +245,7 @@ def async_omni_stage_configs_only_client():
|
||||
# AsyncOmni exposes stage_configs on the engine instance.
|
||||
app.state.args = Namespace(
|
||||
default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
|
||||
max_generated_image_size=4096, # 64*64
|
||||
max_generated_image_size=1024 * 1792,
|
||||
)
|
||||
return TestClient(app)
|
||||
|
||||
@@ -392,6 +392,18 @@ def test_image_edits_async_omni_stage_configs_only(async_omni_stage_configs_only
|
||||
assert len(captured) == 2
|
||||
|
||||
|
||||
def test_generate_images_max_size_rejected(async_omni_test_client):
|
||||
"""Test that a size exceeding max_generated_image_size returns 400."""
|
||||
response = async_omni_test_client.post(
|
||||
"/v1/images/generations",
|
||||
json={
|
||||
"prompt": "a cat",
|
||||
"size": "2048x2048", # 4,194,304 pixels > max_generated_image_size (1,048,576)
|
||||
},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
|
||||
|
||||
def test_generate_multiple_images(test_client):
|
||||
"""Test generating multiple images"""
|
||||
response = test_client.post(
|
||||
@@ -982,12 +994,13 @@ def test_image_edit_parameter_default_single_stage(test_client):
|
||||
assert captured_sampling_params.num_inference_steps == 4
|
||||
assert captured_sampling_params.guidance_scale == 7.5
|
||||
|
||||
# Size exceeding max_generated_image_size (1024*1792) returns 400
|
||||
response = test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes_1)],
|
||||
data={
|
||||
"prompt": "hello world.",
|
||||
"size": "96x96",
|
||||
"size": "2048x2048",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
|
||||
@@ -385,3 +385,18 @@ class TestStreamingSpeechWebSocket:
|
||||
|
||||
speech_service.engine_client.abort.assert_awaited_once_with("req-abort")
|
||||
assert websocket.send_json.await_count == 2
|
||||
|
||||
|
||||
class TestGeneratePcmChunksContract:
|
||||
"""Guard: _generate_pcm_chunks must exist on OmniOpenAIServingSpeech.
|
||||
|
||||
The WebSocket handler calls speech_service._generate_pcm_chunks()
|
||||
at runtime. If the method is removed, all WS TTS streaming breaks
|
||||
with an AttributeError. This test catches that at CI time.
|
||||
"""
|
||||
|
||||
def test_generate_pcm_chunks_defined(self):
|
||||
assert hasattr(OmniOpenAIServingSpeech, "_generate_pcm_chunks")
|
||||
assert asyncio.iscoroutinefunction(OmniOpenAIServingSpeech._generate_pcm_chunks) or callable(
|
||||
OmniOpenAIServingSpeech._generate_pcm_chunks
|
||||
)
|
||||
@@ -4,7 +4,7 @@ import pytest
|
||||
|
||||
from vllm_omni.entrypoints.openai.text_splitter import SentenceSplitter
|
||||
|
||||
pytestmark = [pytest.mark.openai, pytest.mark.speech]
|
||||
pytestmark = [pytest.mark.openai, pytest.mark.speech, pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class TestSentenceSplitterEnglish:
|
||||
|
||||
@@ -6,8 +6,6 @@ from pytest_mock import MockerFixture
|
||||
|
||||
from vllm_omni.entrypoints.stage_utils import set_stage_devices
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_dummy_torch(call_log):
|
||||
class _Props:
|
||||
@@ -55,6 +53,8 @@ def _make_mock_platform(mocker, device_type: str = "cuda", env_var: str = "CUDA_
|
||||
return mock_platform
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
|
||||
def test_set_stage_devices_respects_logical_ids(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
|
||||
# Preserve an existing logical mapping and ensure devices "0,1" map through it.
|
||||
@@ -75,6 +75,8 @@ def test_set_stage_devices_respects_logical_ids(mocker: MockerFixture, monkeypat
|
||||
assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
|
||||
def test_set_stage_devices_handles_not_enough_devices(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
|
||||
# Preserve an existing logical mapping and ensure devices "0,1" map through it.
|
||||
@@ -90,9 +92,10 @@ def test_set_stage_devices_handles_not_enough_devices(mocker: MockerFixture, mon
|
||||
mock_platform,
|
||||
)
|
||||
|
||||
# Raise since we need 4 GPUs, but we only have 2 visible
|
||||
with pytest.raises(ValueError):
|
||||
set_stage_devices(stage_id=0, devices="0,1,2,3")
|
||||
# Keep the logical mapping and resolve to the visible subset.
|
||||
set_stage_devices(stage_id=0, devices="0,1,2,3")
|
||||
|
||||
assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
|
||||
|
||||
@@ -8,6 +8,8 @@ import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from tests.utils import hardware_test
|
||||
|
||||
|
||||
class TestPreLookaheadLayer:
|
||||
"""Tests for PreLookaheadLayer."""
|
||||
@@ -18,6 +20,8 @@ class TestPreLookaheadLayer:
|
||||
|
||||
return PreLookaheadLayer(in_channels=512, channels=512, pre_lookahead_len=3)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_forward_shape(self, layer):
|
||||
"""Test that output shape matches input shape."""
|
||||
batch, seq_len, channels = 2, 10, 512
|
||||
@@ -27,6 +31,8 @@ class TestPreLookaheadLayer:
|
||||
|
||||
assert out.shape == x.shape
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_forward_with_context(self, layer):
|
||||
"""Test forward with context for streaming."""
|
||||
batch, seq_len, channels = 1, 10, 512
|
||||
@@ -38,6 +44,8 @@ class TestPreLookaheadLayer:
|
||||
|
||||
assert out.shape == x.shape
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_residual_connection(self, layer):
|
||||
"""Test that residual connection is applied."""
|
||||
batch, seq_len, channels = 1, 5, 512
|
||||
@@ -59,6 +67,8 @@ class TestDiTAttention:
|
||||
|
||||
return DiTAttention(dim=512, heads=8, dim_head=64, dropout=0.0)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_forward_shape(self, attention):
|
||||
"""Test attention output shape."""
|
||||
batch, seq_len, dim = 2, 16, 512
|
||||
@@ -68,6 +78,8 @@ class TestDiTAttention:
|
||||
|
||||
assert out.shape == x.shape
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_forward_with_mask(self, attention):
|
||||
"""Test attention with mask."""
|
||||
batch, seq_len, dim = 2, 16, 512
|
||||
@@ -81,6 +93,8 @@ class TestDiTAttention:
|
||||
# Masked positions should be zero
|
||||
assert torch.allclose(out[:, -3:], torch.zeros_like(out[:, -3:]))
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_qkv_projections(self, attention):
|
||||
"""Test that Q/K/V projections exist and have correct dimensions."""
|
||||
assert hasattr(attention, "to_q")
|
||||
@@ -100,6 +114,8 @@ class TestDiTBlock:
|
||||
|
||||
return DiTBlock(dim=512, heads=8, dim_head=64, ff_mult=4, dropout=0.0)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_forward_shape(self, block):
|
||||
"""Test block output shape."""
|
||||
batch, seq_len, dim = 2, 16, 512
|
||||
@@ -110,6 +126,8 @@ class TestDiTBlock:
|
||||
|
||||
assert out.shape == x.shape
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_adalayernorm_modulation(self, block):
|
||||
"""Test that AdaLayerNorm modulates based on timestep."""
|
||||
batch, seq_len, dim = 1, 8, 512
|
||||
@@ -144,6 +162,8 @@ class TestDiT:
|
||||
long_skip_connection=True,
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_forward_shape(self, dit):
|
||||
"""Test DiT forward output shape."""
|
||||
batch, mel_dim, seq_len = 1, 80, 32
|
||||
@@ -158,6 +178,8 @@ class TestDiT:
|
||||
|
||||
assert out.shape == (batch, mel_dim, seq_len)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_timestep_embedding(self, dit):
|
||||
"""Test that different timesteps produce different outputs."""
|
||||
batch, mel_dim, seq_len = 1, 80, 16
|
||||
@@ -190,6 +212,8 @@ class TestCFM:
|
||||
|
||||
return DummyEstimator()
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_causal_conditional_cfm_forward(self, dummy_estimator):
|
||||
"""Test CausalConditionalCFM forward pass."""
|
||||
from omegaconf import DictConfig
|
||||
@@ -228,6 +252,8 @@ class TestCFM:
|
||||
class TestSDPAFallback:
|
||||
"""Test SDPA fallback for float32 inputs."""
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_float32_uses_sdpa(self):
|
||||
"""Test that float32 inputs use SDPA fallback."""
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
|
||||
@@ -5,6 +5,8 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class TestMakePadMask:
|
||||
"""Tests for make_pad_mask utility."""
|
||||
|
||||
@@ -7,6 +7,7 @@ Regression tests for https://github.com/vllm-project/vllm-omni/issues/1862
|
||||
|
||||
from collections.abc import Mapping
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.config.stage_config import StageConfigFactory
|
||||
@@ -15,6 +16,8 @@ from vllm_omni.diffusion.data import (
|
||||
OmniDiffusionConfig,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
|
||||
"""Simulate the real path: create_default_diffusion → OmniDiffusionConfig.
|
||||
|
||||
@@ -1,185 +0,0 @@
|
||||
"""Tests for CUDAGraphWrapper.__getattr__ performance optimization.
|
||||
|
||||
This module tests that the patched CUDAGraphWrapper avoids expensive __repr__
|
||||
calls when hasattr() is used for non-existent attributes. The original vLLM
|
||||
implementation includes {self.runnable} in the AttributeError message, which
|
||||
triggers model tree traversal and can take ~6ms on large models.
|
||||
"""
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm_omni.worker.gpu_model_runner import CUDAGraphWrapper
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class SlowReprModel(nn.Module):
|
||||
"""A mock model with artificially slow __repr__ to detect unwanted calls."""
|
||||
|
||||
def __init__(self, repr_delay_ms: float = 10.0):
|
||||
super().__init__()
|
||||
self.linear = nn.Linear(16, 16)
|
||||
self.repr_delay_ms = repr_delay_ms
|
||||
self.repr_call_count = 0
|
||||
|
||||
def forward(self, x):
|
||||
return self.linear(x)
|
||||
|
||||
def __repr__(self):
|
||||
self.repr_call_count += 1
|
||||
# Simulate expensive repr by sleeping
|
||||
time.sleep(self.repr_delay_ms / 1000.0)
|
||||
return f"SlowReprModel(delay={self.repr_delay_ms}ms)"
|
||||
|
||||
|
||||
class MockCUDAGraphWrapper:
|
||||
"""A minimal mock that mimics CUDAGraphWrapper structure for CPU testing."""
|
||||
|
||||
def __init__(self, runnable):
|
||||
# Store in __dict__ directly to avoid triggering __getattr__
|
||||
object.__setattr__(self, "runnable", runnable)
|
||||
|
||||
def __getattr__(self, key: str):
|
||||
# This is the optimized implementation we're testing
|
||||
runnable = object.__getattribute__(self, "runnable")
|
||||
if hasattr(runnable, key):
|
||||
return getattr(runnable, key)
|
||||
# Key optimization: DO NOT include {self.runnable} in error message
|
||||
# as it triggers expensive __repr__ on large models
|
||||
raise AttributeError(f"Attribute {key} not exists in the runnable of cudagraph wrapper")
|
||||
|
||||
|
||||
def test_hasattr_nonexistent_does_not_trigger_repr():
|
||||
"""Verify that hasattr for non-existent attributes doesn't call __repr__."""
|
||||
model = SlowReprModel(repr_delay_ms=100.0) # Very slow repr
|
||||
wrapper = MockCUDAGraphWrapper(model)
|
||||
|
||||
# Reset counter
|
||||
model.repr_call_count = 0
|
||||
|
||||
# Call hasattr for non-existent attribute multiple times
|
||||
for _ in range(10):
|
||||
result = hasattr(wrapper, "nonexistent_attribute_xyz")
|
||||
assert result is False
|
||||
|
||||
# __repr__ should never have been called
|
||||
assert model.repr_call_count == 0, (
|
||||
f"__repr__ was called {model.repr_call_count} times when checking "
|
||||
"for non-existent attributes. This indicates the AttributeError "
|
||||
"message contains {self.runnable} which triggers expensive repr."
|
||||
)
|
||||
|
||||
|
||||
def test_hasattr_nonexistent_is_fast():
|
||||
"""Verify that hasattr for non-existent attributes is fast (<1ms per call)."""
|
||||
model = SlowReprModel(repr_delay_ms=100.0)
|
||||
wrapper = MockCUDAGraphWrapper(model)
|
||||
|
||||
num_iterations = 100
|
||||
start = time.perf_counter()
|
||||
for _ in range(num_iterations):
|
||||
hasattr(wrapper, "nonexistent_attribute_xyz")
|
||||
elapsed_ms = (time.perf_counter() - start) * 1000
|
||||
|
||||
avg_ms = elapsed_ms / num_iterations
|
||||
# If __repr__ were being called, each would take ~100ms
|
||||
# We expect <1ms per call with the fix
|
||||
assert avg_ms < 1.0, (
|
||||
f"hasattr for non-existent attribute took {avg_ms:.2f}ms on average. "
|
||||
"Expected <1ms. This suggests __repr__ is being triggered."
|
||||
)
|
||||
|
||||
|
||||
def test_hasattr_existing_attribute_works():
|
||||
"""Verify that hasattr for existing attributes returns True and works correctly."""
|
||||
model = SlowReprModel()
|
||||
wrapper = MockCUDAGraphWrapper(model)
|
||||
|
||||
# 'forward' exists on nn.Module
|
||||
assert hasattr(wrapper, "forward") is True
|
||||
|
||||
# 'linear' exists on our model
|
||||
assert hasattr(wrapper, "linear") is True
|
||||
|
||||
# Can actually access the attribute
|
||||
linear = wrapper.linear
|
||||
assert isinstance(linear, nn.Linear)
|
||||
|
||||
|
||||
def test_getattr_existing_attribute_returns_value():
|
||||
"""Verify that getattr for existing attributes returns the correct value."""
|
||||
model = SlowReprModel()
|
||||
wrapper = MockCUDAGraphWrapper(model)
|
||||
|
||||
# Access forward method
|
||||
forward_method = wrapper.forward
|
||||
assert callable(forward_method)
|
||||
|
||||
# Access linear layer
|
||||
linear = wrapper.linear
|
||||
assert isinstance(linear, nn.Linear)
|
||||
assert linear.in_features == 16
|
||||
assert linear.out_features == 16
|
||||
|
||||
|
||||
def test_getattr_nonexistent_raises_attribute_error():
|
||||
"""Verify that getattr for non-existent attributes raises AttributeError."""
|
||||
model = SlowReprModel()
|
||||
wrapper = MockCUDAGraphWrapper(model)
|
||||
|
||||
with pytest.raises(AttributeError) as exc_info:
|
||||
_ = wrapper.nonexistent_attribute
|
||||
|
||||
# Verify error message format (should NOT contain model repr)
|
||||
error_msg = str(exc_info.value)
|
||||
assert "nonexistent_attribute" in error_msg
|
||||
assert "cudagraph wrapper" in error_msg
|
||||
# Should NOT contain the slow repr output
|
||||
assert "SlowReprModel(delay=" not in error_msg
|
||||
|
||||
|
||||
def test_attribute_error_message_does_not_contain_runnable_repr():
|
||||
"""Explicitly verify the error message doesn't trigger runnable repr."""
|
||||
model = SlowReprModel(repr_delay_ms=100.0)
|
||||
wrapper = MockCUDAGraphWrapper(model)
|
||||
model.repr_call_count = 0
|
||||
|
||||
try:
|
||||
_ = wrapper.nonexistent_attr
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
# __repr__ should not have been called during error construction
|
||||
assert model.repr_call_count == 0, (
|
||||
"AttributeError message construction triggered __repr__. The error message should not include {self.runnable}."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
def test_real_cudagraph_wrapper_hasattr_performance():
|
||||
"""Test the actual CUDAGraphWrapper from vllm_omni (requires CUDA)."""
|
||||
from vllm.config import CUDAGraphMode
|
||||
|
||||
model = SlowReprModel(repr_delay_ms=50.0).cuda()
|
||||
model.repr_call_count = 0
|
||||
|
||||
# Create actual CUDAGraphWrapper
|
||||
try:
|
||||
wrapper = CUDAGraphWrapper(model, runtime_mode=CUDAGraphMode.NONE)
|
||||
except Exception:
|
||||
pytest.skip("Could not create CUDAGraphWrapper")
|
||||
|
||||
# Test hasattr performance
|
||||
num_iterations = 50
|
||||
start = time.perf_counter()
|
||||
for _ in range(num_iterations):
|
||||
hasattr(wrapper, "nonexistent_xyz")
|
||||
elapsed_ms = (time.perf_counter() - start) * 1000
|
||||
|
||||
avg_ms = elapsed_ms / num_iterations
|
||||
assert avg_ms < 1.0, f"Real CUDAGraphWrapper hasattr took {avg_ms:.2f}ms avg. Expected <1ms with the optimization."
|
||||
assert model.repr_call_count == 0, f"__repr__ called {model.repr_call_count} times"
|
||||
+17
-1
@@ -24,11 +24,27 @@ except ModuleNotFoundError as exc: # pragma: no cover - optional dependency
|
||||
from vllm_omni.transformers_utils import configs as _configs # noqa: F401, E402
|
||||
|
||||
from .config import OmniModelConfig
|
||||
from .entrypoints import AsyncOmni, Omni
|
||||
|
||||
from .version import __version__, __version_tuple__ # isort:skip
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
# Lazy import for AsyncOmni and Omni to avoid pulling in heavy
|
||||
# dependencies (vllm model_loader → fused_moe → pynvml) at package
|
||||
# import time. This prevents crashes in lightweight subprocesses
|
||||
# (e.g. model-architecture inspection) that lack a CUDA context.
|
||||
# See: https://github.com/vllm-project/vllm-omni/issues/1793
|
||||
if name == "AsyncOmni":
|
||||
from .entrypoints.async_omni import AsyncOmni
|
||||
|
||||
return AsyncOmni
|
||||
if name == "Omni":
|
||||
from .entrypoints.omni import Omni
|
||||
|
||||
return Omni
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"__version__",
|
||||
"__version_tuple__",
|
||||
|
||||
@@ -185,7 +185,7 @@ def calculate_metrics(
|
||||
# Note : this may inflate the output token count slightly
|
||||
output_len = len(tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids)
|
||||
actual_output_lens.append(output_len)
|
||||
total_input += input_requests[i].prompt_len
|
||||
total_input += outputs[i].prompt_len
|
||||
tpot = 0
|
||||
if output_len > 1:
|
||||
latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft
|
||||
|
||||
@@ -190,6 +190,10 @@ async def async_request_openai_chat_omni_completions(
|
||||
if metrics := data.get("metrics"):
|
||||
output.output_tokens = metrics.get("num_tokens_out", 0)
|
||||
|
||||
if usage := data.get("usage"):
|
||||
if (pt := usage.get("prompt_tokens")) is not None:
|
||||
output.prompt_len = pt
|
||||
|
||||
output.latency = timestamp - st
|
||||
output.generated_text = generated_text
|
||||
if generated_audio is not None:
|
||||
|
||||
@@ -209,3 +209,13 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
layout="BNSD",
|
||||
)
|
||||
return output
|
||||
|
||||
def forward_musa(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata = None,
|
||||
) -> torch.Tensor:
|
||||
# XXX (MUSA): MUSA uses the same implementation as XPU (mate only provides flash_attn_varlen_func)
|
||||
return self.forward_xpu(query, key, value, attn_metadata)
|
||||
@@ -12,6 +12,8 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# Adapted from https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_flash_attention_utils.py
|
||||
from functools import lru_cache
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
@@ -38,8 +40,10 @@ elif current_omni_platform.is_xpu():
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
pass
|
||||
elif current_omni_platform.is_musa():
|
||||
# XXX (MUSA): Add MUSA-specific Flash Attention when available
|
||||
pass
|
||||
try:
|
||||
from mate import flash_attn_varlen_func # noqa: F401
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
pass
|
||||
else:
|
||||
# CUDA: try FA3 -> FA2 fallback chain
|
||||
# Try FA3 from fa3-fwd PyPI package
|
||||
@@ -76,6 +80,12 @@ else:
|
||||
HAS_FLASH_ATTN = flash_attn_func is not None or flash_attn_varlen_func is not None
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def is_mate_available() -> bool:
|
||||
"""Check if MATE (MUSA Flash Attention) is available."""
|
||||
return current_omni_platform.is_musa() and flash_attn_varlen_func is not None
|
||||
|
||||
|
||||
def _index_first_axis(tensor, indices):
|
||||
"""
|
||||
A local implementation of the PyTorch indexing operation `tensor[indices]` on the first axis,
|
||||
|
||||
@@ -48,16 +48,7 @@ def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
|
||||
coefficients=config.coefficients,
|
||||
)
|
||||
transformer = pipeline.bagel
|
||||
original_forward_flow = transformer._forward_flow
|
||||
|
||||
import types
|
||||
|
||||
def forward_alias(self, *args, **kwargs):
|
||||
return original_forward_flow(*args, **kwargs)
|
||||
|
||||
transformer.forward = types.MethodType(forward_alias, transformer)
|
||||
apply_teacache_hook(transformer, teacache_config)
|
||||
transformer._forward_flow = transformer.forward
|
||||
pipeline.transformer = transformer
|
||||
|
||||
logger.info(
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import types
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
@@ -74,15 +73,8 @@ class BagelAdapter:
|
||||
|
||||
@staticmethod
|
||||
def install_hook(transformer: Any, hook: DataCollectionHook) -> None:
|
||||
original_forward_flow = transformer._forward_flow
|
||||
|
||||
def forward_alias(self, *args, **kwargs):
|
||||
return original_forward_flow(*args, **kwargs)
|
||||
|
||||
transformer.forward = types.MethodType(forward_alias, transformer)
|
||||
registry = HookRegistry.get_or_create(transformer)
|
||||
registry.register_hook(hook._HOOK_NAME, hook)
|
||||
transformer._forward_flow = transformer.forward
|
||||
|
||||
|
||||
class StableAudioAdapter:
|
||||
|
||||
@@ -9,6 +9,7 @@ from dataclasses import dataclass, field, fields
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from pydantic import model_validator
|
||||
from typing_extensions import Self
|
||||
from vllm.config.utils import config
|
||||
@@ -701,10 +702,12 @@ class DiffusionOutput:
|
||||
Final output (after pipeline completion)
|
||||
"""
|
||||
|
||||
output: torch.Tensor | None = None
|
||||
trajectory_timesteps: list[torch.Tensor] | None = None
|
||||
trajectory_latents: torch.Tensor | None = None
|
||||
trajectory_decoded: list[torch.Tensor] | None = None
|
||||
# Fields may be replaced with SHM handle dicts by ipc.pack_diffusion_output_shm
|
||||
output: torch.Tensor | dict | None = None
|
||||
trajectory_timesteps: torch.Tensor | dict | None = None
|
||||
trajectory_latents: torch.Tensor | dict | None = None
|
||||
trajectory_log_probs: torch.Tensor | dict | None = None
|
||||
trajectory_decoded: list[Image.Image] | None = None
|
||||
error: str | None = None
|
||||
aborted: bool = False
|
||||
abort_message: str | None = None
|
||||
|
||||
@@ -145,8 +145,12 @@ class DiffusionEngine:
|
||||
postprocess_start_time = time.perf_counter()
|
||||
outputs = self.post_process_func(output_data) if self.post_process_func is not None else output_data
|
||||
audio_payload = None
|
||||
model_audio_sample_rate = None
|
||||
model_fps = None
|
||||
if isinstance(outputs, dict):
|
||||
audio_payload = outputs.get("audio")
|
||||
model_audio_sample_rate = outputs.get("audio_sample_rate")
|
||||
model_fps = outputs.get("fps")
|
||||
outputs = outputs.get("video", outputs)
|
||||
postprocess_time = time.perf_counter() - postprocess_start_time
|
||||
logger.info(f"Post-processing completed in {postprocess_time:.4f} seconds")
|
||||
@@ -192,6 +196,10 @@ class DiffusionEngine:
|
||||
prompt=prompt,
|
||||
metrics=metrics,
|
||||
latents=output.trajectory_latents,
|
||||
trajectory_latents=output.trajectory_latents,
|
||||
trajectory_timesteps=output.trajectory_timesteps,
|
||||
trajectory_log_probs=output.trajectory_log_probs,
|
||||
trajectory_decoded=output.trajectory_decoded,
|
||||
multimodal_output={"audio": request_audio_payload},
|
||||
final_output_type="audio",
|
||||
stage_durations=output.stage_durations,
|
||||
@@ -202,6 +210,10 @@ class DiffusionEngine:
|
||||
mm_output = {}
|
||||
if audio_payload is not None:
|
||||
mm_output["audio"] = audio_payload
|
||||
if model_audio_sample_rate is not None:
|
||||
mm_output["audio_sample_rate"] = model_audio_sample_rate
|
||||
if model_fps is not None:
|
||||
mm_output["fps"] = model_fps
|
||||
return [
|
||||
OmniRequestOutput.from_diffusion(
|
||||
request_id=request_id,
|
||||
@@ -209,6 +221,10 @@ class DiffusionEngine:
|
||||
prompt=prompt,
|
||||
metrics=metrics,
|
||||
latents=output.trajectory_latents,
|
||||
trajectory_latents=output.trajectory_latents,
|
||||
trajectory_timesteps=output.trajectory_timesteps,
|
||||
trajectory_log_probs=output.trajectory_log_probs,
|
||||
trajectory_decoded=output.trajectory_decoded,
|
||||
custom_output=output.custom_output or {},
|
||||
multimodal_output=mm_output,
|
||||
stage_durations=output.stage_durations,
|
||||
@@ -240,6 +256,10 @@ class DiffusionEngine:
|
||||
prompt=prompt,
|
||||
metrics=metrics,
|
||||
latents=output.trajectory_latents,
|
||||
trajectory_latents=output.trajectory_latents,
|
||||
trajectory_timesteps=output.trajectory_timesteps,
|
||||
trajectory_log_probs=output.trajectory_log_probs,
|
||||
trajectory_decoded=output.trajectory_decoded,
|
||||
multimodal_output={"audio": request_audio_payload},
|
||||
final_output_type="audio",
|
||||
stage_durations=output.stage_durations,
|
||||
@@ -260,6 +280,10 @@ class DiffusionEngine:
|
||||
if num_outputs == 1:
|
||||
sliced_audio = sliced_audio[0]
|
||||
mm_output["audio"] = sliced_audio
|
||||
if model_audio_sample_rate is not None:
|
||||
mm_output["audio_sample_rate"] = model_audio_sample_rate
|
||||
if model_fps is not None:
|
||||
mm_output["fps"] = model_fps
|
||||
results.append(
|
||||
OmniRequestOutput.from_diffusion(
|
||||
request_id=request_id,
|
||||
@@ -267,6 +291,10 @@ class DiffusionEngine:
|
||||
prompt=prompt,
|
||||
metrics=metrics,
|
||||
latents=output.trajectory_latents,
|
||||
trajectory_latents=output.trajectory_latents,
|
||||
trajectory_timesteps=output.trajectory_timesteps,
|
||||
trajectory_log_probs=output.trajectory_log_probs,
|
||||
trajectory_decoded=output.trajectory_decoded,
|
||||
custom_output=output.custom_output or {},
|
||||
multimodal_output=mm_output,
|
||||
stage_durations=output.stage_durations,
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from contextlib import nullcontext
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
@@ -15,11 +16,38 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
|
||||
GridSpec,
|
||||
TileTask,
|
||||
)
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
|
||||
class OmniAutoencoderKLWan(AutoencoderKLWan):
|
||||
def _execution_context(self):
|
||||
try:
|
||||
first_param = next(self.parameters())
|
||||
except StopIteration:
|
||||
return nullcontext()
|
||||
|
||||
dtype = first_param.dtype
|
||||
if dtype not in (torch.float16, torch.bfloat16):
|
||||
return nullcontext()
|
||||
|
||||
return current_omni_platform.create_autocast_context(
|
||||
device_type=first_param.device.type,
|
||||
dtype=dtype,
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
def encode(self, x: torch.Tensor, return_dict: bool = True):
|
||||
with self._execution_context():
|
||||
return super().encode(x, return_dict=return_dict)
|
||||
|
||||
def decode(self, z: torch.Tensor, return_dict: bool = True):
|
||||
with self._execution_context():
|
||||
return super().decode(z, return_dict=return_dict)
|
||||
|
||||
|
||||
class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
|
||||
@classmethod
|
||||
def from_pretrained(cls, *args: Any, **kwargs: Any):
|
||||
model = super().from_pretrained(*args, **kwargs)
|
||||
@@ -84,11 +112,12 @@ class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
|
||||
"""Decode a single latent tile into RGB space."""
|
||||
self.clear_cache()
|
||||
time = []
|
||||
for k in range(len(task.tensor)):
|
||||
self._conv_idx = [0]
|
||||
tile = self.post_quant_conv(task.tensor[k])
|
||||
decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0))
|
||||
time.append(decoded)
|
||||
with self._execution_context():
|
||||
for k in range(len(task.tensor)):
|
||||
self._conv_idx = [0]
|
||||
tile = self.post_quant_conv(task.tensor[k])
|
||||
decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0))
|
||||
time.append(decoded)
|
||||
result = torch.cat(time, dim=2)
|
||||
return result
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any
|
||||
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.utils.fa import is_mate_available
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -52,6 +53,10 @@ class PackagesEnvChecker:
|
||||
"""Check if flash attention is available and compatible."""
|
||||
platform = current_omni_platform
|
||||
|
||||
# MUSA uses MATE for flash attention
|
||||
if platform.is_musa():
|
||||
return is_mate_available()
|
||||
|
||||
# Flash attention requires CUDA-like platforms (CUDA or ROCm)
|
||||
if not platform.is_cuda_alike():
|
||||
return False
|
||||
|
||||
@@ -78,13 +78,29 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
|
||||
return tensor
|
||||
|
||||
|
||||
def _pack_tensor_if_large(val: torch.Tensor) -> torch.Tensor | dict:
|
||||
"""Replace a tensor with an SHM handle if it exceeds the threshold."""
|
||||
if val.nelement() * val.element_size() > _SHM_TENSOR_THRESHOLD:
|
||||
return _tensor_to_shm(val)
|
||||
return val
|
||||
|
||||
|
||||
def _unpack_if_shm_handle(val: object) -> object:
|
||||
"""Reconstruct a tensor from an SHM handle dict, or return as-is."""
|
||||
if isinstance(val, dict) and val.get("__tensor_shm__"):
|
||||
return _tensor_from_shm(val)
|
||||
return val
|
||||
|
||||
|
||||
def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
|
||||
if output.output is not None and isinstance(output.output, torch.Tensor):
|
||||
if output.output.nelement() * output.output.element_size() > _SHM_TENSOR_THRESHOLD:
|
||||
output.output = _tensor_to_shm(output.output)
|
||||
output.output = _pack_tensor_if_large(output.output)
|
||||
if output.trajectory_latents is not None and isinstance(output.trajectory_latents, torch.Tensor):
|
||||
if output.trajectory_latents.nelement() * output.trajectory_latents.element_size() > _SHM_TENSOR_THRESHOLD:
|
||||
output.trajectory_latents = _tensor_to_shm(output.trajectory_latents)
|
||||
output.trajectory_latents = _pack_tensor_if_large(output.trajectory_latents)
|
||||
if output.trajectory_timesteps is not None and isinstance(output.trajectory_timesteps, torch.Tensor):
|
||||
output.trajectory_timesteps = _pack_tensor_if_large(output.trajectory_timesteps)
|
||||
if output.trajectory_log_probs is not None and isinstance(output.trajectory_log_probs, torch.Tensor):
|
||||
output.trajectory_log_probs = _pack_tensor_if_large(output.trajectory_log_probs)
|
||||
return output
|
||||
|
||||
|
||||
@@ -104,10 +120,10 @@ def pack_diffusion_output_shm(output: object) -> object:
|
||||
|
||||
|
||||
def _unpack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
|
||||
if isinstance(output.output, dict) and output.output.get("__tensor_shm__"):
|
||||
output.output = _tensor_from_shm(output.output)
|
||||
if isinstance(output.trajectory_latents, dict) and output.trajectory_latents.get("__tensor_shm__"):
|
||||
output.trajectory_latents = _tensor_from_shm(output.trajectory_latents)
|
||||
output.output = _unpack_if_shm_handle(output.output)
|
||||
output.trajectory_latents = _unpack_if_shm_handle(output.trajectory_latents)
|
||||
output.trajectory_timesteps = _unpack_if_shm_handle(output.trajectory_timesteps)
|
||||
output.trajectory_log_probs = _unpack_if_shm_handle(output.trajectory_log_probs)
|
||||
return output
|
||||
|
||||
|
||||
|
||||
@@ -145,6 +145,14 @@ class RotaryEmbedding(CustomOp):
|
||||
) -> torch.Tensor:
|
||||
return self.forward_native(x, cos, sin)
|
||||
|
||||
def forward_musa(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return self.forward_native(x, cos, sin)
|
||||
|
||||
def forward_native(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
|
||||
@@ -1655,6 +1655,9 @@ class Bagel(nn.Module):
|
||||
cfg_img_past_key_values: NaiveCache | None = None,
|
||||
cfg_img_key_values_lens: torch.IntTensor | None = None,
|
||||
cfg_img_packed_key_value_indexes: torch.LongTensor | None = None,
|
||||
return_trajectory_latents: bool = False,
|
||||
scheduler: object | None = None,
|
||||
scheduler_kwargs: dict | None = None,
|
||||
):
|
||||
x_t = packed_init_noises
|
||||
|
||||
@@ -1663,6 +1666,14 @@ class Bagel(nn.Module):
|
||||
dts = timesteps[:-1] - timesteps[1:]
|
||||
timesteps = timesteps[:-1]
|
||||
|
||||
# Optional trajectory recording for RL rollout data collection
|
||||
trajectory_latents: list[torch.Tensor] | None = [] if return_trajectory_latents else None
|
||||
trajectory_timesteps: list[torch.Tensor] | None = [] if return_trajectory_latents else None
|
||||
trajectory_log_probs: list[torch.Tensor] | None = (
|
||||
[] if (return_trajectory_latents and scheduler is not None) else None
|
||||
)
|
||||
_sched_kw = scheduler_kwargs or {}
|
||||
|
||||
use_cfg_text = cfg_text_scale > 1.0
|
||||
use_cfg_img = cfg_img_scale > 1.0
|
||||
|
||||
@@ -1699,6 +1710,9 @@ class Bagel(nn.Module):
|
||||
cfg_img_past_key_values=cfg_img_past_key_values,
|
||||
cfg_img_key_values_lens=cfg_img_key_values_lens,
|
||||
cfg_img_packed_key_value_indexes=cfg_img_packed_key_value_indexes,
|
||||
return_trajectory_latents=return_trajectory_latents,
|
||||
scheduler=scheduler,
|
||||
scheduler_kwargs=scheduler_kwargs,
|
||||
)
|
||||
|
||||
# ── SP + CFG: sequential single-branch forwards ──
|
||||
@@ -1720,7 +1734,7 @@ class Bagel(nn.Module):
|
||||
packed_seqlens=packed_seqlens,
|
||||
)
|
||||
|
||||
v_t = self._forward_flow_single_branch(
|
||||
v_t = self.forward_single_branch(
|
||||
**common,
|
||||
packed_indexes=packed_indexes,
|
||||
packed_position_ids=packed_position_ids,
|
||||
@@ -1730,7 +1744,7 @@ class Bagel(nn.Module):
|
||||
)
|
||||
|
||||
if cfg_text_scale_ > 1.0:
|
||||
cfg_text_v_t = self._forward_flow_single_branch(
|
||||
cfg_text_v_t = self.forward_single_branch(
|
||||
**common,
|
||||
packed_indexes=cfg_text_packed_query_indexes,
|
||||
packed_position_ids=cfg_text_packed_position_ids,
|
||||
@@ -1740,7 +1754,7 @@ class Bagel(nn.Module):
|
||||
)
|
||||
cfg_img_v_t = None
|
||||
if cfg_img_scale_ > 1.0:
|
||||
cfg_img_v_t = self._forward_flow_single_branch(
|
||||
cfg_img_v_t = self.forward_single_branch(
|
||||
**common,
|
||||
packed_indexes=cfg_img_packed_query_indexes,
|
||||
packed_position_ids=cfg_img_packed_position_ids,
|
||||
@@ -1758,16 +1772,25 @@ class Bagel(nn.Module):
|
||||
cfg_renorm_min,
|
||||
)
|
||||
|
||||
x_t = x_t - v_t.to(x_t.device) * dts[i]
|
||||
if scheduler is not None:
|
||||
out = scheduler.step(v_t.to(x_t.device), timesteps[i], x_t, dts[i], **_sched_kw)
|
||||
x_t = out.prev_sample
|
||||
if trajectory_log_probs is not None and out.log_prob is not None:
|
||||
trajectory_log_probs.append(out.log_prob)
|
||||
else:
|
||||
x_t = x_t - v_t.to(x_t.device) * dts[i]
|
||||
if return_trajectory_latents:
|
||||
trajectory_latents.append(x_t.clone())
|
||||
trajectory_timesteps.append(timesteps[i] - dts[i])
|
||||
|
||||
unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
|
||||
return unpacked_latent
|
||||
return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
|
||||
|
||||
# ── SP without CFG: direct single-branch loop ──
|
||||
if use_sp:
|
||||
for i, t in enumerate(timesteps):
|
||||
timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
|
||||
v_t = self._forward_flow_single_branch(
|
||||
v_t = self.forward_single_branch(
|
||||
x_t=x_t,
|
||||
timestep=timestep,
|
||||
packed_vae_token_indexes=packed_vae_token_indexes,
|
||||
@@ -1781,10 +1804,20 @@ class Bagel(nn.Module):
|
||||
past_key_values=past_key_values,
|
||||
packed_key_value_indexes=packed_key_value_indexes,
|
||||
)
|
||||
x_t = x_t - v_t.to(x_t.device) * dts[i]
|
||||
if scheduler is not None:
|
||||
out = scheduler.step(v_t.to(x_t.device), timesteps[i], x_t, dts[i], **_sched_kw)
|
||||
x_t = out.prev_sample
|
||||
out_log_prob = getattr(out, "log_prob", None)
|
||||
if trajectory_log_probs is not None and out_log_prob is not None:
|
||||
trajectory_log_probs.append(out_log_prob)
|
||||
else:
|
||||
x_t = x_t - v_t.to(x_t.device) * dts[i]
|
||||
if return_trajectory_latents:
|
||||
trajectory_latents.append(x_t.clone())
|
||||
trajectory_timesteps.append(timesteps[i] - dts[i])
|
||||
|
||||
unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
|
||||
return unpacked_latent
|
||||
return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
|
||||
|
||||
# ── Batched CFG mode (cfg_parallel_size=1, no SP) ──
|
||||
cfg_batched = None
|
||||
@@ -1850,7 +1883,7 @@ class Bagel(nn.Module):
|
||||
else:
|
||||
cfg_text_scale_ = 1.0
|
||||
cfg_img_scale_ = 1.0
|
||||
v_t = self._forward_flow(
|
||||
v_t = self.forward(
|
||||
x_t=x_t,
|
||||
timestep=timestep,
|
||||
packed_vae_token_indexes=packed_vae_token_indexes,
|
||||
@@ -1870,10 +1903,19 @@ class Bagel(nn.Module):
|
||||
cfg_batched=cfg_batched,
|
||||
)
|
||||
|
||||
x_t = x_t - v_t.to(x_t.device) * dts[i] # velocity pointing from data to noise
|
||||
if scheduler is not None:
|
||||
out = scheduler.step(v_t.to(x_t.device), timesteps[i], x_t, dts[i], **_sched_kw)
|
||||
x_t = out.prev_sample
|
||||
if trajectory_log_probs is not None and out.log_prob is not None:
|
||||
trajectory_log_probs.append(out.log_prob)
|
||||
else:
|
||||
x_t = x_t - v_t.to(x_t.device) * dts[i] # velocity pointing from data to noise
|
||||
if return_trajectory_latents:
|
||||
trajectory_latents.append(x_t.clone())
|
||||
trajectory_timesteps.append(timesteps[i] - dts[i])
|
||||
|
||||
unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
|
||||
return unpacked_latent
|
||||
return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
|
||||
|
||||
def _generate_image_parallel(
|
||||
self,
|
||||
@@ -1905,6 +1947,9 @@ class Bagel(nn.Module):
|
||||
cfg_img_past_key_values: NaiveCache | None,
|
||||
cfg_img_key_values_lens: torch.IntTensor | None,
|
||||
cfg_img_packed_key_value_indexes: torch.LongTensor | None,
|
||||
return_trajectory_latents: bool = False,
|
||||
scheduler: object | None = None,
|
||||
scheduler_kwargs: dict | None = None,
|
||||
):
|
||||
"""CFG parallel denoising loop: each rank computes one CFG branch.
|
||||
|
||||
@@ -1961,13 +2006,20 @@ class Bagel(nn.Module):
|
||||
else:
|
||||
raise RuntimeError(f"Unexpected cfg_rank={cfg_rank} for Bagel 3-branch CFG parallel")
|
||||
|
||||
trajectory_latents: list[torch.Tensor] | None = [] if return_trajectory_latents else None
|
||||
trajectory_timesteps: list[torch.Tensor] | None = [] if return_trajectory_latents else None
|
||||
trajectory_log_probs: list[torch.Tensor] | None = (
|
||||
[] if (return_trajectory_latents and scheduler is not None) else None
|
||||
)
|
||||
_sched_kw = scheduler_kwargs or {}
|
||||
|
||||
for i, t in enumerate(timesteps):
|
||||
timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
|
||||
use_cfg_this_step = t > cfg_interval[0] and t <= cfg_interval[1] and cfg_text_scale > 1.0
|
||||
|
||||
if use_cfg_this_step:
|
||||
# CFG interval: each rank computes its own branch
|
||||
local_v_t = self._forward_flow_single_branch(
|
||||
local_v_t = self.forward_single_branch(
|
||||
x_t=x_t,
|
||||
timestep=timestep,
|
||||
packed_vae_token_indexes=packed_vae_token_indexes,
|
||||
@@ -1994,7 +2046,7 @@ class Bagel(nn.Module):
|
||||
)
|
||||
else:
|
||||
# Outside CFG interval: all ranks compute with gen inputs, no comm
|
||||
v_t = self._forward_flow_single_branch(
|
||||
v_t = self.forward_single_branch(
|
||||
x_t=x_t,
|
||||
timestep=timestep,
|
||||
packed_vae_token_indexes=packed_vae_token_indexes,
|
||||
@@ -2009,10 +2061,19 @@ class Bagel(nn.Module):
|
||||
packed_key_value_indexes=packed_key_value_indexes,
|
||||
)
|
||||
|
||||
x_t = x_t - v_t.to(x_t.device) * dts[i]
|
||||
if scheduler is not None:
|
||||
out = scheduler.step(v_t.to(x_t.device), timesteps[i], x_t, dts[i], **_sched_kw)
|
||||
x_t = out.prev_sample
|
||||
if trajectory_log_probs is not None and out.log_prob is not None:
|
||||
trajectory_log_probs.append(out.log_prob)
|
||||
else:
|
||||
x_t = x_t - v_t.to(x_t.device) * dts[i]
|
||||
if return_trajectory_latents:
|
||||
trajectory_latents.append(x_t.clone())
|
||||
trajectory_timesteps.append(timesteps[i] - dts[i])
|
||||
|
||||
unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
|
||||
return unpacked_latent
|
||||
return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
|
||||
|
||||
@staticmethod
|
||||
def _combine_cfg(
|
||||
@@ -2067,7 +2128,7 @@ class Bagel(nn.Module):
|
||||
|
||||
return v_t
|
||||
|
||||
def _forward_flow_single_branch(
|
||||
def forward_single_branch(
|
||||
self,
|
||||
x_t: torch.Tensor,
|
||||
timestep: torch.LongTensor,
|
||||
@@ -2197,7 +2258,7 @@ class Bagel(nn.Module):
|
||||
v_t = v_t[packed_vae_token_indexes]
|
||||
return v_t
|
||||
|
||||
def _forward_flow(
|
||||
def forward(
|
||||
self,
|
||||
x_t: torch.Tensor,
|
||||
timestep: torch.LongTensor,
|
||||
|
||||
@@ -159,6 +159,9 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
self.od_config = od_config
|
||||
self.device = get_local_device()
|
||||
|
||||
self.scheduler: object | None = None
|
||||
self.scheduler_kwargs: dict = {}
|
||||
|
||||
model = od_config.model
|
||||
local_files_only = os.path.exists(model)
|
||||
if local_files_only:
|
||||
@@ -630,7 +633,7 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
enabled=self.device.type != "cpu",
|
||||
dtype=self.od_config.dtype,
|
||||
):
|
||||
latents = self.bagel.generate_image(
|
||||
latents, trajectory_latents, trajectory_timesteps, trajectory_log_probs = self.bagel.generate_image(
|
||||
past_key_values=gen_context["past_key_values"],
|
||||
cfg_text_past_key_values=cfg_text_context["past_key_values"],
|
||||
cfg_img_past_key_values=cfg_img_context["past_key_values"],
|
||||
@@ -650,11 +653,36 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
cfg_img_packed_query_indexes=generation_input_cfg_img["cfg_packed_query_indexes"],
|
||||
cfg_img_key_values_lens=generation_input_cfg_img["cfg_key_values_lens"],
|
||||
cfg_img_packed_key_value_indexes=generation_input_cfg_img["cfg_packed_key_value_indexes"],
|
||||
return_trajectory_latents=req.sampling_params.return_trajectory_latents,
|
||||
scheduler=self.scheduler,
|
||||
scheduler_kwargs=self.scheduler_kwargs,
|
||||
)
|
||||
|
||||
img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape)
|
||||
|
||||
# Build trajectory output when requested
|
||||
trajectory_latents_stacked: torch.Tensor | None = None
|
||||
trajectory_timesteps_stacked: torch.Tensor | None = None
|
||||
trajectory_decoded: list[Image.Image] | None = None
|
||||
if trajectory_latents:
|
||||
trajectory_latents_stacked = torch.stack(trajectory_latents)
|
||||
trajectory_timesteps_stacked = torch.stack(trajectory_timesteps)
|
||||
if req.sampling_params.return_trajectory_decoded:
|
||||
trajectory_decoded = [
|
||||
self._decode_image_from_latent(self.bagel, self.vae, lat, image_shape) for lat in trajectory_latents
|
||||
]
|
||||
|
||||
trajectory_log_probs_stacked: torch.Tensor | None = None
|
||||
if trajectory_log_probs:
|
||||
trajectory_log_probs_stacked = torch.stack(trajectory_log_probs)
|
||||
|
||||
return DiffusionOutput(
|
||||
output=img, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
|
||||
output=img,
|
||||
trajectory_latents=trajectory_latents_stacked,
|
||||
trajectory_timesteps=trajectory_timesteps_stacked,
|
||||
trajectory_log_probs=trajectory_log_probs_stacked,
|
||||
trajectory_decoded=trajectory_decoded,
|
||||
stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
|
||||
)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
|
||||
@@ -582,6 +582,7 @@ class LongCatImageTransformer2DModel(nn.Module):
|
||||
"""
|
||||
|
||||
_repeated_blocks = ["LongCatImageTransformerBlock", "LongCatImageSingleTransformerBlock"]
|
||||
_layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
|
||||
|
||||
# Sequence Parallelism for LongCat (following diffusers' _cp_plan pattern)
|
||||
_sp_plan = {
|
||||
|
||||
@@ -28,8 +28,6 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
get_cfg_group,
|
||||
get_classifier_free_guidance_rank,
|
||||
get_classifier_free_guidance_world_size,
|
||||
)
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
@@ -122,6 +120,31 @@ def calculate_shift(
|
||||
return mu
|
||||
|
||||
|
||||
class _VideoAudioScheduler:
|
||||
"""Composite scheduler dispatching to video and audio schedulers."""
|
||||
|
||||
def __init__(self, video_scheduler, audio_scheduler):
|
||||
self.video_scheduler = video_scheduler
|
||||
self.audio_scheduler = audio_scheduler
|
||||
|
||||
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
|
||||
video_out = self.video_scheduler.step(
|
||||
noise_pred[0],
|
||||
t[0],
|
||||
latents[0],
|
||||
return_dict=False,
|
||||
generator=generator,
|
||||
)[0]
|
||||
audio_out = self.audio_scheduler.step(
|
||||
noise_pred[1],
|
||||
t[1],
|
||||
latents[1],
|
||||
return_dict=False,
|
||||
generator=generator,
|
||||
)[0]
|
||||
return ((video_out, audio_out),)
|
||||
|
||||
|
||||
class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
def __init__(
|
||||
self,
|
||||
@@ -542,6 +565,10 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2)
|
||||
return latents
|
||||
|
||||
@staticmethod
|
||||
def _unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor:
|
||||
return latents[:, :num_frames]
|
||||
|
||||
def prepare_latents(
|
||||
self,
|
||||
batch_size: int = 1,
|
||||
@@ -597,25 +624,49 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
noise_scale: float = 0.0,
|
||||
dtype: torch.dtype | None = None,
|
||||
device: torch.device | None = None,
|
||||
generator: torch.Generator | None = None,
|
||||
generator: torch.Generator | list[torch.Generator] | None = None,
|
||||
latents: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, int]:
|
||||
) -> tuple[torch.Tensor, int, int]:
|
||||
original_latent_length = audio_latent_length
|
||||
padded_latent_length = original_latent_length
|
||||
|
||||
latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
|
||||
|
||||
sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1)
|
||||
if sp_size > 1:
|
||||
padded_latent_length += (sp_size - (original_latent_length % sp_size)) % sp_size
|
||||
|
||||
if latents is not None:
|
||||
if latents.ndim == 4:
|
||||
# latents are of shape [B, C, L, M], need to be packed
|
||||
latents = self._pack_audio_latents(latents)
|
||||
if latents.ndim != 3:
|
||||
raise ValueError(
|
||||
f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is [batch_size, num_seq, num_features]." # noqa
|
||||
f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is "
|
||||
"[batch_size, num_seq, num_features] or [batch_size, num_channels, audio_length, mel_bins]."
|
||||
)
|
||||
latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std)
|
||||
latents = self._create_noised_state(latents, noise_scale, generator)
|
||||
return latents.to(device=device, dtype=dtype)
|
||||
|
||||
# TODO: confirm whether this logic is correct
|
||||
latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
|
||||
if latents.shape[1] not in {original_latent_length, padded_latent_length}:
|
||||
raise ValueError(
|
||||
"Provided `audio_latents` has incompatible audio frame count "
|
||||
f"{latents.shape[1]}; expected {original_latent_length} or {padded_latent_length}."
|
||||
)
|
||||
|
||||
shape = (batch_size, num_channels_latents, audio_latent_length, latent_mel_bins)
|
||||
if latents.shape[1] == original_latent_length and padded_latent_length > original_latent_length:
|
||||
padding = torch.zeros(
|
||||
latents.shape[0],
|
||||
padded_latent_length - original_latent_length,
|
||||
latents.shape[2],
|
||||
dtype=latents.dtype,
|
||||
device=latents.device,
|
||||
)
|
||||
latents = torch.cat([latents, padding], dim=1)
|
||||
|
||||
return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length
|
||||
|
||||
shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins)
|
||||
|
||||
if isinstance(generator, list) and len(generator) != batch_size:
|
||||
raise ValueError(
|
||||
@@ -625,7 +676,7 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
|
||||
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
latents = self._pack_audio_latents(latents)
|
||||
return latents
|
||||
return latents, original_latent_length, padded_latent_length
|
||||
|
||||
@property
|
||||
def guidance_scale(self):
|
||||
@@ -655,147 +706,44 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
def interrupt(self):
|
||||
return self._interrupt
|
||||
|
||||
def _is_cfg_parallel_enabled(self, do_true_cfg: bool) -> bool:
|
||||
return do_true_cfg and get_classifier_free_guidance_world_size() > 1
|
||||
|
||||
def _transformer_cache_context(self, context_name: str):
|
||||
cache_context = getattr(self.transformer, "cache_context", None)
|
||||
if callable(cache_context):
|
||||
return cache_context(context_name)
|
||||
return nullcontext()
|
||||
|
||||
def _predict_noise_av(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
def predict_noise(self, **kwargs):
|
||||
with self._transformer_cache_context("cond_uncond"):
|
||||
noise_pred_video, noise_pred_audio = self.transformer(**kwargs)
|
||||
return noise_pred_video, noise_pred_audio
|
||||
|
||||
def predict_noise_av_maybe_with_cfg(
|
||||
self,
|
||||
do_true_cfg: bool,
|
||||
true_cfg_scale: float,
|
||||
positive_kwargs: dict[str, Any],
|
||||
negative_kwargs: dict[str, Any] | None,
|
||||
guidance_rescale: float = 0.0,
|
||||
cfg_normalize: bool = False,
|
||||
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
|
||||
if do_true_cfg:
|
||||
cfg_parallel_ready = get_classifier_free_guidance_world_size() > 1
|
||||
|
||||
if cfg_parallel_ready:
|
||||
cfg_group = get_cfg_group()
|
||||
cfg_rank = get_classifier_free_guidance_rank()
|
||||
|
||||
if cfg_rank == 0:
|
||||
noise_pred_video, noise_pred_audio = self._predict_noise_av(**positive_kwargs)
|
||||
else:
|
||||
noise_pred_video, noise_pred_audio = self._predict_noise_av(**negative_kwargs)
|
||||
|
||||
noise_pred_video = noise_pred_video.float()
|
||||
noise_pred_audio = noise_pred_audio.float()
|
||||
|
||||
gathered_video = cfg_group.all_gather(noise_pred_video, separate_tensors=True)
|
||||
gathered_audio = cfg_group.all_gather(noise_pred_audio, separate_tensors=True)
|
||||
|
||||
if cfg_rank == 0:
|
||||
noise_pred_video_text = gathered_video[0]
|
||||
noise_pred_video_uncond = gathered_video[1]
|
||||
noise_pred_audio_text = gathered_audio[0]
|
||||
noise_pred_audio_uncond = gathered_audio[1]
|
||||
|
||||
noise_pred_video = self.combine_cfg_noise(
|
||||
noise_pred_video_text,
|
||||
noise_pred_video_uncond,
|
||||
true_cfg_scale,
|
||||
cfg_normalize,
|
||||
)
|
||||
noise_pred_audio = self.combine_cfg_noise(
|
||||
noise_pred_audio_text,
|
||||
noise_pred_audio_uncond,
|
||||
true_cfg_scale,
|
||||
cfg_normalize,
|
||||
)
|
||||
|
||||
if guidance_rescale > 0:
|
||||
noise_pred_video = rescale_noise_cfg(
|
||||
noise_pred_video,
|
||||
noise_pred_video_text,
|
||||
guidance_rescale=guidance_rescale,
|
||||
)
|
||||
noise_pred_audio = rescale_noise_cfg(
|
||||
noise_pred_audio,
|
||||
noise_pred_audio_text,
|
||||
guidance_rescale=guidance_rescale,
|
||||
)
|
||||
return noise_pred_video, noise_pred_audio
|
||||
return None, None
|
||||
|
||||
noise_pred_video_text, noise_pred_audio_text = self._predict_noise_av(**positive_kwargs)
|
||||
noise_pred_video_uncond, noise_pred_audio_uncond = self._predict_noise_av(**negative_kwargs)
|
||||
|
||||
noise_pred_video_text = noise_pred_video_text.float()
|
||||
noise_pred_audio_text = noise_pred_audio_text.float()
|
||||
noise_pred_video_uncond = noise_pred_video_uncond.float()
|
||||
noise_pred_audio_uncond = noise_pred_audio_uncond.float()
|
||||
|
||||
noise_pred_video = self.combine_cfg_noise(
|
||||
noise_pred_video_text,
|
||||
noise_pred_video_uncond,
|
||||
true_cfg_scale,
|
||||
cfg_normalize,
|
||||
)
|
||||
noise_pred_audio = self.combine_cfg_noise(
|
||||
noise_pred_audio_text,
|
||||
noise_pred_audio_uncond,
|
||||
true_cfg_scale,
|
||||
cfg_normalize,
|
||||
)
|
||||
|
||||
if guidance_rescale > 0:
|
||||
noise_pred_video = rescale_noise_cfg(
|
||||
noise_pred_video,
|
||||
noise_pred_video_text,
|
||||
guidance_rescale=guidance_rescale,
|
||||
)
|
||||
noise_pred_audio = rescale_noise_cfg(
|
||||
noise_pred_audio,
|
||||
noise_pred_audio_text,
|
||||
guidance_rescale=guidance_rescale,
|
||||
)
|
||||
|
||||
return noise_pred_video, noise_pred_audio
|
||||
|
||||
noise_pred_video, noise_pred_audio = self._predict_noise_av(**positive_kwargs)
|
||||
return noise_pred_video.float(), noise_pred_audio.float()
|
||||
|
||||
def _scheduler_step_video_audio_maybe_with_cfg(
|
||||
def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, true_cfg_scale, cfg_normalize=False):
|
||||
"""Per-element CFG combine with guidance_rescale support."""
|
||||
(video_pos, audio_pos) = positive_noise_pred
|
||||
(video_neg, audio_neg) = negative_noise_pred
|
||||
video_combined = super().combine_cfg_noise(video_pos, video_neg, true_cfg_scale, cfg_normalize)
|
||||
audio_combined = super().combine_cfg_noise(audio_pos, audio_neg, true_cfg_scale, cfg_normalize)
|
||||
if self._guidance_rescale and self._guidance_rescale > 0:
|
||||
video_combined = rescale_noise_cfg(video_combined, video_pos, guidance_rescale=self._guidance_rescale)
|
||||
audio_combined = rescale_noise_cfg(audio_combined, audio_pos, guidance_rescale=self._guidance_rescale)
|
||||
return (video_combined, audio_combined)
|
||||
|
||||
def _synchronize_cfg_parallel_step_output(
|
||||
self,
|
||||
noise_pred_video: torch.Tensor | None,
|
||||
noise_pred_audio: torch.Tensor | None,
|
||||
t: torch.Tensor,
|
||||
latents: torch.Tensor,
|
||||
audio_latents: torch.Tensor,
|
||||
audio_scheduler: FlowMatchEulerDiscreteScheduler,
|
||||
latents: tuple[torch.Tensor, torch.Tensor],
|
||||
do_true_cfg: bool,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
cfg_parallel_ready = self._is_cfg_parallel_enabled(do_true_cfg)
|
||||
if not (do_true_cfg and get_classifier_free_guidance_world_size() > 1):
|
||||
return latents
|
||||
|
||||
if cfg_parallel_ready:
|
||||
cfg_group = get_cfg_group()
|
||||
cfg_rank = get_classifier_free_guidance_rank()
|
||||
|
||||
if cfg_rank == 0:
|
||||
latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
|
||||
audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
|
||||
|
||||
latents = latents.contiguous()
|
||||
audio_latents = audio_latents.contiguous()
|
||||
cfg_group.broadcast(latents, src=0)
|
||||
cfg_group.broadcast(audio_latents, src=0)
|
||||
return latents, audio_latents
|
||||
|
||||
latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
|
||||
audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
|
||||
return latents, audio_latents
|
||||
# Without this sync, CUDA async execution causes non-deterministic
|
||||
# numerical drift across denoising steps in CFG parallel mode,
|
||||
# producing different video outputs across runs.
|
||||
latents = tuple(tensor.contiguous() for tensor in latents)
|
||||
device = next((tensor.device for tensor in latents if tensor.is_cuda), None)
|
||||
if device is not None:
|
||||
torch.cuda.current_stream(device).synchronize()
|
||||
return latents
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(
|
||||
@@ -828,6 +776,8 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
attention_kwargs: dict[str, Any] | None = None,
|
||||
max_sequence_length: int | None = None,
|
||||
) -> DiffusionOutput:
|
||||
# Extract prompt/negative_prompt from request.
|
||||
# Input format: req.prompts is a list of str or dict with "prompt"/"negative_prompt" keys.
|
||||
prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
|
||||
if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
|
||||
negative_prompt = None
|
||||
@@ -869,6 +819,7 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
else req.sampling_params.extra_args.get("audio_latents", audio_latents)
|
||||
)
|
||||
|
||||
# Override with pre-computed embeddings if provided in request.
|
||||
req_prompt_embeds = [_get_prompt_field(p, "prompt_embeds") for p in req.prompts]
|
||||
if any(p is not None for p in req_prompt_embeds):
|
||||
prompt_embeds = torch.stack(req_prompt_embeds) # type: ignore[arg-type]
|
||||
@@ -939,20 +890,17 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
max_sequence_length=max_sequence_length,
|
||||
device=device,
|
||||
)
|
||||
cfg_parallel_ready = self._is_cfg_parallel_enabled(self.do_classifier_free_guidance)
|
||||
if self.do_classifier_free_guidance and not cfg_parallel_ready:
|
||||
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
|
||||
prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
|
||||
|
||||
# Compute positive prompt connectors
|
||||
additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0
|
||||
connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
|
||||
prompt_embeds, additive_attention_mask, additive_mask=True
|
||||
)
|
||||
|
||||
# Compute negative prompt connectors when CFG is enabled
|
||||
negative_connector_prompt_embeds = None
|
||||
negative_connector_audio_prompt_embeds = None
|
||||
negative_connector_attention_mask = None
|
||||
if cfg_parallel_ready:
|
||||
if self.do_classifier_free_guidance:
|
||||
negative_additive_attention_mask = (
|
||||
1 - negative_prompt_attention_mask.to(negative_prompt_embeds.dtype)
|
||||
) * -1000000.0
|
||||
@@ -1027,20 +975,7 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
num_channels_latents_audio = (
|
||||
self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8
|
||||
)
|
||||
|
||||
# padding audio_latents if needed
|
||||
sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1)
|
||||
if sp_size > 1:
|
||||
pad_len = (sp_size - (audio_num_frames % sp_size)) % sp_size
|
||||
if pad_len > 0:
|
||||
if audio_latents is not None:
|
||||
pad_shape = list(audio_latents.shape)
|
||||
pad_shape[2] = pad_len
|
||||
padding = torch.zeros(pad_shape, dtype=audio_latents.dtype, device=audio_latents.device)
|
||||
audio_latents = torch.cat([audio_latents, padding], dim=2)
|
||||
audio_num_frames += pad_len
|
||||
|
||||
audio_latents = self.prepare_audio_latents(
|
||||
audio_latents, original_audio_num_frames, padded_audio_num_frames = self.prepare_audio_latents(
|
||||
batch_size * num_videos_per_prompt,
|
||||
num_channels_latents=num_channels_latents_audio,
|
||||
audio_latent_length=audio_num_frames,
|
||||
@@ -1061,6 +996,7 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
self.scheduler.config.get("max_shift", 2.05),
|
||||
)
|
||||
audio_scheduler = copy.deepcopy(self.scheduler)
|
||||
video_audio_scheduler = _VideoAudioScheduler(self.scheduler, audio_scheduler)
|
||||
_ = retrieve_timesteps(
|
||||
audio_scheduler,
|
||||
num_inference_steps,
|
||||
@@ -1083,12 +1019,10 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate
|
||||
)
|
||||
audio_coords = self.transformer.audio_rope.prepare_audio_coords(
|
||||
audio_latents.shape[0], audio_num_frames, audio_latents.device
|
||||
audio_latents.shape[0], padded_audio_num_frames, audio_latents.device
|
||||
)
|
||||
# Duplicate the positional ids as well if using CFG
|
||||
if self.do_classifier_free_guidance and not cfg_parallel_ready:
|
||||
video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) # Repeat twice in batch dim
|
||||
audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1))
|
||||
# No coord duplication needed: mixin handles CFG via separate forward calls,
|
||||
# not batch=2. Each forward gets batch=1 coords directly.
|
||||
|
||||
with self.progress_bar(total=len(timesteps)) as pbar:
|
||||
for i, t in enumerate(timesteps):
|
||||
@@ -1097,119 +1031,60 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
|
||||
self._current_timestep = t
|
||||
|
||||
if cfg_parallel_ready:
|
||||
latent_model_input = latents.to(prompt_embeds.dtype)
|
||||
audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
|
||||
timestep = t.expand(latent_model_input.shape[0])
|
||||
latent_model_input = latents.to(prompt_embeds.dtype)
|
||||
audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
|
||||
timestep = t.expand(latent_model_input.shape[0])
|
||||
do_true_cfg = self.do_classifier_free_guidance
|
||||
|
||||
positive_kwargs = {
|
||||
"hidden_states": latent_model_input,
|
||||
"audio_hidden_states": audio_latent_model_input,
|
||||
"encoder_hidden_states": connector_prompt_embeds,
|
||||
"audio_encoder_hidden_states": connector_audio_prompt_embeds,
|
||||
"timestep": timestep,
|
||||
"encoder_attention_mask": connector_attention_mask,
|
||||
"audio_encoder_attention_mask": connector_attention_mask,
|
||||
"num_frames": latent_num_frames,
|
||||
"height": latent_height,
|
||||
"width": latent_width,
|
||||
"fps": frame_rate,
|
||||
"audio_num_frames": audio_num_frames,
|
||||
"video_coords": video_coords,
|
||||
"audio_coords": audio_coords,
|
||||
"attention_kwargs": attention_kwargs,
|
||||
"return_dict": False,
|
||||
}
|
||||
negative_kwargs = {
|
||||
"hidden_states": latent_model_input,
|
||||
"audio_hidden_states": audio_latent_model_input,
|
||||
positive_kwargs = {
|
||||
"hidden_states": latent_model_input,
|
||||
"audio_hidden_states": audio_latent_model_input,
|
||||
"encoder_hidden_states": connector_prompt_embeds,
|
||||
"audio_encoder_hidden_states": connector_audio_prompt_embeds,
|
||||
"timestep": timestep,
|
||||
"encoder_attention_mask": connector_attention_mask,
|
||||
"audio_encoder_attention_mask": connector_attention_mask,
|
||||
"num_frames": latent_num_frames,
|
||||
"height": latent_height,
|
||||
"width": latent_width,
|
||||
"fps": frame_rate,
|
||||
"audio_num_frames": padded_audio_num_frames,
|
||||
"video_coords": video_coords,
|
||||
"audio_coords": audio_coords,
|
||||
"attention_kwargs": attention_kwargs,
|
||||
"return_dict": False,
|
||||
}
|
||||
negative_kwargs = (
|
||||
{
|
||||
**positive_kwargs,
|
||||
"encoder_hidden_states": negative_connector_prompt_embeds,
|
||||
"audio_encoder_hidden_states": negative_connector_audio_prompt_embeds,
|
||||
"timestep": timestep,
|
||||
"encoder_attention_mask": negative_connector_attention_mask,
|
||||
"audio_encoder_attention_mask": negative_connector_attention_mask,
|
||||
"num_frames": latent_num_frames,
|
||||
"height": latent_height,
|
||||
"width": latent_width,
|
||||
"fps": frame_rate,
|
||||
"audio_num_frames": audio_num_frames,
|
||||
"video_coords": video_coords,
|
||||
"audio_coords": audio_coords,
|
||||
"attention_kwargs": attention_kwargs,
|
||||
"return_dict": False,
|
||||
}
|
||||
if do_true_cfg
|
||||
else None
|
||||
)
|
||||
|
||||
noise_pred_video, noise_pred_audio = self.predict_noise_av_maybe_with_cfg(
|
||||
do_true_cfg=True,
|
||||
true_cfg_scale=guidance_scale,
|
||||
positive_kwargs=positive_kwargs,
|
||||
negative_kwargs=negative_kwargs,
|
||||
guidance_rescale=guidance_rescale,
|
||||
cfg_normalize=False,
|
||||
)
|
||||
noise_pred_video, noise_pred_audio = self.predict_noise_maybe_with_cfg(
|
||||
do_true_cfg=do_true_cfg,
|
||||
true_cfg_scale=guidance_scale,
|
||||
positive_kwargs=positive_kwargs,
|
||||
negative_kwargs=negative_kwargs,
|
||||
cfg_normalize=False,
|
||||
)
|
||||
|
||||
latents, audio_latents = self._scheduler_step_video_audio_maybe_with_cfg(
|
||||
noise_pred_video,
|
||||
noise_pred_audio,
|
||||
t,
|
||||
latents,
|
||||
audio_latents,
|
||||
audio_scheduler,
|
||||
do_true_cfg=True,
|
||||
)
|
||||
else:
|
||||
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
|
||||
latent_model_input = latent_model_input.to(prompt_embeds.dtype)
|
||||
audio_latent_model_input = (
|
||||
torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents
|
||||
)
|
||||
audio_latent_model_input = audio_latent_model_input.to(prompt_embeds.dtype)
|
||||
|
||||
timestep = t.expand(latent_model_input.shape[0])
|
||||
|
||||
with self._transformer_cache_context("cond_uncond"):
|
||||
noise_pred_video, noise_pred_audio = self.transformer(
|
||||
hidden_states=latent_model_input,
|
||||
audio_hidden_states=audio_latent_model_input,
|
||||
encoder_hidden_states=connector_prompt_embeds,
|
||||
audio_encoder_hidden_states=connector_audio_prompt_embeds,
|
||||
timestep=timestep,
|
||||
encoder_attention_mask=connector_attention_mask,
|
||||
audio_encoder_attention_mask=connector_attention_mask,
|
||||
num_frames=latent_num_frames,
|
||||
height=latent_height,
|
||||
width=latent_width,
|
||||
fps=frame_rate,
|
||||
audio_num_frames=audio_num_frames,
|
||||
video_coords=video_coords,
|
||||
audio_coords=audio_coords,
|
||||
attention_kwargs=attention_kwargs,
|
||||
return_dict=False,
|
||||
)
|
||||
noise_pred_video = noise_pred_video.float()
|
||||
noise_pred_audio = noise_pred_audio.float()
|
||||
|
||||
if self.do_classifier_free_guidance:
|
||||
noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2)
|
||||
noise_pred_video = noise_pred_video_uncond + guidance_scale * (
|
||||
noise_pred_video_text - noise_pred_video_uncond
|
||||
)
|
||||
|
||||
noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2)
|
||||
noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (
|
||||
noise_pred_audio_text - noise_pred_audio_uncond
|
||||
)
|
||||
|
||||
if guidance_rescale > 0:
|
||||
noise_pred_video = rescale_noise_cfg(
|
||||
noise_pred_video, noise_pred_video_text, guidance_rescale=guidance_rescale
|
||||
)
|
||||
noise_pred_audio = rescale_noise_cfg(
|
||||
noise_pred_audio, noise_pred_audio_text, guidance_rescale=guidance_rescale
|
||||
)
|
||||
|
||||
latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
|
||||
audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
|
||||
latents, audio_latents = self.scheduler_step_maybe_with_cfg(
|
||||
(noise_pred_video, noise_pred_audio),
|
||||
(t, t),
|
||||
(latents, audio_latents),
|
||||
do_true_cfg=do_true_cfg,
|
||||
per_request_scheduler=video_audio_scheduler,
|
||||
)
|
||||
latents, audio_latents = self._synchronize_cfg_parallel_step_output(
|
||||
(latents, audio_latents),
|
||||
do_true_cfg=do_true_cfg,
|
||||
)
|
||||
|
||||
pbar.update()
|
||||
|
||||
@@ -1225,10 +1100,15 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
|
||||
)
|
||||
|
||||
audio_latents = self._unpad_audio_latents(audio_latents, original_audio_num_frames)
|
||||
audio_latents = self._denormalize_audio_latents(
|
||||
audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std
|
||||
)
|
||||
audio_latents = self._unpack_audio_latents(audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins)
|
||||
audio_latents = self._unpack_audio_latents(
|
||||
audio_latents,
|
||||
original_audio_num_frames,
|
||||
num_mel_bins=latent_mel_bins,
|
||||
)
|
||||
|
||||
if output_type == "latent":
|
||||
video = latents
|
||||
|
||||
@@ -14,7 +14,7 @@ import torch
|
||||
import torch.nn as nn
|
||||
from diffusers import FlowMatchEulerDiscreteScheduler
|
||||
from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
|
||||
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg, retrieve_timesteps
|
||||
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
|
||||
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from diffusers.video_processor import VideoProcessor
|
||||
@@ -22,7 +22,6 @@ from vllm.logger import init_logger
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.parallel_state import get_cfg_group, get_classifier_free_guidance_rank
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
@@ -46,6 +45,32 @@ def get_ltx2_post_process_func(od_config: OmniDiffusionConfig):
|
||||
return _get_ltx2_post_process_func(od_config)
|
||||
|
||||
|
||||
class _I2VVideoAudioScheduler:
|
||||
"""Composite scheduler for I2V: uses _step_video_latents_i2v for video, standard step for audio."""
|
||||
|
||||
def __init__(self, pipeline, audio_scheduler, latent_num_frames, latent_height, latent_width):
|
||||
self.video_scheduler = pipeline.scheduler
|
||||
self.audio_scheduler = audio_scheduler
|
||||
self._pipeline = pipeline
|
||||
self._latent_num_frames = latent_num_frames
|
||||
self._latent_height = latent_height
|
||||
self._latent_width = latent_width
|
||||
|
||||
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
|
||||
video_out = self._pipeline._step_video_latents_i2v(
|
||||
noise_pred[0],
|
||||
latents[0],
|
||||
t[0],
|
||||
self._latent_num_frames,
|
||||
self._latent_height,
|
||||
self._latent_width,
|
||||
)
|
||||
audio_out = self.audio_scheduler.step(noise_pred[1], t[1], latents[1], return_dict=False, generator=generator)[
|
||||
0
|
||||
]
|
||||
return ((video_out, audio_out),)
|
||||
|
||||
|
||||
class LTX2ImageToVideoPipeline(LTX2Pipeline):
|
||||
support_image_input = True
|
||||
|
||||
@@ -287,6 +312,8 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
|
||||
attention_kwargs: dict[str, Any] | None = None,
|
||||
max_sequence_length: int | None = None,
|
||||
) -> DiffusionOutput:
|
||||
# Extract prompt/negative_prompt from request.
|
||||
# Input format: req.prompts is a list of str or dict with "prompt"/"negative_prompt" keys.
|
||||
prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
|
||||
if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
|
||||
negative_prompt = None
|
||||
@@ -328,6 +355,7 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
|
||||
else req.sampling_params.extra_args.get("audio_latents", audio_latents)
|
||||
)
|
||||
|
||||
# Override with pre-computed embeddings if provided in request.
|
||||
req_prompt_embeds = [_get_prompt_field(p, "prompt_embeds") for p in req.prompts]
|
||||
if any(p is not None for p in req_prompt_embeds):
|
||||
prompt_embeds = torch.stack(req_prompt_embeds) # type: ignore[arg-type]
|
||||
@@ -429,20 +457,17 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
|
||||
max_sequence_length=max_sequence_length,
|
||||
device=device,
|
||||
)
|
||||
cfg_parallel_ready = self._is_cfg_parallel_enabled(self.do_classifier_free_guidance)
|
||||
if self.do_classifier_free_guidance and not cfg_parallel_ready:
|
||||
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
|
||||
prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
|
||||
|
||||
# Compute positive prompt connectors
|
||||
additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0
|
||||
connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
|
||||
prompt_embeds, additive_attention_mask, additive_mask=True
|
||||
)
|
||||
|
||||
# Compute negative prompt connectors when CFG is enabled
|
||||
negative_connector_prompt_embeds = None
|
||||
negative_connector_audio_prompt_embeds = None
|
||||
negative_connector_attention_mask = None
|
||||
if cfg_parallel_ready:
|
||||
if self.do_classifier_free_guidance:
|
||||
negative_additive_attention_mask = (
|
||||
1 - negative_prompt_attention_mask.to(negative_prompt_embeds.dtype)
|
||||
) * -1000000.0
|
||||
@@ -500,8 +525,6 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
|
||||
generator,
|
||||
latents,
|
||||
)
|
||||
if self.do_classifier_free_guidance and not cfg_parallel_ready:
|
||||
conditioning_mask = torch.cat([conditioning_mask, conditioning_mask])
|
||||
|
||||
duration_s = num_frames / frame_rate
|
||||
audio_latents_per_second = (
|
||||
@@ -529,20 +552,7 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
|
||||
num_channels_latents_audio = (
|
||||
self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8
|
||||
)
|
||||
|
||||
# padding audio_latents if needed
|
||||
sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1)
|
||||
if sp_size > 1:
|
||||
pad_len = (sp_size - (audio_num_frames % sp_size)) % sp_size
|
||||
if pad_len > 0:
|
||||
if audio_latents is not None:
|
||||
pad_shape = list(audio_latents.shape)
|
||||
pad_shape[2] = pad_len
|
||||
padding = torch.zeros(pad_shape, dtype=audio_latents.dtype, device=audio_latents.device)
|
||||
audio_latents = torch.cat([audio_latents, padding], dim=2)
|
||||
audio_num_frames += pad_len
|
||||
|
||||
audio_latents = self.prepare_audio_latents(
|
||||
audio_latents, original_audio_num_frames, padded_audio_num_frames = self.prepare_audio_latents(
|
||||
batch_size * num_videos_per_prompt,
|
||||
num_channels_latents=num_channels_latents_audio,
|
||||
audio_latent_length=audio_num_frames,
|
||||
@@ -585,12 +595,17 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
|
||||
latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate
|
||||
)
|
||||
audio_coords = self.transformer.audio_rope.prepare_audio_coords(
|
||||
audio_latents.shape[0], audio_num_frames, audio_latents.device
|
||||
audio_latents.shape[0], padded_audio_num_frames, audio_latents.device
|
||||
)
|
||||
# Duplicate the positional ids as well if using CFG
|
||||
if self.do_classifier_free_guidance and not cfg_parallel_ready:
|
||||
video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) # Repeat twice in batch dim
|
||||
audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1))
|
||||
|
||||
i2v_scheduler = _I2VVideoAudioScheduler(
|
||||
pipeline=self,
|
||||
audio_scheduler=audio_scheduler,
|
||||
latent_num_frames=latent_num_frames,
|
||||
latent_height=latent_height,
|
||||
latent_width=latent_width,
|
||||
)
|
||||
# No coord duplication needed: mixin handles CFG via separate forward calls.
|
||||
|
||||
with self.progress_bar(total=len(timesteps)) as pbar:
|
||||
for i, t in enumerate(timesteps):
|
||||
@@ -599,140 +614,62 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
|
||||
|
||||
self._current_timestep = t
|
||||
|
||||
if cfg_parallel_ready:
|
||||
latent_model_input = latents.to(prompt_embeds.dtype)
|
||||
audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
|
||||
latent_model_input = latents.to(prompt_embeds.dtype)
|
||||
audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
|
||||
timestep = t.expand(latent_model_input.shape[0])
|
||||
video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
|
||||
do_true_cfg = self.do_classifier_free_guidance
|
||||
|
||||
timestep = t.expand(latent_model_input.shape[0])
|
||||
video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
|
||||
|
||||
positive_kwargs = {
|
||||
"hidden_states": latent_model_input,
|
||||
"audio_hidden_states": audio_latent_model_input,
|
||||
"encoder_hidden_states": connector_prompt_embeds,
|
||||
"audio_encoder_hidden_states": connector_audio_prompt_embeds,
|
||||
"timestep": video_timestep,
|
||||
"audio_timestep": timestep,
|
||||
"encoder_attention_mask": connector_attention_mask,
|
||||
"audio_encoder_attention_mask": connector_attention_mask,
|
||||
"num_frames": latent_num_frames,
|
||||
"height": latent_height,
|
||||
"width": latent_width,
|
||||
"fps": frame_rate,
|
||||
"audio_num_frames": audio_num_frames,
|
||||
"video_coords": video_coords,
|
||||
"audio_coords": audio_coords,
|
||||
"attention_kwargs": attention_kwargs,
|
||||
"return_dict": False,
|
||||
}
|
||||
negative_kwargs = {
|
||||
"hidden_states": latent_model_input,
|
||||
"audio_hidden_states": audio_latent_model_input,
|
||||
positive_kwargs = {
|
||||
"hidden_states": latent_model_input,
|
||||
"audio_hidden_states": audio_latent_model_input,
|
||||
"encoder_hidden_states": connector_prompt_embeds,
|
||||
"audio_encoder_hidden_states": connector_audio_prompt_embeds,
|
||||
"timestep": video_timestep,
|
||||
"audio_timestep": timestep,
|
||||
"encoder_attention_mask": connector_attention_mask,
|
||||
"audio_encoder_attention_mask": connector_attention_mask,
|
||||
"num_frames": latent_num_frames,
|
||||
"height": latent_height,
|
||||
"width": latent_width,
|
||||
"fps": frame_rate,
|
||||
"audio_num_frames": padded_audio_num_frames,
|
||||
"video_coords": video_coords,
|
||||
"audio_coords": audio_coords,
|
||||
"attention_kwargs": attention_kwargs,
|
||||
"return_dict": False,
|
||||
}
|
||||
negative_kwargs = (
|
||||
{
|
||||
**positive_kwargs,
|
||||
"encoder_hidden_states": negative_connector_prompt_embeds,
|
||||
"audio_encoder_hidden_states": negative_connector_audio_prompt_embeds,
|
||||
"timestep": video_timestep,
|
||||
"audio_timestep": timestep,
|
||||
"encoder_attention_mask": negative_connector_attention_mask,
|
||||
"audio_encoder_attention_mask": negative_connector_attention_mask,
|
||||
"num_frames": latent_num_frames,
|
||||
"height": latent_height,
|
||||
"width": latent_width,
|
||||
"fps": frame_rate,
|
||||
"audio_num_frames": audio_num_frames,
|
||||
"video_coords": video_coords,
|
||||
"audio_coords": audio_coords,
|
||||
"attention_kwargs": attention_kwargs,
|
||||
"return_dict": False,
|
||||
}
|
||||
if do_true_cfg
|
||||
else None
|
||||
)
|
||||
|
||||
noise_pred_video, noise_pred_audio = self.predict_noise_av_maybe_with_cfg(
|
||||
do_true_cfg=True,
|
||||
true_cfg_scale=guidance_scale,
|
||||
positive_kwargs=positive_kwargs,
|
||||
negative_kwargs=negative_kwargs,
|
||||
guidance_rescale=guidance_rescale,
|
||||
cfg_normalize=False,
|
||||
)
|
||||
noise_pred_video, noise_pred_audio = self.predict_noise_maybe_with_cfg(
|
||||
do_true_cfg=do_true_cfg,
|
||||
true_cfg_scale=guidance_scale,
|
||||
positive_kwargs=positive_kwargs,
|
||||
negative_kwargs=negative_kwargs,
|
||||
cfg_normalize=False,
|
||||
)
|
||||
|
||||
if get_classifier_free_guidance_rank() == 0:
|
||||
latents = self._step_video_latents_i2v(
|
||||
noise_pred_video,
|
||||
latents,
|
||||
t,
|
||||
latent_num_frames,
|
||||
latent_height,
|
||||
latent_width,
|
||||
)
|
||||
audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
|
||||
|
||||
cfg_group = get_cfg_group()
|
||||
latents = latents.contiguous()
|
||||
audio_latents = audio_latents.contiguous()
|
||||
cfg_group.broadcast(latents, src=0)
|
||||
cfg_group.broadcast(audio_latents, src=0)
|
||||
else:
|
||||
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
|
||||
latent_model_input = latent_model_input.to(prompt_embeds.dtype)
|
||||
audio_latent_model_input = (
|
||||
torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents
|
||||
)
|
||||
audio_latent_model_input = audio_latent_model_input.to(prompt_embeds.dtype)
|
||||
|
||||
timestep = t.expand(latent_model_input.shape[0])
|
||||
video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
|
||||
|
||||
with self._transformer_cache_context("cond_uncond"):
|
||||
noise_pred_video, noise_pred_audio = self.transformer(
|
||||
hidden_states=latent_model_input,
|
||||
audio_hidden_states=audio_latent_model_input,
|
||||
encoder_hidden_states=connector_prompt_embeds,
|
||||
audio_encoder_hidden_states=connector_audio_prompt_embeds,
|
||||
timestep=video_timestep,
|
||||
audio_timestep=timestep,
|
||||
encoder_attention_mask=connector_attention_mask,
|
||||
audio_encoder_attention_mask=connector_attention_mask,
|
||||
num_frames=latent_num_frames,
|
||||
height=latent_height,
|
||||
width=latent_width,
|
||||
fps=frame_rate,
|
||||
audio_num_frames=audio_num_frames,
|
||||
video_coords=video_coords,
|
||||
audio_coords=audio_coords,
|
||||
attention_kwargs=attention_kwargs,
|
||||
return_dict=False,
|
||||
)
|
||||
noise_pred_video = noise_pred_video.float()
|
||||
noise_pred_audio = noise_pred_audio.float()
|
||||
|
||||
if self.do_classifier_free_guidance:
|
||||
noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2)
|
||||
noise_pred_video = noise_pred_video_uncond + guidance_scale * (
|
||||
noise_pred_video_text - noise_pred_video_uncond
|
||||
)
|
||||
|
||||
noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2)
|
||||
noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (
|
||||
noise_pred_audio_text - noise_pred_audio_uncond
|
||||
)
|
||||
|
||||
if guidance_rescale > 0:
|
||||
noise_pred_video = rescale_noise_cfg(
|
||||
noise_pred_video, noise_pred_video_text, guidance_rescale=guidance_rescale
|
||||
)
|
||||
noise_pred_audio = rescale_noise_cfg(
|
||||
noise_pred_audio, noise_pred_audio_text, guidance_rescale=guidance_rescale
|
||||
)
|
||||
|
||||
latents = self._step_video_latents_i2v(
|
||||
noise_pred_video,
|
||||
latents,
|
||||
t,
|
||||
latent_num_frames,
|
||||
latent_height,
|
||||
latent_width,
|
||||
)
|
||||
|
||||
audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
|
||||
latents, audio_latents = self.scheduler_step_maybe_with_cfg(
|
||||
(noise_pred_video, noise_pred_audio),
|
||||
(t, t),
|
||||
(latents, audio_latents),
|
||||
do_true_cfg=do_true_cfg,
|
||||
per_request_scheduler=i2v_scheduler,
|
||||
)
|
||||
latents, audio_latents = self._synchronize_cfg_parallel_step_output(
|
||||
(latents, audio_latents),
|
||||
do_true_cfg=do_true_cfg,
|
||||
)
|
||||
|
||||
pbar.update()
|
||||
|
||||
@@ -748,10 +685,15 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
|
||||
latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
|
||||
)
|
||||
|
||||
audio_latents = self._unpad_audio_latents(audio_latents, original_audio_num_frames)
|
||||
audio_latents = self._denormalize_audio_latents(
|
||||
audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std
|
||||
)
|
||||
audio_latents = self._unpack_audio_latents(audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins)
|
||||
audio_latents = self._unpack_audio_latents(
|
||||
audio_latents,
|
||||
original_audio_num_frames,
|
||||
num_mel_bins=latent_mel_bins,
|
||||
)
|
||||
|
||||
if output_type == "latent":
|
||||
video = latents
|
||||
|
||||
@@ -1624,7 +1624,12 @@ def get_magi_human_post_process_func(*args, **kwargs):
|
||||
def post_process(output):
|
||||
if isinstance(output, tuple) and len(output) == 2:
|
||||
video, audio = output
|
||||
return {"video": video, "audio": audio}
|
||||
return {
|
||||
"video": video,
|
||||
"audio": audio,
|
||||
"audio_sample_rate": 44100,
|
||||
"fps": 25,
|
||||
}
|
||||
return output
|
||||
|
||||
return post_process
|
||||
|
||||
@@ -114,6 +114,8 @@ class NextStepConfig(LlamaConfig):
|
||||
|
||||
|
||||
class NextStepModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = ["layers"]
|
||||
|
||||
def __init__(self, config: NextStepConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
|
||||
@@ -366,6 +366,7 @@ class OvisImageTransformer2DModel(nn.Module):
|
||||
"""
|
||||
|
||||
_repeated_blocks = ["OvisImageTransformerBlock", "OvisImageSingleTransformerBlock"]
|
||||
_layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
||||
@@ -387,6 +387,7 @@ class SD3Transformer2DModel(nn.Module):
|
||||
"""
|
||||
|
||||
_repeated_blocks = ["SD3TransformerBlock"]
|
||||
_layerwise_offload_blocks_attrs = ["transformer_blocks"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
||||
@@ -272,7 +272,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
|
||||
model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
|
||||
).to(self.device)
|
||||
self.vae = DistributedAutoencoderKLWan.from_pretrained(
|
||||
model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
|
||||
model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
|
||||
).to(self.device)
|
||||
|
||||
# Initialize transformers with correct config (weights loaded via load_weights)
|
||||
|
||||
@@ -217,7 +217,7 @@ class Wan22I2VPipeline(
|
||||
|
||||
# VAE
|
||||
self.vae = DistributedAutoencoderKLWan.from_pretrained(
|
||||
model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
|
||||
model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
|
||||
).to(self.device)
|
||||
|
||||
# Transformers (weights loaded via load_weights)
|
||||
|
||||
@@ -24,13 +24,13 @@ from typing import Any, cast
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
from diffusers import AutoencoderKLWan
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from torch import nn
|
||||
from transformers import AutoTokenizer, UMT5EncoderModel
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
@@ -174,8 +174,8 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress
|
||||
).to(self.device)
|
||||
|
||||
# VAE
|
||||
self.vae = AutoencoderKLWan.from_pretrained(
|
||||
model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
|
||||
self.vae = OmniAutoencoderKLWan.from_pretrained(
|
||||
model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
|
||||
).to(self.device)
|
||||
|
||||
# Single transformer (TI2V uses dense 5B model, not MoE)
|
||||
|
||||
@@ -21,9 +21,10 @@ import os
|
||||
from collections.abc import Callable, Iterable
|
||||
from typing import Any
|
||||
|
||||
import PIL.Image
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
|
||||
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
|
||||
from diffusers.utils import logging
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
@@ -59,7 +60,7 @@ def get_post_process_func(
|
||||
vae_config = json.load(f)
|
||||
vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8
|
||||
|
||||
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
|
||||
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2, do_convert_rgb=True)
|
||||
|
||||
def post_process_func(
|
||||
images: torch.Tensor,
|
||||
@@ -83,6 +84,20 @@ def calculate_shift(
|
||||
return mu
|
||||
|
||||
|
||||
# Copied from diffusers
|
||||
def retrieve_latents(
|
||||
encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
|
||||
):
|
||||
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
|
||||
return encoder_output.latent_dist.sample(generator)
|
||||
elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
|
||||
return encoder_output.latent_dist.mode()
|
||||
elif hasattr(encoder_output, "latents"):
|
||||
return encoder_output.latents
|
||||
else:
|
||||
raise AttributeError("Could not access latents of provided encoder_output")
|
||||
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
|
||||
def retrieve_timesteps(
|
||||
scheduler,
|
||||
@@ -187,6 +202,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
|
||||
)
|
||||
|
||||
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2, do_convert_rgb=True)
|
||||
|
||||
def encode_prompt(
|
||||
self,
|
||||
prompt: str | list[str],
|
||||
@@ -282,12 +299,45 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
device,
|
||||
generator,
|
||||
latents=None,
|
||||
image=None,
|
||||
timestep=None,
|
||||
):
|
||||
height = 2 * (int(height) // (self.vae_scale_factor * 2))
|
||||
width = 2 * (int(width) // (self.vae_scale_factor * 2))
|
||||
|
||||
shape = (batch_size, num_channels_latents, height, width)
|
||||
|
||||
if image is not None:
|
||||
if latents is not None:
|
||||
return latents.to(device=device, dtype=dtype)
|
||||
|
||||
image = image.to(device=device, dtype=dtype)
|
||||
if image.shape[1] != num_channels_latents:
|
||||
if isinstance(generator, list):
|
||||
image_latents = [
|
||||
retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i])
|
||||
for i in range(image.shape[0])
|
||||
]
|
||||
image_latents = torch.cat(image_latents, dim=0)
|
||||
else:
|
||||
image_latents = retrieve_latents(self.vae.encode(image), generator=generator)
|
||||
|
||||
image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
|
||||
else:
|
||||
image_latents = image
|
||||
|
||||
if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
|
||||
additional_image_per_prompt = batch_size // image_latents.shape[0]
|
||||
image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
|
||||
elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
|
||||
raise ValueError(
|
||||
f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
|
||||
)
|
||||
|
||||
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
latents = self.scheduler.scale_noise(image_latents, timestep, noise)
|
||||
return latents
|
||||
|
||||
if latents is None:
|
||||
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
else:
|
||||
@@ -296,6 +346,14 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
latents = latents.to(device)
|
||||
return latents
|
||||
|
||||
def get_timesteps(self, num_inference_steps, strength, device):
|
||||
init_timestep = min(num_inference_steps * strength, num_inference_steps)
|
||||
t_start = int(max(num_inference_steps - init_timestep, 0))
|
||||
timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
|
||||
if hasattr(self.scheduler, "set_begin_index"):
|
||||
self.scheduler.set_begin_index(t_start * self.scheduler.order)
|
||||
return timesteps, num_inference_steps - t_start
|
||||
|
||||
@property
|
||||
def guidance_scale(self):
|
||||
return self._guidance_scale
|
||||
@@ -320,6 +378,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
self,
|
||||
req: OmniDiffusionRequest,
|
||||
prompt: str | list[str] | None = None,
|
||||
image: PipelineImageInput = None,
|
||||
strength: float = 0.6,
|
||||
height: int = 1024,
|
||||
width: int = 1024,
|
||||
num_inference_steps: int = 50,
|
||||
@@ -347,6 +407,11 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
prompt (`str` or `list[str]`, *optional*):
|
||||
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
|
||||
instead.
|
||||
image (`PipelineImageInput`, *optional*):
|
||||
The image to use for img2img generation. If provided, the pipeline
|
||||
will perform img2img instead of text-to-image.
|
||||
strength (`float`, *optional*, defaults to 0.6):
|
||||
Indicates extent to transform the reference `image`. Must be between 0 and 1.
|
||||
height (`int`, *optional*, defaults to 1024):
|
||||
The height in pixels of the generated image.
|
||||
width (`int`, *optional*, defaults to 1024):
|
||||
@@ -425,6 +490,34 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
elif req.prompts:
|
||||
negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
|
||||
|
||||
# Handle img2img: extract image from request
|
||||
if image is None and req.prompts:
|
||||
if len(req.prompts) > 1:
|
||||
logger.warning(
|
||||
"This model only supports a single prompt for img2img, not a batched request. "
|
||||
"Taking only the first image for now."
|
||||
)
|
||||
first_prompt = req.prompts[0]
|
||||
if not isinstance(first_prompt, str):
|
||||
raw_image = first_prompt.get("multi_modal_data", {}).get("image")
|
||||
if raw_image is not None:
|
||||
if isinstance(raw_image, list):
|
||||
image = [PIL.Image.open(im) if isinstance(im, str) else raw_image[0] for im in raw_image[:1]]
|
||||
else:
|
||||
image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else raw_image
|
||||
|
||||
# strength is currently only applicable for Z-Image I2I; other pipelines ignore this parameter
|
||||
strength = req.sampling_params.strength if req.sampling_params.strength is not None else strength
|
||||
if strength is not None and image is None:
|
||||
logger.warning(
|
||||
"strength parameter (%.2f) is only applicable for image-to-image (I2I) generation. "
|
||||
"It will be ignored for text-to-image (T2I) generation.",
|
||||
strength,
|
||||
)
|
||||
strength = None
|
||||
if image is not None and strength is not None and (strength < 0 or strength > 1):
|
||||
raise ValueError(f"The value of strength should be in [0.0, 1.0] but is {strength}")
|
||||
|
||||
height = req.sampling_params.height or height
|
||||
width = req.sampling_params.width or width
|
||||
num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
|
||||
@@ -491,16 +584,71 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
# 4. Prepare latent variables
|
||||
num_channels_latents = self.transformer.in_channels
|
||||
|
||||
latents = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
torch.float32,
|
||||
device,
|
||||
generator,
|
||||
latents,
|
||||
)
|
||||
# img2img mode: prepare latents from input image
|
||||
if image is not None:
|
||||
# Handle image list - take first image
|
||||
if isinstance(image, list):
|
||||
image = image[0]
|
||||
|
||||
# Prepare image for VAE encoding using image_processor
|
||||
if not isinstance(image, torch.Tensor):
|
||||
init_image = self.image_processor.preprocess(image, height, width)
|
||||
image = init_image.to(dtype=torch.float32, device=device)
|
||||
|
||||
# Initialize scheduler kwargs for img2img
|
||||
mu = calculate_shift(
|
||||
(height // self.vae_scale_factor // 2) * (width // self.vae_scale_factor // 2),
|
||||
self.scheduler.config.get("base_image_seq_len", 256),
|
||||
self.scheduler.config.get("max_image_seq_len", 4096),
|
||||
self.scheduler.config.get("base_shift", 0.5),
|
||||
self.scheduler.config.get("max_shift", 1.15),
|
||||
)
|
||||
self.scheduler.sigma_min = 0.0
|
||||
scheduler_kwargs = {"mu": mu}
|
||||
|
||||
# First initialize timesteps in scheduler
|
||||
timesteps, num_inference_steps = retrieve_timesteps(
|
||||
self.scheduler,
|
||||
num_inference_steps,
|
||||
device,
|
||||
sigmas=sigmas,
|
||||
**scheduler_kwargs,
|
||||
)
|
||||
|
||||
# Then adjust timesteps based on strength
|
||||
timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
|
||||
|
||||
if num_inference_steps < 1:
|
||||
raise ValueError(
|
||||
f"After adjusting the num_inference_steps by strength parameter: "
|
||||
f"{strength}, the number of pipeline steps is {num_inference_steps} "
|
||||
f"which is < 1 and not appropriate for this pipeline."
|
||||
)
|
||||
latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
|
||||
|
||||
latents = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds[0].dtype,
|
||||
device,
|
||||
generator,
|
||||
latents,
|
||||
image,
|
||||
latent_timestep,
|
||||
)
|
||||
else:
|
||||
latents = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
torch.float32,
|
||||
device,
|
||||
generator,
|
||||
latents,
|
||||
)
|
||||
|
||||
# Repeat prompt_embeds for num_images_per_prompt
|
||||
if num_images_per_prompt > 1:
|
||||
@@ -509,25 +657,28 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)]
|
||||
|
||||
actual_batch_size = batch_size * num_images_per_prompt
|
||||
image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2)
|
||||
|
||||
# 5. Prepare timesteps
|
||||
mu = calculate_shift(
|
||||
image_seq_len,
|
||||
self.scheduler.config.get("base_image_seq_len", 256),
|
||||
self.scheduler.config.get("max_image_seq_len", 4096),
|
||||
self.scheduler.config.get("base_shift", 0.5),
|
||||
self.scheduler.config.get("max_shift", 1.15),
|
||||
)
|
||||
self.scheduler.sigma_min = 0.0
|
||||
scheduler_kwargs = {"mu": mu}
|
||||
timesteps, num_inference_steps = retrieve_timesteps(
|
||||
self.scheduler,
|
||||
num_inference_steps,
|
||||
device,
|
||||
sigmas=sigmas,
|
||||
**scheduler_kwargs,
|
||||
)
|
||||
if image is None:
|
||||
image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2)
|
||||
mu = calculate_shift(
|
||||
image_seq_len,
|
||||
self.scheduler.config.get("base_image_seq_len", 256),
|
||||
self.scheduler.config.get("max_image_seq_len", 4096),
|
||||
self.scheduler.config.get("base_shift", 0.5),
|
||||
self.scheduler.config.get("max_shift", 1.15),
|
||||
)
|
||||
self.scheduler.sigma_min = 0.0
|
||||
scheduler_kwargs = {"mu": mu}
|
||||
|
||||
timesteps, num_inference_steps = retrieve_timesteps(
|
||||
self.scheduler,
|
||||
num_inference_steps,
|
||||
device,
|
||||
sigmas=sigmas,
|
||||
**scheduler_kwargs,
|
||||
)
|
||||
|
||||
num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
|
||||
self._num_timesteps = len(timesteps)
|
||||
|
||||
|
||||
@@ -21,7 +21,7 @@ class PipelineModules:
|
||||
class ModuleDiscovery:
|
||||
"""Discovers pipeline components for offloading"""
|
||||
|
||||
DIT_ATTRS = ["transformer", "transformer_2", "dit", "sr_dit", "language_model", "transformer_blocks"]
|
||||
DIT_ATTRS = ["transformer", "transformer_2", "dit", "sr_dit", "language_model", "transformer_blocks", "model"]
|
||||
ENCODER_ATTRS = ["text_encoder", "text_encoder_2", "text_encoder_3", "image_encoder"]
|
||||
VAE_ATTRS = ["vae", "audio_vae"]
|
||||
|
||||
|
||||
@@ -26,6 +26,8 @@ class OmniDiffusionRequest:
|
||||
sampling_params: OmniDiffusionSamplingParams
|
||||
|
||||
request_ids: list[str] = field(default_factory=list)
|
||||
request_id: str | None = None
|
||||
kv_sender_info: dict | None = None
|
||||
|
||||
def __post_init__(self):
|
||||
"""Initialize dependent fields after dataclass initialization."""
|
||||
|
||||
@@ -179,6 +179,7 @@ class StageDiffusionClient:
|
||||
request_id: str,
|
||||
prompt: OmniPromptType,
|
||||
sampling_params: OmniDiffusionSamplingParams,
|
||||
kv_sender_info: dict[int, dict[str, Any]] | None = None,
|
||||
) -> None:
|
||||
self._request_socket.send(
|
||||
self._encoder.encode(
|
||||
@@ -187,6 +188,7 @@ class StageDiffusionClient:
|
||||
"request_id": request_id,
|
||||
"prompt": prompt,
|
||||
"sampling_params": self._sampling_params_to_dict(sampling_params),
|
||||
"kv_sender_info": kv_sender_info,
|
||||
}
|
||||
)
|
||||
)
|
||||
@@ -198,6 +200,7 @@ class StageDiffusionClient:
|
||||
request_id: str,
|
||||
prompts: list[OmniPromptType],
|
||||
sampling_params: OmniDiffusionSamplingParams,
|
||||
kv_sender_info: dict[int, dict[str, Any]] | None = None,
|
||||
) -> None:
|
||||
"""Submit a list of prompts as a single batched engine call.
|
||||
|
||||
@@ -206,7 +209,12 @@ class StageDiffusionClient:
|
||||
*request_id*.
|
||||
"""
|
||||
task = asyncio.create_task(
|
||||
self._run_batch(request_id, prompts, sampling_params),
|
||||
self._run_batch(
|
||||
request_id,
|
||||
prompts,
|
||||
sampling_params,
|
||||
kv_sender_info,
|
||||
),
|
||||
name=f"diffusion-batch-{request_id}",
|
||||
)
|
||||
self._tasks[request_id] = task
|
||||
@@ -216,6 +224,7 @@ class StageDiffusionClient:
|
||||
request_id: str,
|
||||
prompts: list[OmniPromptType],
|
||||
sampling_params: OmniDiffusionSamplingParams,
|
||||
kv_sender_info: dict[int, dict[str, Any]] | None = None,
|
||||
) -> None:
|
||||
try:
|
||||
self._request_socket.send(
|
||||
@@ -225,6 +234,7 @@ class StageDiffusionClient:
|
||||
"request_id": request_id,
|
||||
"prompts": prompts,
|
||||
"sampling_params": self._sampling_params_to_dict(sampling_params),
|
||||
"kv_sender_info": kv_sender_info,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
@@ -14,8 +14,10 @@ from multiprocessing.process import BaseProcess
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import msgspec
|
||||
import torch
|
||||
import zmq
|
||||
import zmq.asyncio
|
||||
from PIL import Image
|
||||
from vllm.logger import init_logger
|
||||
from vllm.transformers_utils.config import get_hf_file_to_dict
|
||||
from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx
|
||||
@@ -128,6 +130,7 @@ class StageDiffusionProc:
|
||||
request_id: str,
|
||||
prompt: Any,
|
||||
sampling_params_dict: dict,
|
||||
kv_sender_info: dict[str, Any] | None = None,
|
||||
) -> OmniRequestOutput:
|
||||
"""Build a diffusion request and run DiffusionEngine.step()."""
|
||||
sampling_params = self._reconstruct_sampling_params(sampling_params_dict)
|
||||
@@ -136,6 +139,8 @@ class StageDiffusionProc:
|
||||
prompts=[prompt],
|
||||
sampling_params=sampling_params,
|
||||
request_ids=[request_id],
|
||||
request_id=request_id,
|
||||
kv_sender_info=kv_sender_info,
|
||||
)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
@@ -150,6 +155,7 @@ class StageDiffusionProc:
|
||||
request_id: str,
|
||||
prompts: list[Any],
|
||||
sampling_params_dict: dict,
|
||||
kv_sender_info: dict[str, Any] | None = None,
|
||||
) -> OmniRequestOutput:
|
||||
"""Build a batched diffusion request and run DiffusionEngine.step().
|
||||
|
||||
@@ -163,7 +169,9 @@ class StageDiffusionProc:
|
||||
request = OmniDiffusionRequest(
|
||||
prompts=prompts,
|
||||
sampling_params=sampling_params,
|
||||
request_ids=[request_id] * len(prompts),
|
||||
request_ids=[f"{request_id}-{i}" for i in range(len(prompts))],
|
||||
request_id=request_id,
|
||||
kv_sender_info=kv_sender_info,
|
||||
)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
@@ -174,8 +182,13 @@ class StageDiffusionProc:
|
||||
merged_mm: dict[str, Any] = {}
|
||||
merged_metrics: dict[str, Any] = {}
|
||||
merged_durations: dict[str, float] = {}
|
||||
merged_custom: dict[str, Any] = {}
|
||||
peak_mem = 0.0
|
||||
latents = None
|
||||
trajectory_latents: list[torch.Tensor] | None = None
|
||||
trajectory_timesteps: list[torch.Tensor] | None = None
|
||||
trajectory_log_probs: torch.Tensor | None = None
|
||||
trajectory_decoded: list[Image.Image] | None = None
|
||||
final_output_type = "image"
|
||||
|
||||
for r in results:
|
||||
@@ -183,9 +196,18 @@ class StageDiffusionProc:
|
||||
merged_mm.update(r._multimodal_output)
|
||||
merged_metrics.update(r.metrics)
|
||||
merged_durations.update(r.stage_durations)
|
||||
merged_custom.update(r._custom_output)
|
||||
peak_mem = max(peak_mem, r.peak_memory_mb)
|
||||
if latents is None and r.latents is not None:
|
||||
latents = r.latents
|
||||
if trajectory_latents is None:
|
||||
trajectory_latents = r.trajectory_latents
|
||||
if trajectory_timesteps is None:
|
||||
trajectory_timesteps = r.trajectory_timesteps
|
||||
if trajectory_log_probs is None:
|
||||
trajectory_log_probs = r.trajectory_log_probs
|
||||
if trajectory_decoded is None:
|
||||
trajectory_decoded = r.trajectory_decoded
|
||||
if r.final_output_type != "image":
|
||||
final_output_type = r.final_output_type
|
||||
|
||||
@@ -195,6 +217,11 @@ class StageDiffusionProc:
|
||||
prompt=prompts[0] if len(prompts) == 1 else None,
|
||||
metrics=merged_metrics,
|
||||
latents=latents,
|
||||
trajectory_latents=trajectory_latents,
|
||||
trajectory_timesteps=trajectory_timesteps,
|
||||
trajectory_log_probs=trajectory_log_probs,
|
||||
trajectory_decoded=trajectory_decoded,
|
||||
custom_output=merged_custom or None,
|
||||
multimodal_output=merged_mm or None,
|
||||
final_output_type=final_output_type,
|
||||
stage_durations=merged_durations,
|
||||
@@ -325,10 +352,20 @@ class StageDiffusionProc:
|
||||
|
||||
tasks: dict[str, asyncio.Task] = {}
|
||||
|
||||
async def _dispatch_request(request_id: str, prompt: Any, sampling_params_dict: dict) -> None:
|
||||
async def _dispatch_request(
|
||||
request_id: str,
|
||||
prompt: Any,
|
||||
sampling_params_dict: dict,
|
||||
kv_sender_info: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""Process a single diffusion request and send the response."""
|
||||
try:
|
||||
result = await self._process_request(request_id, prompt, sampling_params_dict)
|
||||
result = await self._process_request(
|
||||
request_id,
|
||||
prompt,
|
||||
sampling_params_dict,
|
||||
kv_sender_info=kv_sender_info,
|
||||
)
|
||||
await response_socket.send(encoder.encode({"type": "result", "output": result}))
|
||||
except DiffusionRequestAbortedError as e:
|
||||
logger.info(
|
||||
@@ -363,6 +400,7 @@ class StageDiffusionProc:
|
||||
request_id,
|
||||
msg["prompt"],
|
||||
msg["sampling_params"],
|
||||
msg.get("kv_sender_info"),
|
||||
)
|
||||
)
|
||||
tasks[request_id] = task
|
||||
@@ -370,9 +408,19 @@ class StageDiffusionProc:
|
||||
elif msg_type == "add_batch_request":
|
||||
request_id = msg["request_id"]
|
||||
|
||||
async def _dispatch_batch(rid: str, prompts: list, sp_dict: dict) -> None:
|
||||
async def _dispatch_batch(
|
||||
rid: str,
|
||||
prompts: list,
|
||||
sp_dict: dict,
|
||||
kv_sender_info: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
try:
|
||||
result = await self._process_batch_request(rid, prompts, sp_dict)
|
||||
result = await self._process_batch_request(
|
||||
rid,
|
||||
prompts,
|
||||
sp_dict,
|
||||
kv_sender_info=kv_sender_info,
|
||||
)
|
||||
await response_socket.send(encoder.encode({"type": "result", "output": result}))
|
||||
except DiffusionRequestAbortedError as e:
|
||||
logger.info(
|
||||
@@ -399,6 +447,7 @@ class StageDiffusionProc:
|
||||
request_id,
|
||||
msg["prompts"],
|
||||
msg["sampling_params"],
|
||||
msg.get("kv_sender_info"),
|
||||
)
|
||||
)
|
||||
tasks[request_id] = task
|
||||
|
||||
Loaded 100 of 134 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user